From 4ff31614cf32c5142cb9e4cada7181f574d1f703 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Tue, 27 Oct 2026 11:12:28 -0700 Subject: [PATCH 01/12] feat(bundle): agent-env bundle register writes a bundle into the configured stores The bundle is planned as a run plans it, then each @local id it writes takes a registry id: - an entry takes , the prefix the folder's name and a dash by default; - a task's output takes -; - a derived id keeps its suffix; - a shipped agent named by its alias takes the alias. Each task and eval names the version of what it references. Each image the bundle builds is built for linux/amd64, pushed under a tag of its sources and registered by digest with its build context. A ledger in the target store's bundle_registrations makes a rerun write only what changed. Problems are refused before anything is written: an id the store holds that the bundle didn't register (unless --adopt), an @local id left in a document, or, for a shared store, a credential written out or a URL only this machine reaches. A step that leaves its agent unset stays unset. Co-Authored-By: Claude Opus 5.5 --- .../artifact/artifacts/docker_image.py | 31 +- src/agent_env/bundle/_secrets.py | 24 ++ src/agent_env/bundle/ledger.py | 48 ++- src/agent_env/bundle/materialize.py | 94 +++- src/agent_env/bundle/plan.py | 16 +- src/agent_env/bundle/register.py | 406 ++++++++++++++++++ src/agent_env/cli/bundle.py | 60 ++- tst/unit/bundle/register_test.py | 396 +++++++++++++++++ 8 files changed, 1032 insertions(+), 43 deletions(-) create mode 100644 src/agent_env/bundle/register.py create mode 100644 tst/unit/bundle/register_test.py diff --git a/src/agent_env/artifact/artifacts/docker_image.py b/src/agent_env/artifact/artifacts/docker_image.py index a2bdc8377..663281ee8 100644 --- a/src/agent_env/artifact/artifacts/docker_image.py +++ b/src/agent_env/artifact/artifacts/docker_image.py @@ -22,9 +22,9 @@ from agent_env.store.ids import fs_safe, image_repository, is_local_id, local_image_repository from agent_env.store.image_store.local_registry_image_store import LocalRegistryImageStore from agent_env.store.image_store.oci_registry_credentials import is_loopback_host, names_registry, registry_host_from_ref -from agent_env.store.image_store.registry_api import pin_digest, pinned_from +from agent_env.store.image_store.registry_api import NoSuchImageError, pin_digest, pinned_from from agent_env.utils.docker.context import BuildContext -from agent_env.utils.docker.build import DEFAULT_BUILD_PLATFORM, normalize_platform +from agent_env.utils.docker.build import DEFAULT_BUILD_PLATFORM, build_image, normalize_platform logger = logging.getLogger(__name__) @@ -561,6 +561,33 @@ class GitHubBuildResult: docker_context_github_url: str | None +def push_context_image(id: str, context_path: str, dockerfile_path: str, + before: Callable[[], None] | None = None) -> tuple[str, str]: + """The ref and digest of the image of ``id``'s build context in ``id``'s image store, under the tag of its sources: + built for linux/amd64 with this machine's Docker and pushed, calling ``before`` first, unless the store serves the + tag already. Raises DockerBuildError when the build fails, RuntimeError when the push does, and ValueError when the + registry can't be read.""" + from agent_env.config import get_config + + context = _context_of(Path(context_path), dockerfile_path) + tag = "src-" + context.source_digest(DEFAULT_BUILD_PLATFORM).removeprefix("sha256:") + store = get_config().get_image_store_for(id) + # ECR, unlike the OCI spec, refuses two separators in a row, which a derived id (``__env_image``) holds. + repository = image_repository(id).replace("__", "-") + ref = store.image_ref(repository, tag) + store.ensure_repository(repository) + try: + return ref, pin_digest(ref, store.auth(ref)).rpartition("@")[2] + except NoSuchImageError: + pass + if before is not None: + before() + local = f"{local_image_repository(id)}:{tag}" + build_image(Path(context_path) / context.dockerfile, Path(context_path), local, platform=DEFAULT_BUILD_PLATFORM) + _push_local_image(local, ref, store) + return ref, pin_digest(ref, store.auth(ref)).rpartition("@")[2] + + def _push_local_image(local_tag: str, image_ref: str, image_store) -> None: auth = image_store.auth(image_ref) if auth is not None: diff --git a/src/agent_env/bundle/_secrets.py b/src/agent_env/bundle/_secrets.py index 22f1590c0..1a5a279cf 100644 --- a/src/agent_env/bundle/_secrets.py +++ b/src/agent_env/bundle/_secrets.py @@ -6,6 +6,7 @@ from typing import Any from agent_env.config.describe import is_secret_key +from agent_env.config.loader import SECRET_REF_PREFIX # The fields that map environment variable names to values, where a credential is most often set. _ENV_VAR_MAPS = ("env_vars", "default_env_vars") @@ -58,3 +59,26 @@ def _masked(text: str) -> tuple[str, list[str]]: if count: found.append(what) return text, found + + +def credentials_in(value: Any, path: str = "") -> list[tuple[str, str]]: + """Each credential ``value`` holds written out, by its path, and what it is: an env var whose name looks like a + secret's, set to anything but a ``secret:`` reference, and any value holding a credential ``scrub`` masks.""" + found = [] + items = value.items() if isinstance(value, dict) else enumerate(value) if isinstance(value, list) else () + for key, item in items: + at = f"{path}.{key}" if path and isinstance(key, str) else f"{path}[{key}]" if isinstance(key, int) else key + if isinstance(value, dict) and key in _ENV_VAR_MAPS and isinstance(item, dict): + for name, setting in item.items(): + if not isinstance(setting, str) or not setting or setting.startswith(SECRET_REF_PREFIX): + continue + if is_secret_key(name): + found.append((f"{at}.{name}", "its name looks like a secret's, and it's set")) + elif kinds := _masked(setting)[1]: + found.append((f"{at}.{name}", f"it holds {kinds[0]}")) + elif isinstance(item, str): + if kinds := _masked(item)[1]: + found.append((at, f"it holds {', '.join(kinds)}")) + else: + found += credentials_in(item, at) + return found diff --git a/src/agent_env/bundle/ledger.py b/src/agent_env/bundle/ledger.py index 219c79d97..a6e3c9dba 100644 --- a/src/agent_env/bundle/ledger.py +++ b/src/agent_env/bundle/ledger.py @@ -27,7 +27,7 @@ from agent_env.env.store import ENVS_COLLECTION from agent_env.eval.store import EVALS_COLLECTION from agent_env.store import Filter, Sort -from agent_env.store.document_store import LocalSqliteDocumentStore +from agent_env.store.document_store import DocumentStore, LocalSqliteDocumentStore from agent_env.store.local_state import holding_locks from agent_env.task.store import TASKS_COLLECTION @@ -79,7 +79,11 @@ def next_version(self) -> int: class Ledger: """One plan's rows in the ``@local`` namespace's store.""" - def __init__(self, store: LocalSqliteDocumentStore, plan: Plan): + collection = LEDGER_COLLECTION + # Whether a task or eval is written naming the versions of what it references, so is hashed with them. + pins_references = False + + def __init__(self, store: DocumentStore, plan: Plan): self._store = store self._plan = plan @@ -130,10 +134,10 @@ def digest(self, write: Write, needs: Mapping[tuple[str, str], int]) -> Digest | inputs["files"][key] = _file_sha256(path) + (_mode(path) if built else "") if (published := self._plan.published.get(write.id)) is not None: inputs["published"] = published.pinned - if write.kind not in (BundleKind.TASK, BundleKind.EVAL): + if self.pins_references or write.kind not in (BundleKind.TASK, BundleKind.EVAL): # A document records the versions of what it references (an env its images, a universe its - # environments), so one written anew means it must be written again. A task or eval names its - # references without a version. + # environments), so one written anew means it must be written again. A task or eval a run writes names + # its references without a version. for store, id in write.needs: inputs["needs"][f"{store} {id}"] = str(needs[store, id]) for (kind, id), version in unpinned_store_refs(self._plan, write).items(): @@ -146,36 +150,44 @@ def record(self, check: Check, write: Callable[[], int]) -> int: A pending row marks the attempt until it's done, with its version once ``write`` returns. One an interrupted run of this bundle left is kept by ``check`` when its version can be, and dropped here.""" key = {"store": check.write.kind.store, "id": check.write.id, "bundle": self._bundle(check.write)} - self._store.ensure_index(LEDGER_COLLECTION, ["store", "id", "status"]) - while self._store.delete(LEDGER_COLLECTION, Filter.of(**key, status="pending")): + self._store.ensure_index(self.collection, ["store", "id", "status"]) + while self._store.delete(self.collection, Filter.of(**key, status="pending")): pass pending = {**key, "status": "pending", "scheme": SCHEME, "digest": check.digest.value if check.digest else None, - "inputs": check.digest.inputs if check.digest else None, "at": _now()} - self._store.insert(LEDGER_COLLECTION, pending) + "inputs": check.digest.inputs if check.digest else None, "at": _now(), **self.row_fields()} + self._store.insert(self.collection, pending) version = write() # Stamped before the check below, which can take seconds, so a run interrupted during it leaves the version # this attempt wrote, and the next run keeps it rather than writing it again. - self._store.replace(LEDGER_COLLECTION, Filter.of(**key, status="pending"), {**pending, "version": version}) + self._store.replace(self.collection, Filter.of(**key, status="pending"), {**pending, "version": version}) done = {**pending, "status": "done", "version": version, "at": _now()} if check.digest is not None and self.digest(check.write, check.needs) != check.digest: # A file changed during the write, so what the version holds is unknown: the next run writes it again. done.update(digest=None, inputs=None) - self._store.replace(LEDGER_COLLECTION, Filter.of(**key, status="pending"), done, upsert=True) + self._store.replace(self.collection, Filter.of(**key, status="pending"), done, upsert=True) return version def adopt(self, check: Check) -> None: """Record the version an interrupted run wrote, which ``check`` keeps, as that run would have.""" key = {"store": check.write.kind.store, "id": check.write.id, "bundle": self._bundle(check.write)} done = {**key, "status": "done", "scheme": SCHEME, "digest": check.digest.value, "inputs": check.digest.inputs, - "version": check.version, "at": _now()} - self._store.replace(LEDGER_COLLECTION, Filter.of(**key, status="pending"), done, upsert=True) + "version": check.version, "at": _now(), **self.row_fields()} + self._store.replace(self.collection, Filter.of(**key, status="pending"), done, upsert=True) + + def row_fields(self) -> dict[str, Any]: + """What each row records beyond the write: nothing here.""" + return {} + + def latest_row(self, store: str, id: str) -> dict | None: + """The latest done row for (``store``, ``id``), whichever bundle wrote it.""" + return self._latest(store, id) def _orphan(self, write: Write, stored: int | None, digest: Digest) -> int | None: """The version an interrupted run of this bundle wrote and didn't record, when it's still the store's latest and was made from what ``write`` is now made from.""" - pending = self._find(LEDGER_COLLECTION, Filter.of(store=write.kind.store, id=write.id, - bundle=self._bundle(write), status="pending")) + pending = self._find(self.collection, Filter.of(store=write.kind.store, id=write.id, + bundle=self._bundle(write), status="pending")) if pending is None or pending.get("version") is None or pending["version"] != stored: return None return stored if (pending["scheme"], pending["digest"]) == (SCHEME, digest.value) else None @@ -185,15 +197,15 @@ def _bundle(self, write: Write) -> str: return self._plan.bundle_of(write.source.entry).id_root def _latest(self, store: str, id: str) -> dict | None: - return self._find(LEDGER_COLLECTION, Filter.of(store=store, id=id, status="done")) + return self._find(self.collection, Filter.of(store=store, id=id, status="done")) def _stored_version(self, write: Write) -> int | None: doc = self._find(_COLLECTIONS[write.kind.store], Filter.of(id=write.id)) return doc["version"] if doc else None def _find(self, collection: str, filter: Filter) -> dict | None: - """The highest version matching ``filter``, read without creating the store.""" - if not self._store.path.exists(): + """The highest version matching ``filter``, read without creating a local store.""" + if isinstance(self._store, LocalSqliteDocumentStore) and not self._store.path.exists(): return None return self._store.find_one(collection, filter, sort=Sort.by("version")) diff --git a/src/agent_env/bundle/materialize.py b/src/agent_env/bundle/materialize.py index feffc6310..6e9f0c1cc 100644 --- a/src/agent_env/bundle/materialize.py +++ b/src/agent_env/bundle/materialize.py @@ -24,9 +24,9 @@ from agent_env.a2a_agent import A2AAgent from agent_env.artifact import Artifact -from agent_env.artifact.artifacts.docker_image import DockerImageArtifact +from agent_env.artifact.artifacts.docker_image import DockerImageArtifact, push_context_image from agent_env.artifact.registry import canonical_type, get_artifact_registry -from agent_env.entity.refs import EntityRef, RefRole, ref_sites +from agent_env.entity.refs import EntityKind, EntityRef, RefRole, ref_sites from agent_env.env.registry import get_env_registry from agent_env.eval import Eval, EvalTask from agent_env.providers.sandbox_providers.host_images import build, build_staged @@ -42,7 +42,15 @@ from .ledger import Ledger, locked_ids, materializing from .parse import BundleError, BundleKind from .plan import Plan, Write, artifact_writer, env_writer, unpinned_store_refs -from .resolve import BuiltImage, BundleImage, ReferencedImage, build_step +from .resolve import BuiltImage, BundleImage, ReferencedImage, Reference, build_step + +@dataclass(frozen=True) +class Registering: + """How a registration writes a plan: through ``ledger``, which keeps its rows in the target store, each task and + eval naming the version of each entity it references, and each image the bundle builds built here and pushed.""" + + ledger: Ledger + @dataclass(frozen=True) class Materialized: @@ -79,6 +87,7 @@ def materialize( on_build: Callable[[Write], None] | None = None, on_write: Callable[[Materialized], None] | None = None, build_here: frozenset[str] = frozenset(), + registering: Registering | None = None, ) -> Materialization: """Write ``plan``'s entities, tasks and evals. The first write that fails stops it, and every earlier one stays: they're in the ledger, so the next run reuses them. ``on_wait`` is called when another run @@ -90,12 +99,14 @@ def materialize( ``dry_run`` checks and preflights what a run would, and writes nothing: no entity, ledger row or lock. Each write gets the version it would reuse or the store's next, which another run can take first. A step - reading what the run would write first isn't preflighted, since the store doesn't hold it yet.""" + reading what the run would write first isn't preflighted, since the store doesn't hold it yet. + + ``registering`` writes as a registration does (``Registering``).""" _refuse_unwritable(plan) if not namespace_routing_enabled(): raise RuntimeError("materializing a bundle needs namespace routing, which the agent-env CLI turns on; " "call it inside agent_env.store.routing.namespace_routing()") - ledger = Ledger.for_plan(plan) + ledger = registering.ledger if registering is not None else Ledger.for_plan(plan) entities = [write for write in plan.writes if write.kind not in (BundleKind.TASK, BundleKind.EVAL)] tasks = [write for write in plan.writes if write.kind is BundleKind.TASK] evals = [write for write in plan.writes if write.kind is BundleKind.EVAL] @@ -119,14 +130,18 @@ def through_ledger(write: Write, write_fn: Callable[[], int]) -> None: with nullcontext() if dry_run else materializing(locked_ids(plan), on_wait): for write in entities: here = write.id in build_here and not dry_run - through_ledger(write, lambda: _write_entity(plan, write, here, on_build)) + if registering is not None and isinstance(write.source, BuiltImage) and write.id not in plan.published: + through_ledger(write, lambda: _write_pushed_image(plan, write, on_build)) + else: + through_ledger(write, lambda: _write_entity(plan, write, here, on_build)) if here and done[_key(write)].reused: _build_reused(plan, write, done[_key(write)].version, on_build) unwritten = {_key(item.write) for item in done.values() if not item.reused} if dry_run else set() + versions = {key: item.version for key, item in done.items()} if registering is not None else None built, problems, unchecked = {}, [], [] for write in tasks: with _noted(plan, write, "preflighting"): - built[write.id] = _build(write) + built[write.id] = _build(write, plan, versions) found, skipped = _preflight(plan, write, built[write.id], unwritten) problems.extend(f"{_path(plan, write)}: {problem}" for problem in found) unchecked.extend((write, step) for step in skipped) @@ -135,7 +150,9 @@ def through_ledger(write: Write, write_fn: Callable[[], int]) -> None: for write in tasks: through_ledger(write, lambda: Task.put(id=write.id, steps=built[write.id].steps).version) for write in evals: - through_ledger(write, lambda: _WRITERS[write.kind](plan, write)) + if versions is not None: + versions.update({key: item.version for key, item in done.items()}) + through_ledger(write, lambda: _write_eval(plan, write, versions)) return Materialization(plan, tuple(done[_key(write)] for write in plan.writes), tuple(unchecked)) @@ -170,6 +187,23 @@ def _write_build_context(plan: Plan, write: Write, here: bool, on_build: Callabl raise BundleError([f"{_path(plan, write)}: {e}"]) from None +def _write_pushed_image(plan: Plan, write: Write, on_build: Callable[[Write], None] | None) -> int: + """Register the image an entry's Dockerfile describes by its digest in the configured image store, built here and + pushed unless the store holds it already (``push_context_image``), with its build context beside it.""" + image = write.source + with staged_build_context(plan.bundle_of(image.entry), image.entry, image.dockerfile) as staged: + try: + ref, digest = push_context_image(write.id, str(staged.path), staged.dockerfile, _before(on_build, write)) + return DockerImageArtifact.put_ref( + write.id, description=_image_description(plan, write), image_name=ref, image_digest=digest, + context_path=str(staged.path), dockerfile_path=staged.dockerfile, + ).version + except DockerBuildError as e: + raise BundleError([f"{_path(plan, write)}: {_tail(str(e))}"]) from None + except (ValueError, RuntimeError) as e: # the push, the registry's answer, or the object store + raise BundleError([f"{_path(plan, write)}: {e}"]) from None + + def _write_published_image(plan: Plan, write: Write) -> int: """Register the image published for the sources an entry's Dockerfile builds, by its digest, with its build context beside it, which a sandbox that can't pull the image builds it from; nothing is built here. Refused when the @@ -276,13 +310,28 @@ def _pinned(plan: Plan, write: Write, refs: tuple[EntityRef, ...]) -> Any: return config -def _write_eval(plan: Plan, write: Write) -> int: - """An unpinned task ref, the bundle's own or a store's, is written without a version: the eval runs - its latest.""" - tasks = [EvalTask(ref.id, ref.version) for ref in write.source.references] +def _write_eval(plan: Plan, write: Write, versions: dict[tuple[str, str], int] | None = None) -> int: + """An unpinned task ref, the bundle's own or a store's, is written without a version, so the eval runs its latest; + given the ``versions`` written, as a registration is, each names the version it reads.""" + tasks = [EvalTask(ref.id, ref.version if versions is None else _version(plan, ref, versions)) + for ref in write.source.references] return Eval.put(id=write.id, tasks=tasks).version +def _version(plan: Plan, ref: Reference, versions: dict[tuple[str, str], int]) -> int: + """The version ``ref`` reads: its own, the one written for the bundle's entity it names, or the store's latest the + plan read.""" + if ref.version is not None: + return ref.version + if ref.local is not None: + return versions[_KINDS[ref.kind], ref.id] + return plan.store_latest[ref.kind, ref.id] + + +# The store each kind a reference names is written to, by the key ``done`` holds it under. +_KINDS = {EntityKind.ENV: "env", EntityKind.AGENT: "agent", EntityKind.ARTIFACT: "artifact", EntityKind.TASK: "task"} + + # The writers this release has, by kind; any other kind is refused before anything is written. Tasks are # written separately, once every one of them is preflighted, and evals after them, since they name the tasks. _WRITERS: dict[BundleKind, Callable[[Plan, Write], int]] = { @@ -321,8 +370,25 @@ def _noted(plan: Plan, write: Write, doing: str) -> Iterator[None]: raise -def _build(write: Write) -> Task: - return Task(id=write.id, version=None, steps=[build_step(step) for step in write.source.config]) +def _build(write: Write, plan: Plan, versions: dict[tuple[str, str], int] | None) -> Task: + """The task ``write`` puts: as resolved, or given the ``versions`` written, with each reference that takes a version + naming the one it reads.""" + steps = write.source.config if versions is None else pinned_steps(plan, write, versions) + return Task(id=write.id, version=None, steps=[build_step(step) for step in steps]) + + +def pinned_steps(plan: Plan, write: Write, versions: dict[tuple[str, str], int]) -> list[dict]: + """A copy of the task ``write``'s steps with each reference it reads, the bundle's own entity or a store's, naming + the version it reads (``_version``), where its field takes a version.""" + steps = copy.deepcopy(write.source.config) + references = {(ref.kind, ref.id): ref for ref in write.source.references} + registry = get_task_step_registry() + for step in steps: + for site in ref_sites(registry[step["type"]].entity_refs or (), step): + ref = references.get((site.ref.kind, site.value)) + if site.ref.role is RefRole.INPUT and site.version_key is not None and site.version is None and ref: + site.rewrite(site.value, _version(plan, ref, versions)) + return steps def _preflight(plan: Plan, write: Write, task: Task, diff --git a/src/agent_env/bundle/plan.py b/src/agent_env/bundle/plan.py index 0b56d3908..cb50c8477 100644 --- a/src/agent_env/bundle/plan.py +++ b/src/agent_env/bundle/plan.py @@ -187,13 +187,14 @@ def path(self, entry: BundleEntry) -> str: def plan_bundle(resolved: ResolvedBundle, *, tasks: Sequence[str] = (), evals: Sequence[str] = (), - installed: Installed | None = None, agent: str | None = None) -> Plan: + installed: Installed | None = None, agent: str | None = None, default_agents: bool = True) -> Plan: """Plan a run of ``tasks`` and ``evals``, each a name or id in this bundle. With neither, the run is every eval or, in a bundle without evals, every task. ``installed`` lists the installed bundles whose entities a reference pulls in (``installed_bundles``); without it, an id under one is a store id. ``agent`` names - the agent the solver runs as, as a task's reference would, leaving the tasks as written (``Plan.agent``). Raises - BundleError listing every problem found.""" - return _Planner(resolved, installed).plan(tasks, evals, agent) + the agent the solver runs as, as a task's reference would, leaving the tasks as written (``Plan.agent``). Without + ``default_agents``, a step that leaves its agent unset pulls in no configured default, as for a task written to + run wherever another's default applies. Raises BundleError listing every problem found.""" + return _Planner(resolved, installed, default_agents).plan(tasks, evals, agent) def plan_entities(resolved: ResolvedBundle, ids: Sequence[str], *, installed: Installed | None = None) -> Plan: @@ -220,8 +221,9 @@ class _Pullable: class _Planner: - def __init__(self, resolved: ResolvedBundle, installed: Installed | None = None): + def __init__(self, resolved: ResolvedBundle, installed: Installed | None = None, default_agents: bool = True): self.resolved = resolved + self.default_agents = default_agents self.problems: list[str] = [] self.nodes: dict[_Key, ResolvedEntry | BundleImage] = {} self.rank: dict[_Key, int] = {} @@ -271,7 +273,7 @@ def _localized(self, entry: ResolvedEntry) -> ResolvedEntry: reference to the bundle's own entity points at that; for a task, with a reference for each step that deploys the configured default agent when this bundle or an installed one defines it.""" references = [ref if ref.local is not None else self._pulled(ref) for ref in entry.references] - if entry.entry.kind is BundleKind.TASK: + if entry.entry.kind is BundleKind.TASK and self.default_agents: references += self._default_agents(entry) if references == list(entry.references): return entry @@ -325,7 +327,7 @@ def _resolve_aliases(self) -> dict[str, AliasUse]: aliased[ref.kind, ref.id] = use.target_id ref = self._pulled(replace(ref, id=use.target_id)) references.append(ref) - if _leaves_an_agent_unset(entry): + if self.default_agents and _leaves_an_agent_unset(entry): if default is _UNREAD: default = self._alias_use(self.alias_table, EntityKind.AGENT, get_config().get_default_a2a_agent_id(), None, self.alias_uses, diff --git a/src/agent_env/bundle/register.py b/src/agent_env/bundle/register.py new file mode 100644 index 000000000..cf8739af2 --- /dev/null +++ b/src/agent_env/bundle/register.py @@ -0,0 +1,406 @@ +"""Register a bundle folder into the configured stores under registry ids, so whatever runs a store's tasks runs what the +folder ran. + +A registration plans every task and eval in the folder, or those selected, as a run plans them, but without the +configured default agent a step that leaves its agent unset falls back to: the registered task leaves it unset, so +whatever runs it takes its own default. Each ``@local`` id the plan writes then takes a registry id (``registry_ids``), +and each task and eval is written naming the version of each entity it references, so a registered version never +changes meaning. Each image the bundle builds is built here for linux/amd64 and pushed to the configured image store +under a tag of its sources, with its build context beside it; one published for its sources is registered as published. + +Nothing is written when any problem is found before the first write: an id the rules can't give, two writes given one +id, an id the store holds that this bundle didn't register (unless adopted), an ``@local`` id left in what would be +written, a store reference the store doesn't hold, or, for a store that isn't this machine's, a credential written out +or a URL that only reaches this machine. A dry run checks the same and writes, builds and pushes nothing. Each write +goes through a ledger kept in the target store, so registering again writes only what changed. Two registrations +racing for one id can't both write a version of it: the loser fails, and running it again reuses what the winner wrote. +""" + +from __future__ import annotations + +import re +import subprocess +import unicodedata +from collections.abc import Callable, Iterator, Mapping, Sequence +from dataclasses import dataclass, replace +from importlib.metadata import version as package_version +from pathlib import Path +from typing import Any +from urllib.parse import urlsplit + +from agent_env.config import get_config +from agent_env.entity.refs import RefRole, ref_sites +from agent_env.plugins._bundles import AGENT_ALIASES, ENV_ALIASES +from agent_env.store.document_store import DocumentStore, LocalSqliteDocumentStore +from agent_env.store.ids import LOCAL_PREFIX, is_local_id +from agent_env.store.image_store.oci_registry_credentials import is_loopback_host +from agent_env.store.routing import configured_store, namespace_routing +from agent_env.task_step.registry import get_task_step_registry + +from ._secrets import credentials_in +from .installed import installed_bundles +from .ledger import Ledger +from .materialize import Materialization, Materialized, Registering, materialize +from .parse import Bundle, BundleEntry, BundleError, BundleKind, parse_bundle +from .plan import AliasUse, Plan, Write, plan_bundle, plan_entities +from .run import _label +from .resolve import BundleImage, Reference, ResolvedBundle, ResolvedEntry, resolve_bundle + +REGISTRATIONS = "bundle_registrations" +# A registry id: what a hub route, an object key and an image repository all take. +_REGISTRY_ID = re.compile(r"[a-z0-9][a-z0-9._-]*") +_NOT_IN_ID = re.compile(r"[^a-z0-9._-]+") +_URL = re.compile(r"\b(?:file:///?[^\s\"']*|[a-z][a-z0-9+.-]*://[^\s/\"']+)", re.IGNORECASE) + + +@dataclass(frozen=True) +class Registration: + """What registering a bundle wrote, or in a dry run would write, under which registry ids.""" + + plan: Plan # under the registry ids + materialization: Materialization + ids: Mapping[str, str] # each ``@local`` id the plan writes or its tasks' outputs take, and its registry id + floating: tuple[str, ...] # each reference written without a version, so whatever runs it reads its latest + bundle: str # the bundle's identity, which its ledger rows name + shared: bool # the target is a store other than this machine's own + + +def register_bundle( + root: Path | str, + *, + tasks: Sequence[str] = (), + evals: Sequence[str] = (), + prefix: str | None = None, + adopt: Sequence[str] = (), + dry_run: bool = False, + on_progress: Callable[[str], None] | None = None, +) -> Registration: + """Register the bundle at ``root``: ``tasks`` and ``evals`` by name, or with neither every task and eval it has, + or every entity when it has none, with what they reference. ``prefix`` opens each entry's registry id, the folder's + name and a dash unless given. ``adopt`` names ids the store holds that this bundle may write new versions of though + it didn't register them. ``on_progress`` is given a line for each write, as it's checked or written, and before + each image is built. + + It turns namespace routing on while it runs. Raises BundleError listing every problem found before anything is + written; a write that fails stops it, and every earlier one stays, in the ledger, for the next run to reuse.""" + with namespace_routing(): + bundle = parse_bundle(Path(root)) + resolved = resolve_bundle(bundle) + plan = _plan(resolved, tasks, evals) + ids, problems = registry_ids(plan, _prefix(bundle, prefix)) + if problems: + raise BundleError(problems) + registered = translated(plan, ids) + target = configured_store(get_config().get_document_store()) + shared = shared_store(target) + identity, commit, dirty = bundle_identity(bundle.root) + ledger = RegistrationLedger(target, registered, identity, commit, dirty) + problems = [*_unowned(registered, ledger, frozenset(adopt)), *_unregistrable(registered, shared)] + if problems: + raise BundleError(problems) + say = on_progress or (lambda line: None) + materialization = materialize( + registered, dry_run=dry_run, registering=Registering(ledger), + on_wait=lambda: say("waiting for another agent-env run to finish writing this bundle's ids"), + on_build=lambda write: say(f" {_label(registered, write)}: building with docker and pushing"), + on_write=lambda done: say(f" {registered_line(registered, done)}")) + return Registration(registered, materialization, ids, tuple(_floating(registered)), identity, shared) + + +def shared_store(store: DocumentStore) -> bool: + """Whether ``store`` is other than this machine's own, so others read what's written to it.""" + return not isinstance(store, LocalSqliteDocumentStore) + + +def registered_line(plan: Plan, done: Materialized) -> str: + """A write, as a registration names it: where it's written from, the registry id and version, and why.""" + why = "unchanged" if done.reused else "; ".join(done.reasons) + return f"{_label(plan, done.write)} -> {done.write.kind.store} {done.write.id} v{done.version} ({why})" + + +def _plan(resolved: ResolvedBundle, tasks: Sequence[str], evals: Sequence[str]) -> Plan: + """The plan of ``tasks`` and ``evals``, or of every task and eval, or of every entity of a bundle with none.""" + entries = resolved.bundle.entries + if not tasks and not evals: + tasks = [entry.name for entry in entries if entry.kind is BundleKind.TASK] + evals = [entry.name for entry in entries if entry.kind is BundleKind.EVAL] + if not tasks and not evals: + return plan_entities(resolved, [entry.id for entry in entries], installed=installed_bundles) + return plan_bundle(resolved, tasks=tasks, evals=evals, installed=installed_bundles, default_agents=False) + + +def _prefix(bundle: Bundle, prefix: str | None) -> str: + """``prefix``, or the folder's name as an id and a dash.""" + if prefix is not None: + if prefix and not _REGISTRY_ID.fullmatch(prefix): + raise BundleError([f"--prefix {prefix!r}: an id opens with a lowercase letter or digit and holds only " + "lowercase letters, digits and . _ -"]) + return prefix + name = _slug(bundle.name) + if not name: + raise BundleError([f"{bundle.name!r}: the folder's name gives no id; pass --prefix"]) + return f"{name}-" + + +def registry_ids(plan: Plan, prefix: str) -> tuple[dict[str, str], list[str]]: + """The registry id each ``@local`` id the plan writes takes, and each problem giving them. + + The bundle's own entry is ````, its name lowercased with each run of what an id can't hold made + ``-``. An installed bundle's agent or env pulled in by a reference takes the first alias its bundle gives it, + which the reference named it by. An id derived from another (``__env_image``) keeps its suffix after the + other's registry id, and a task's output is ``-``.""" + ids: dict[str, str] = {} + problems: list[str] = [] + for write in plan.writes: + entry = write.source.entry + if entry.id in ids or (isinstance(write.source, BundleImage) and write.id != entry.id): + continue + if entry in plan.origins: + alias = _alias(plan.origins[entry], entry) + if alias is None: + problems.append(f"{plan.path(entry)}: comes with an installed bundle and has no alias to register it " + "under; name an entity the store holds instead") + continue + ids[entry.id] = alias + elif name := _slug(entry.name): + ids[entry.id] = f"{prefix}{name}" + else: + problems.append(f"{plan.path(entry)}: its name gives no id; rename it") + for write in plan.writes: + if write.id not in ids and (derived := _derived(write.id, ids)) is not None: + ids[write.id] = derived + source = write.source + if isinstance(source, ResolvedEntry) and source.entry.id in ids: + for output in source.outputs: + name = _slug(output.id.removeprefix(f"{source.entry.id}/").rsplit("/", 1)[-1]) + ids[output.id] = f"{ids[source.entry.id]}-{name}" + by_target: dict[str, list[str]] = {} + for local, registry_id in ids.items(): + by_target.setdefault(registry_id, []).append(local) + if not _REGISTRY_ID.fullmatch(registry_id.replace("__", "_")): + problems.append(f"{registry_id!r}: isn't an id a registry takes; rename what it's made from") + for registry_id, locals_ in by_target.items(): + if len(locals_) > 1: + problems.append(f"{', '.join(sorted(locals_))} would all be registered as {registry_id!r}; rename one") + return ids, problems + + +def _alias(bundle: Bundle, entry: BundleEntry) -> str | None: + table = {BundleKind.AGENT: AGENT_ALIASES, BundleKind.ENV: ENV_ALIASES}.get(entry.kind) + aliases = bundle.aliases.get(table, {}).get(entry.name, ()) if table else () + return aliases[0] if aliases else None + + +def _derived(entity_id: str, ids: Mapping[str, str]) -> str | None: + """The registry id of ``entity_id`` when it's derived from an id in ``ids``: that id's, with the same suffix.""" + base = max((known for known in ids if entity_id.startswith(f"{known}__")), key=len, default=None) + return None if base is None else ids[base] + entity_id[len(base):] + + +def _slug(name: str) -> str: + """``name`` lowercased, in ASCII, with each run of what an id can't hold made ``-``, and no ``-``, ``.`` or ``_`` + at its ends.""" + ascii_name = unicodedata.normalize("NFKD", name).encode("ascii", "ignore").decode().lower() + return _NOT_IN_ID.sub("-", ascii_name.replace("__", "-")).strip("-._") + + +def translated(plan: Plan, ids: Mapping[str, str]) -> Plan: + """``plan`` writing each entity under its registry id, and naming it so wherever the plan's documents do.""" + + def entity(entry: BundleEntry) -> BundleEntry: + if entry not in entries: + entries[entry] = replace(entry, id=ids.get(entry.id, entry.id)) + return entries[entry] + + def image(source: BundleImage) -> BundleImage: + return replace(source, id=ids.get(source.id, source.id), entry=entity(source.entry)) + + def reference(ref: Reference) -> Reference: + local = ref.local + if isinstance(local, BundleImage): + local = image(local) + elif isinstance(local, BundleEntry): + local = entity(local) + return replace(ref, id=ids.get(ref.id, ref.id), local=local) + + def resolved(source: ResolvedEntry) -> ResolvedEntry: + return ResolvedEntry(entity(source.entry), _renamed(source.config, ids), + tuple(reference(ref) for ref in source.references), + tuple(replace(output, id=ids.get(output.id, output.id)) for output in source.outputs)) + + def source_of(source: ResolvedEntry | BundleImage) -> ResolvedEntry | BundleImage: + return image(source) if isinstance(source, BundleImage) else resolved(source) + + entries: dict[BundleEntry, BundleEntry] = {} + writes = tuple(Write(write.kind, ids.get(write.id, write.id), source_of(write.source), + tuple((store, ids.get(id, id)) for store, id in write.needs)) for write in plan.writes) + sources = {write.source.entry: write.source for write in writes if isinstance(write.source, ResolvedEntry)} + own = plan.bundle.bundle + bundle = replace(own, entries=tuple(entity(entry) for entry in own.entries)) + resolved_bundle = ResolvedBundle(bundle, tuple(sources.get(entity(entry.entry)) or resolved(entry) + for entry in plan.bundle.entries), + tuple(image(source) for source in plan.bundle.images)) + return replace( + plan, bundle=resolved_bundle, writes=writes, + tasks=tuple(sources[entity(task.entry)] for task in plan.tasks), + evals=tuple(sources[entity(item.entry)] for item in plan.evals), + store_refs=tuple(reference(ref) for ref in plan.store_refs), + origins={entity(entry): origin for entry, origin in plan.origins.items()}, + published={ids.get(id, id): publication for id, publication in plan.published.items()}, + aliases={alias: _aliased(use, ids) for alias, use in plan.aliases.items()}, + env_aliases={alias: _aliased(use, ids) for alias, use in plan.env_aliases.items()}, + ) + + +def _aliased(use: AliasUse, ids: Mapping[str, str]) -> AliasUse: + return replace(use, target_id=ids.get(use.target_id, use.target_id)) + + +def _renamed(value: Any, ids: Mapping[str, str]) -> Any: + """A copy of ``value`` with each string that's an ``@local`` id in ``ids``, or derived from one, its registry id.""" + if isinstance(value, str): + return ids.get(value) or _derived(value, ids) or value if is_local_id(value) else value + if isinstance(value, dict): + return {key: _renamed(item, ids) for key, item in value.items()} + if isinstance(value, list): + return [_renamed(item, ids) for item in value] + return value + + +class RegistrationLedger(Ledger): + """A registration's rows, in the target store: each names the bundle that wrote it, by its identity, and the commit + and state of the folder it was written from. A task or eval is hashed with the versions it names.""" + + collection = REGISTRATIONS + pins_references = True + + def __init__(self, store: DocumentStore, plan: Plan, identity: str, commit: str | None, dirty: bool | None): + super().__init__(store, plan) + self.identity, self.commit, self.dirty = identity, commit, dirty + + def bundle_key(self, write: Write) -> str: + return self._bundle(write) + + def stored_version(self, write: Write) -> int | None: + return self._stored_version(write) + + def _bundle(self, write: Write) -> str: + """The bundle's identity, or for an entity an installed bundle comes with, that bundle's, which every bundle + registering it shares.""" + origin = self._plan.origins.get(write.source.entry) + return f"installed:{origin.id_root.removeprefix(LOCAL_PREFIX)}" if origin is not None else self.identity + + def row_fields(self) -> dict[str, Any]: + return {"commit": self.commit, "dirty": self.dirty, "agent_env": package_version("agentenv-framework")} + + +def bundle_identity(root: Path) -> tuple[str, str | None, bool | None]: + """Who the bundle at ``root`` is, which its ledger rows name: its git remote, without credentials, and its path in + the repo, or else ``folder:`` and its path on this machine; with the commit checked out and whether the folder + holds uncommitted changes, both None outside git.""" + remote = _git(root, "remote", "get-url", "origin") + path_in_repo = _git(root, "rev-parse", "--show-prefix") + if remote is None or path_in_repo is None: + return f"folder:{root}", None, None + status = _git(root, "status", "--porcelain", "--", ".") + return (f"{_remote_name(remote)}:{path_in_repo.rstrip('/')}", _git(root, "rev-parse", "HEAD"), + None if status is None else bool(status)) + + +def _git(root: Path, *args: str) -> str | None: + try: + result = subprocess.run(["git", "-C", str(root), *args], capture_output=True, text=True, timeout=10) + except (OSError, subprocess.TimeoutExpired): + return None + return result.stdout.strip() if result.returncode == 0 else None + + +def _remote_name(remote: str) -> str: + """A git remote as ``host/path``, without its scheme, credentials, port or ``.git``.""" + if "://" in remote: + parts = urlsplit(remote) + host, path = parts.hostname or "", parts.path + else: # scp-like, user@host:path + host, _, path = remote.rpartition("@")[2].partition(":") + return f"{host}/{path.strip('/').removesuffix('.git')}" + + +def _unowned(plan: Plan, ledger: RegistrationLedger, adopt: frozenset[str]) -> Iterator[str]: + """A problem for each write whose id the store holds already, unless this bundle registered it or ``adopt`` names + it.""" + for write in plan.writes: + stored = ledger.stored_version(write) + if stored is None or write.id in adopt: + continue + row = ledger.latest_row(write.kind.store, write.id) + if row is not None and row.get("bundle") == ledger.bundle_key(write): + continue + by = f"registered from {row['bundle']}" if row is not None else "not registered from a bundle" + yield (f"{plan.path(write.source.entry)}: {write.id!r} is in the store already, at v{stored}, {by}; pass " + f"--adopt {write.id} to write a new version of it, or --prefix to register this bundle under other ids") + + +def _unregistrable(plan: Plan, shared: bool) -> Iterator[str]: + """A problem for each ``@local`` id left in what the plan writes and, for a ``shared`` store, each credential + written out and each URL that reaches only this machine.""" + for write in plan.writes: + if not isinstance(write.source, ResolvedEntry): + continue + where = plan.path(write.source.entry) + for path, text in _strings(write.source.config): + if is_local_id(text): + yield f"{where}: {path or 'its value'} names {text!r}, which only this machine's stores hold" + elif shared and (url := _local_url(text)) is not None: + yield f"{where}: {path or 'its value'} names {url}, which only this machine reaches" + if shared: + for path, why in credentials_in(write.source.config): + yield (f"{where}: {path}: {why}, and a shared store shows it to anyone who reads it; leave it out, for " + "whatever runs the task to give it") + + +def _strings(value: Any, path: str = "") -> Iterator[tuple[str, str]]: + if isinstance(value, str): + yield path, value + elif isinstance(value, dict): + for key, item in value.items(): + yield from _strings(item, f"{path}.{key}" if path else str(key)) + elif isinstance(value, list): + for index, item in enumerate(value): + yield from _strings(item, f"{path}[{index}]") + + +def _local_url(text: str) -> str | None: + """The first URL in ``text`` that reaches only this machine: a file, or a loopback host's.""" + for match in _URL.finditer(text): + url = match.group(0) + if url.lower().startswith("file:") or is_loopback_host(urlsplit(url).hostname or ""): + return url + return None + + +def _floating(plan: Plan) -> Iterator[str]: + """Each entity a registered task names in a field that takes no version, so whatever runs it reads its latest.""" + registry = get_task_step_registry() + for write in plan.writes: + if write.kind is not BundleKind.TASK: + continue + for step in write.source.config: + for site in ref_sites(registry[step["type"]].entity_refs or (), step): + if site.ref.role is RefRole.INPUT and site.version_key is None and isinstance(site.value, str): + yield (f"{plan.path(write.source.entry)}: step {step.get('id')!r}: {site.path}: " + f"{site.ref.kind.value} {site.value}") + + +def describe_target() -> list[str]: + """The configured stores a registration writes to, as lines: their kinds, and where each is.""" + config = get_config() + documents = configured_store(config.get_document_store()) + objects = configured_store(config.get_object_store()) + images = configured_store(config.get_image_store()) + where = (str(documents.path) if isinstance(documents, LocalSqliteDocumentStore) + else getattr(getattr(documents, "database", None), "name", None)) + return [f"documents: {type(documents).__name__}" + (f" ({where})" if where else ""), + f"objects: {type(objects).__name__} ({objects.object_url('')})", + f"images: {type(images).__name__}" + (f" ({images.registry_host})" if hasattr(images, "registry_host") + else "")] + diff --git a/src/agent_env/cli/bundle.py b/src/agent_env/cli/bundle.py index c0253bdb2..ec4998c26 100644 --- a/src/agent_env/cli/bundle.py +++ b/src/agent_env/cli/bundle.py @@ -1,4 +1,4 @@ -"""``agent-env bundle``: start a bundle folder from what a store holds.""" +"""``agent-env bundle``: start a bundle folder from what a store holds, and register one into the stores.""" from __future__ import annotations @@ -7,13 +7,15 @@ import click from agent_env.bundle.importer import BundleImport, import_eval, import_task +from agent_env.bundle.parse import BundleKind +from agent_env.bundle.register import describe_target, register_bundle from agent_env.cli.utils import interruptible from agent_env.entity.scopes import FieldScope @click.group() def bundle(): - """Create bundle folders, which ``agent-env run`` runs.""" + """Create bundle folders, which ``agent-env run`` runs, and register them into the configured stores.""" @bundle.command("init") @@ -59,3 +61,57 @@ def _report(imported: BundleImport, folder: Path) -> None: for line in lines: click.echo(f" {line}") click.echo(f"Run it with: agent-env run {folder}") + + +@bundle.command("register") +@click.argument("folder", type=click.Path(exists=True, file_okay=False, path_type=Path)) +@click.option("--task", "tasks", multiple=True, + help="Register this task, by name, with what it uses. Repeatable. With no --task or --eval, every task " + "and eval in the bundle.") +@click.option("--eval", "evals", multiple=True, help="Register this eval, by name, and the tasks it names. Repeatable.") +@click.option("--prefix", default=None, + help="What each entry's registry id opens with; by default the folder's name and a dash.") +@click.option("--adopt", multiple=True, + help="An id the store holds that this bundle didn't register, to write a new version of. Repeatable.") +@click.option("--dry-run", is_flag=True, help="Check everything and say what would be written; write nothing.") +@click.option("--yes", is_flag=True, help="Write to a store other than this machine's own without asking.") +def register(folder: Path, tasks: tuple[str, ...], evals: tuple[str, ...], prefix: str | None, + adopt: tuple[str, ...], dry_run: bool, yes: bool): + """Write the bundle at FOLDER into the configured stores under registry ids, so whatever runs the stores' tasks + runs what the folder runs. + + Each env, agent, artifact, skill, task and eval takes the id PREFIX plus its name, lowercased; a task's outputs + take the task's id and their names. A shipped agent a task names by its alias is registered under the alias when + the store lacks one. Each task and eval names the version of everything it references, so a registered version + never changes meaning; a step that leaves its agent unset is left so, for whatever runs it to give its own. Each + image the bundle builds is built here for linux/amd64 with Docker and pushed to the configured image store, with + its build context beside it. + + It checks everything first and says what it would write: an id the store holds that this bundle didn't register, + an @local id left in what it would write, or, for a store other than this machine's, a credential written out or a + URL only this machine reaches, refuses the whole registration. A store other than this machine's is written only + after you confirm, or with --yes. Registering again writes only what changed since.""" + options = {"tasks": tasks, "evals": evals, "prefix": prefix, "adopt": adopt} + click.echo(f"Registering {folder} into:") + for line in describe_target(): + click.echo(f" {line}") + click.echo("Checking:") + with interruptible(): + planned = register_bundle(folder, dry_run=True, on_progress=click.echo, **options) + for line in planned.floating: + click.echo(f" floating, since its field takes no version: {line}") + if dry_run: + click.echo("Nothing was written: --dry-run.") + return + if all(done.reused for done in planned.materialization.writes): + click.echo("Nothing to write: every version is registered already.") + return + if planned.shared and not yes: + click.confirm("Write these to the stores above?", abort=True) + click.echo("Writing:") + with interruptible(): + registered = register_bundle(folder, on_progress=click.echo, **options) + tasks_written = [done for done in registered.materialization.writes if done.write.kind is BundleKind.TASK] + if tasks_written: + first = tasks_written[0] + click.echo(f"Run a registered task with: agent-env task run --id {first.write.id} --version {first.version}") diff --git a/tst/unit/bundle/register_test.py b/tst/unit/bundle/register_test.py new file mode 100644 index 000000000..e68aa109e --- /dev/null +++ b/tst/unit/bundle/register_test.py @@ -0,0 +1,396 @@ +"""Registering a bundle: each entry under a registry id, each task and eval naming the versions it reads, a ledger in +the target store so registering again writes only what changed, and every problem refused before anything is +written.""" + +import json +import subprocess +from pathlib import Path + +import pytest +from click.testing import CliRunner + +from agent_env.a2a_agent import A2AAgent +from agent_env.artifact import Artifact, FileArtifact +from agent_env.artifact.artifacts import docker_image as docker_image_module +from agent_env.artifact.artifacts.docker_image import DockerImageArtifact, push_context_image +from agent_env.bundle import BundleError +from agent_env.bundle import materialize as materialize_module +from agent_env.bundle import register as register_module +from agent_env.bundle.register import REGISTRATIONS, _remote_name, register_bundle +from agent_env.cli import cli +from agent_env.config import get_config +from agent_env.config.runtime import Config +from agent_env.entity.refs import EntityRef, RefRole +from agent_env.env.env import Env +from agent_env.eval import Eval +from agent_env.store import Filter +from agent_env.task import Task +from agent_env.task_step.task_step import TaskStep +from tst.unit.bundle._support import DIGEST, layout + +BOX = {"id": "box", "type": "deploy_sandbox", "sandbox_name": "box", "sandbox_mode": "vm"} + + +def _load(step_id, artifact): + return {"id": step_id, "type": "load_artifact", "sandbox_name": "box", "artifact_id": artifact} + + +TRIAGE = { + "artifacts/greeting/hello.txt": "hello\n", + "artifacts/docs/a.md": "a\n", + "tasks/t.json": json.dumps([BOX, _load("greet", "greeting"), _load("docs", "docs")]), +} + + +class _Writes(TaskStep): + """A step writing the artifact ``artifact_id``.""" + + type = "writes_register_test" + entity_refs = (EntityRef.artifact("artifact_id", role=RefRole.OUTPUT),) + + def __init__(self, artifact_id=None, **base): + super().__init__(**base) + self.artifact_id = artifact_id + + @classmethod + def from_dict(cls, data): + return cls(**cls._base_from_dict(data), artifact_id=data.get("artifact_id")) + + def to_dict(self): + return {**super().to_dict(), "artifact_id": self.artifact_id} + + async def execute(self, context): + return context + + +@pytest.fixture(autouse=True) +def steps(monkeypatch): + registry = {**Config().task_step_registry(), _Writes.type: _Writes} + monkeypatch.setattr(Config, "task_step_registry", lambda self: registry) + + +_GIT = register_module._git + + +@pytest.fixture(autouse=True) +def outside_git(monkeypatch): + """Each folder is outside git, but the test that says otherwise: none reads a repository it didn't make.""" + monkeypatch.setattr(register_module, "_git", lambda root, *args: None) + + +@pytest.fixture +def home(local_stores, tmp_path, monkeypatch): + monkeypatch.setenv("HOME", str(tmp_path)) + return tmp_path + + +def _steps(task_id, version=None): + return Task.get(task_id, version).to_dict()["steps"] + + +def _rows(): + return get_config().get_document_store().query(REGISTRATIONS, Filter.of(status="done")) + + +def test_each_entry_takes_its_prefix_and_name_and_each_task_names_the_versions_it_reads(home): + registered = register_bundle(layout(home / "triage", TRIAGE)) + + assert {done.write.id: done.version for done in registered.materialization.writes} == { + "triage-docs": 1, "triage-greeting": 1, "triage-t": 1} + steps = _steps("triage-t") + assert [step["artifacts"] for step in steps[1:]] == [[{"id": "triage-greeting", "version": 1}], + [{"id": "triage-docs", "version": 1}]] + assert Artifact.get("triage-docs").version == 1 + + +def test_registering_again_writes_nothing_and_an_edit_rewrites_only_what_it_reaches(home): + root = layout(home / "triage", TRIAGE) + register_bundle(root) + + again = register_bundle(root) + (root / "artifacts/docs/a.md").write_text("changed\n") + edited = register_bundle(root) + + assert all(done.reused for done in again.materialization.writes) + assert {done.write.id: (done.version, done.reused) for done in edited.materialization.writes} == { + "triage-docs": (2, False), "triage-greeting": (1, True), "triage-t": (2, False)} + assert _steps("triage-t")[2]["artifacts"] == [{"id": "triage-docs", "version": 2}] + assert _steps("triage-t", 1)[2]["artifacts"] == [{"id": "triage-docs", "version": 1}] + + +def test_a_dry_run_says_what_it_would_write_and_writes_nothing(home): + planned = register_bundle(layout(home / "triage", TRIAGE), dry_run=True) + + assert {done.write.id: done.version for done in planned.materialization.writes} == { + "triage-docs": 1, "triage-greeting": 1, "triage-t": 1} + assert get_config().get_document_store().count("tasks", Filter()) == 0 and _rows() == [] + + +def test_the_ledger_rows_are_in_the_target_store_and_name_the_bundle_by_its_remote_without_credentials( + local_stores, tmp_path, monkeypatch): + monkeypatch.setattr(register_module, "_git", _GIT) + repo = tmp_path / "repo" + layout(repo / "bundles" / "triage", TRIAGE) + for args in (["init", "-q"], ["remote", "add", "origin", "https://user:token123@github.com/team/repo.git"], + ["add", "."], ["-c", "user.email=a@b.c", "-c", "user.name=a", "commit", "-q", "-m", "x"]): + subprocess.run(["git", "-C", str(repo), *args], check=True) + + registered = register_bundle(repo / "bundles" / "triage") + + assert registered.bundle == "github.com/team/repo:bundles/triage" + rows = _rows() + assert {row["id"] for row in rows} == {"triage-docs", "triage-greeting", "triage-t"} + assert {row["bundle"] for row in rows} == {"github.com/team/repo:bundles/triage"} + assert all(row["commit"] and row["dirty"] is False for row in rows) + assert "token123" not in json.dumps(rows, default=str) + + +@pytest.mark.parametrize(("remote", "name"), [ + ("git@github.com:team/repo.git", "github.com/team/repo"), + ("https://x-access-token:ghs_abc@github.com/team/repo", "github.com/team/repo"), + ("ssh://git@gitlab.example.com:2222/group/sub/repo.git", "gitlab.example.com/group/sub/repo"), +]) +def test_a_remote_is_named_by_its_host_and_path_alone(remote, name): + assert _remote_name(remote) == name + + +def test_an_id_the_store_holds_that_the_bundle_didnt_register_is_refused_unless_adopted(home): + (home / "seed.txt").write_text("seed") + FileArtifact.put("triage-docs", description="someone else's", file_path=str(home / "seed.txt")) + root = layout(home / "triage", TRIAGE) + + with pytest.raises(BundleError) as refused: + register_bundle(root) + adopted = register_bundle(root, adopt=["triage-docs"]) + + assert refused.value.problems == ( + "artifacts/docs: 'triage-docs' is in the store already, at v1, not registered from a bundle; pass --adopt " + "triage-docs to write a new version of it, or --prefix to register this bundle under other ids",) + assert {done.write.id: done.version for done in adopted.materialization.writes}["triage-docs"] == 2 + + +def test_another_folder_of_the_same_name_is_refused_the_ids_this_one_registered(home): + register_bundle(layout(home / "a" / "triage", TRIAGE)) + + with pytest.raises(BundleError) as refused: + register_bundle(layout(home / "b" / "triage", TRIAGE)) + + assert len(refused.value.problems) == 3 + assert all(f"registered from folder:{home / 'a' / 'triage'}" in problem for problem in refused.value.problems) + + +def test_a_prefix_registers_under_other_ids(home): + registered = register_bundle(layout(home / "triage", TRIAGE), prefix="team.") + + assert {done.write.id for done in registered.materialization.writes} == {"team.docs", "team.greeting", "team.t"} + + +def test_a_name_an_id_cant_hold_is_made_one_and_names_that_collide_are_refused(home): + files = {"artifacts/My Docs/a.md": "a\n", "artifacts/my-docs/b.md": "b\n", + "tasks/T.json": json.dumps([BOX, _load("one", "My Docs"), _load("two", "my-docs")])} + + with pytest.raises(BundleError) as refused: + register_bundle(layout(home / "Big Triage", files)) + + assert refused.value.problems == ("@local/~/Big Triage/My Docs, @local/~/Big Triage/my-docs would all be " + "registered as 'big-triage-my-docs'; rename one",) + + +def test_a_bad_prefix_is_refused(home): + with pytest.raises(BundleError) as refused: + register_bundle(layout(home / "triage", TRIAGE), prefix="Team/") + + assert "--prefix 'Team/'" in refused.value.problems[0] + + +def test_a_task_output_is_named_after_its_task_and_isnt_pinned(home): + files = {**TRIAGE, "tasks/t.json": json.dumps([BOX, {"id": "make", "type": _Writes.type, "artifact_id": "report"}, + {**_load("read", "report"), "depends_on": [{"task_step_id": "make"}]}])} + + registered = register_bundle(layout(home / "triage", files)) + + assert registered.ids["@local/~/triage/t/report"] == "triage-t-report" + steps = _steps("triage-t") + assert steps[1]["artifact_id"] == "triage-t-report" + assert steps[2]["artifacts"] == [{"id": "triage-t-report", "version": None}] + + +def test_an_eval_names_the_version_of_each_task(home): + files = {**TRIAGE, "evals/nightly.toml": 'tasks = ["t"]\n'} + root = layout(home / "triage", files) + register_bundle(root) + (root / "artifacts/docs/a.md").write_text("changed\n") + + register_bundle(root) + + assert [(task.task_id, task.task_version) for task in Eval.get("triage-nightly").tasks] == [("triage-t", 2)] + assert [(task.task_id, task.task_version) for task in Eval.get("triage-nightly", 1).tasks] == [("triage-t", 1)] + + +def test_a_store_reference_is_pinned_to_the_version_the_plan_read(home): + (home / "seed.txt").write_text("seed") + FileArtifact.put("shared-seed", description="seed", file_path=str(home / "seed.txt")) + FileArtifact.put("shared-seed", description="seed", file_path=str(home / "seed.txt")) + files = {"tasks/t.json": json.dumps([BOX, _load("seed", "shared-seed")])} + + register_bundle(layout(home / "triage", files)) + + assert _steps("triage-t")[1]["artifacts"] == [{"id": "shared-seed", "version": 2}] + + +def test_a_step_that_leaves_its_agent_unset_is_registered_so_and_writes_no_default_agent(home): + files = {"tasks/t.json": json.dumps([BOX, {"id": "agent", "type": "deploy_agent", "env_ids": []}])} + + registered = register_bundle(layout(home / "triage", files)) + + assert [done.write.id for done in registered.materialization.writes] == ["triage-t"] + assert _steps("triage-t")[1].get("a2a_agent_id") is None + + +def test_a_shipped_agent_named_by_its_alias_is_registered_under_the_alias(home, agents_published): + files = {"tasks/t.json": json.dumps([BOX, {"id": "agent", "type": "deploy_agent", "env_ids": [], + "a2a_agent_id": "claude-code"}])} + + registered = register_bundle(layout(home / "triage", files)) + + written = {done.write.id: done.version for done in registered.materialization.writes} + assert written == {"claude-code__agent_image": 1, "claude-code": 1, "triage-t": 1} + step = _steps("triage-t")[1] + assert (step["a2a_agent_id"], step["a2a_agent_version"]) == ("claude-code", 1) + image = DockerImageArtifact.get("claude-code__agent_image") + assert (image.image_name, image.image_digest) == (agents_published, DIGEST) + assert A2AAgent.get("claude-code").docker_image_artifact.id == "claude-code__agent_image" + + +def test_an_image_the_bundle_builds_is_pushed_and_registered_by_its_digest_with_its_context(home, monkeypatch): + pushed = [] + + def push(entity_id, context_path, dockerfile_path, before=None): + pushed.append((entity_id, dockerfile_path)) + return f"registry.example.com/{entity_id}:src-1", DIGEST + + monkeypatch.setattr(materialize_module, "push_context_image", push) + files = {"envs/crm/Dockerfile": "FROM python:3.12-slim\n", "envs/crm/env.toml": 'environment_name = "crm"\n', + "tasks/t.json": json.dumps([{"id": "env", "type": "deploy_env", "env_id": "crm"}])} + + register_bundle(layout(home / "triage", files)) + + assert [(entity_id, Path(dockerfile).name) for entity_id, dockerfile in pushed] == [ + ("triage-crm__env_image", "Dockerfile")] + image = DockerImageArtifact.get("triage-crm__env_image") + assert (image.image_name, image.image_digest) == ("registry.example.com/triage-crm__env_image:src-1", DIGEST) + assert image.build_context_object_url and image.dockerfile_path == "Dockerfile" + assert Env.get("triage-crm").docker_image_artifact.id == "triage-crm__env_image" + step = _steps("triage-t")[0] + assert (step["env_id"], step["env_version"]) == ("triage-crm", 1) + + +def test_an_at_local_id_left_in_what_would_be_written_is_refused(home): + files = {"tasks/t.json": json.dumps([{**BOX, "metadata": {"from": "@local/~/elsewhere/x"}}])} + + with pytest.raises(BundleError) as refused: + register_bundle(layout(home / "triage", files)) + + assert refused.value.problems == ( + "tasks/t.json: [0].metadata.from names '@local/~/elsewhere/x', which only this machine's stores hold",) + + +def test_a_shared_store_refuses_a_credential_written_out_and_a_url_only_this_machine_reaches(home, monkeypatch): + monkeypatch.setattr(register_module, "shared_store", lambda store: True) + files = {"tasks/t.json": json.dumps([ + {**BOX, "env_vars": {"API_KEY": "abc", "TOKEN": "secret:team_token", "MODE": "fast"}}, + {"id": "agent", "type": "deploy_agent", "env_ids": [], "metadata": {"wiki": "http://localhost:8080/x"}}])} + + with pytest.raises(BundleError) as refused: + register_bundle(layout(home / "triage", files)) + + assert refused.value.problems == ( + "tasks/t.json: [0].env_vars.API_KEY: its name looks like a secret's, and it's set, and a shared store shows " + "it to anyone who reads it; leave it out, for whatever runs the task to give it", + "tasks/t.json: [1].metadata.wiki names http://localhost:8080, which only this machine reaches") + + +def test_the_cli_says_what_it_would_write_asks_before_a_shared_store_and_dry_run_writes_nothing(home, monkeypatch): + root = layout(home / "triage", TRIAGE) + monkeypatch.setattr(register_module, "shared_store", lambda store: True) + + dry = CliRunner().invoke(cli, ["bundle", "register", str(root), "--dry-run"]) + declined = CliRunner().invoke(cli, ["bundle", "register", str(root)], input="n\n") + written = CliRunner().invoke(cli, ["bundle", "register", str(root), "--yes"]) + again = CliRunner().invoke(cli, ["bundle", "register", str(root), "--yes"]) + + assert dry.exit_code == 0, dry.output + assert " artifacts/docs -> artifact triage-docs v1 (new)" in dry.output + assert dry.output.rstrip().endswith("Nothing was written: --dry-run.") + assert declined.exit_code == 1 and "Write these to the stores above?" in declined.output + assert written.exit_code == 0, written.output + assert "Run a registered task with: agent-env task run --id triage-t --version 1" in written.output + assert "Nothing to write: every version is registered already." in again.output + assert [task.version for task in [Task.get("triage-t")]] == [1] + + +def test_a_bundle_with_no_tasks_registers_every_entity(home): + registered = register_bundle(layout(home / "library", {"artifacts/greeting/hello.txt": "hello\n"})) + + assert [done.write.id for done in registered.materialization.writes] == ["library-greeting"] + + +def test_a_failed_write_keeps_what_was_written_and_the_next_run_reuses_it(home, monkeypatch): + root = layout(home / "triage", TRIAGE) + real = Task.put + + def failing(*args, **kwargs): + raise RuntimeError("the store went away") + + monkeypatch.setattr(Task, "put", failing) + with pytest.raises(RuntimeError): + register_bundle(root) + monkeypatch.setattr(Task, "put", real) + + retried = register_bundle(root) + + assert {done.write.id: (done.version, done.reused) for done in retried.materialization.writes} == { + "triage-docs": (1, True), "triage-greeting": (1, True), "triage-t": (1, False)} + assert Artifact.get("triage-docs").version == 1 + + +class _Store: + """An image store that records each repository it's asked to set up.""" + + def __init__(self): + self.repositories = [] + + def image_ref(self, repository, tag): + return f"registry.example.com/{repository}:{tag}" + + def ensure_repository(self, repository): + self.repositories.append(repository) + + def auth(self, ref): + return None + + +def test_an_image_is_built_and_pushed_once_under_the_tag_of_its_sources(home, monkeypatch): + store, built, pushed, served = _Store(), [], [], set() + + def pin(ref, auth): + if ref not in served: + raise docker_image_module.NoSuchImageError(f"{ref}: no such image") + return f"{ref}@{DIGEST}" + + monkeypatch.setattr(get_config(), "get_image_store_for", lambda entity_id: store) + monkeypatch.setattr(docker_image_module, "pin_digest", pin) + monkeypatch.setattr(docker_image_module, "build_image", + lambda dockerfile, context, tag, platform: built.append((tag, platform))) + monkeypatch.setattr(docker_image_module, "_push_local_image", + lambda local, ref, image_store: (pushed.append(ref), served.add(ref))) + context = layout(home / "context", {"Dockerfile": "FROM python:3.12-slim\n"}) + + first = push_context_image("triage-crm__env_image", str(context), "Dockerfile") + again = push_context_image("triage-crm__env_image", str(context), "Dockerfile") + + ref = first[0] + assert first == again == (ref, DIGEST) + assert ref.startswith("registry.example.com/triage-crm-env_image:src-") and "__" not in ref + assert [platform for _, platform in built] == ["linux/amd64"] and pushed == [ref] + assert store.repositories == ["triage-crm-env_image", "triage-crm-env_image"] From 707d7002253e63abd47c0cbafc06dc67fbd18bf0 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Tue, 27 Oct 2026 11:58:30 -0700 Subject: [PATCH 02/12] fix(bundle): a registration that loses a race for a version fails rather than write the next one Two registrations of one bundle on different machines aren't serialized by the local locks, so both wrote every id, the second at v2. The registration ledger now holds each write to the version its check found next (expecting_version): a put of that id then takes exactly that version, through insert_next, or raises ConcurrentModificationError, and a rerun reuses what the winner wrote. Only an agent an installed bundle comes with is registered under its alias: no store env may start on an installed env alias's id, so an installed env is refused before anything is written. Co-Authored-By: Claude Opus 5.5 --- AGENTS.md | 2 +- src/agent_env/artifact/store.py | 4 ++ src/agent_env/bundle/ledger.py | 10 ++-- src/agent_env/bundle/register.py | 18 ++++--- .../store/document_store/__init__.py | 4 ++ .../store/document_store/document_store.py | 30 +++++++++-- tst/unit/bundle/register_test.py | 53 ++++++++++++++++++- 7 files changed, 105 insertions(+), 16 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 651af9ade..7004a1db8 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -114,7 +114,7 @@ and deserialized through a registry. | `a2a_agent/` | The `A2AAgent` entity (`a2a_agent`), its stores and the validator steps. The protocol package provides the agent-side framework. | | `openenv/` | `build_app(bundle)`: the [OpenEnv](https://github.com/openenv/OpenEnv) server for a bundle's tasks, each session playing one episode through the Episode API, with the envs' MCP tools, an added `end_turn` tool and the task's verifier score as the reward. It prepares bundles as `agent-env run` does (`prepare_bundle(played=True)`, without the agent the client plays) and keeps namespace routing on for the app's life. `openenv` comes with the `openenv` extra, a floor with no cap, and this package is the only one that imports it. | | `runner/` | The `[runner]` seam: `Runner.submit()` returns a `RunHandle` (`run_id`, `instance_id`) before the run starts; `LocalRunner` is built in. `runner/store.py` is the run index every runner writes and the explorer reads. | -| `bundle/` | A bundle is a folder of entities (`envs/`, `agents/`, `artifacts/`, `skills/`, `tasks/.json`, `evals/.toml`) that runs without a registry: `agent-env run` plans it, writes it into the `@local` namespace (reusing each version whose inputs haven't changed) and runs its evals or tasks; `agent-env bundle init` writes a stored task or eval as one, with every env, agent and artifact it uses in its own folder (each image as its build context, pruned to what its Dockerfile reads and pinned to the Python and npm versions of the stored image) so it runs with no store. | +| `bundle/` | A bundle is a folder of entities (`envs/`, `agents/`, `artifacts/`, `skills/`, `tasks/.json`, `evals/.toml`) that runs without a registry: `agent-env run` plans it, writes it into the `@local` namespace (reusing each version whose inputs haven't changed) and runs its evals or tasks; `agent-env bundle init` writes a stored task or eval as one, with every env, agent and artifact it uses in its own folder (each image as its build context, pruned to what its Dockerfile reads and pinned to the Python and npm versions of the stored image) so it runs with no store; `agent-env bundle register` writes a folder into the configured stores under registry ids (``), each task and eval naming the versions it reads and each image it builds pushed, through a ledger in the target store so registering again writes only what changed. | | `eval/` | `Eval` (`eval`), a versioned list of task references, and its store. | | `plugins/` | Entry-point plugins (see Extension points): the group names, `inventory()`, which `agent-env plugin list/show/check` report, `load_failures` and `settings`. The top level is public; the underscored modules are internal. | | `explorer/` | Optional local web UI: `agent-env up`, binds loopback `:8234`. | diff --git a/src/agent_env/artifact/store.py b/src/agent_env/artifact/store.py index 716fd01aa..a574afa1b 100644 --- a/src/agent_env/artifact/store.py +++ b/src/agent_env/artifact/store.py @@ -22,6 +22,7 @@ UpdateSpec, VersionedEntityStore, VersionedEntityStoreCache, + expected_version, ) from agent_env.store.ids import key_segment from agent_env.store.query import QueryBuilder, to_document_query @@ -193,6 +194,9 @@ def put_document(self, artifact: _A) -> _A: if artifact.version > 0: versioned.insert_next(doc, artifact.version) return artifact.model_copy() + if (expected := expected_version(artifact.id)) is not None: + versioned.insert_next(doc, expected) + return artifact.model_copy(update={"version": expected}) version = versioned.next_version(artifact.id) last_exc: Optional[DuplicateKeyError] = None for attempt in range(_MAX_VERSION_RETRIES): diff --git a/src/agent_env/bundle/ledger.py b/src/agent_env/bundle/ledger.py index a6e3c9dba..7f05d47fc 100644 --- a/src/agent_env/bundle/ledger.py +++ b/src/agent_env/bundle/ledger.py @@ -13,7 +13,7 @@ import json import stat from collections.abc import Callable, Iterable, Iterator, Mapping -from contextlib import contextmanager +from contextlib import contextmanager, nullcontext from dataclasses import dataclass from datetime import date, datetime, time, timezone from pathlib import Path @@ -27,7 +27,7 @@ from agent_env.env.store import ENVS_COLLECTION from agent_env.eval.store import EVALS_COLLECTION from agent_env.store import Filter, Sort -from agent_env.store.document_store import DocumentStore, LocalSqliteDocumentStore +from agent_env.store.document_store import DocumentStore, LocalSqliteDocumentStore, expecting_version from agent_env.store.local_state import holding_locks from agent_env.task.store import TASKS_COLLECTION @@ -82,6 +82,9 @@ class Ledger: collection = LEDGER_COLLECTION # Whether a task or eval is written naming the versions of what it references, so is hashed with them. pins_references = False + # Whether each write lands on exactly the version its check found next, so one that loses a race to another + # writer raises ConcurrentModificationError rather than add a version after the other's. + strict = False def __init__(self, store: DocumentStore, plan: Plan): self._store = store @@ -157,7 +160,8 @@ def record(self, check: Check, write: Callable[[], int]) -> int: "digest": check.digest.value if check.digest else None, "inputs": check.digest.inputs if check.digest else None, "at": _now(), **self.row_fields()} self._store.insert(self.collection, pending) - version = write() + with expecting_version(check.write.id, check.next_version) if self.strict else nullcontext(): + version = write() # Stamped before the check below, which can take seconds, so a run interrupted during it leaves the version # this attempt wrote, and the next run keeps it rather than writing it again. self._store.replace(self.collection, Filter.of(**key, status="pending"), {**pending, "version": version}) diff --git a/src/agent_env/bundle/register.py b/src/agent_env/bundle/register.py index cf8739af2..da813e22f 100644 --- a/src/agent_env/bundle/register.py +++ b/src/agent_env/bundle/register.py @@ -30,7 +30,7 @@ from agent_env.config import get_config from agent_env.entity.refs import RefRole, ref_sites -from agent_env.plugins._bundles import AGENT_ALIASES, ENV_ALIASES +from agent_env.plugins._bundles import AGENT_ALIASES from agent_env.store.document_store import DocumentStore, LocalSqliteDocumentStore from agent_env.store.ids import LOCAL_PREFIX, is_local_id from agent_env.store.image_store.oci_registry_credentials import is_loopback_host @@ -146,8 +146,9 @@ def registry_ids(plan: Plan, prefix: str) -> tuple[dict[str, str], list[str]]: """The registry id each ``@local`` id the plan writes takes, and each problem giving them. The bundle's own entry is ````, its name lowercased with each run of what an id can't hold made - ``-``. An installed bundle's agent or env pulled in by a reference takes the first alias its bundle gives it, - which the reference named it by. An id derived from another (``__env_image``) keeps its suffix after the + ``-``. An installed bundle's agent pulled in by a reference takes the first alias its bundle gives it, which the + reference named it by; any other entity of an installed bundle has no registry id, and an env's alias can't be + one, since no store env may start on an installed env alias's id. An id derived from another (``__env_image``) keeps its suffix after the other's registry id, and a task's output is ``-``.""" ids: dict[str, str] = {} problems: list[str] = [] @@ -158,8 +159,8 @@ def registry_ids(plan: Plan, prefix: str) -> tuple[dict[str, str], list[str]]: if entry in plan.origins: alias = _alias(plan.origins[entry], entry) if alias is None: - problems.append(f"{plan.path(entry)}: comes with an installed bundle and has no alias to register it " - "under; name an entity the store holds instead") + problems.append(f"{plan.path(entry)}: comes with an installed bundle, and only an agent with an alias " + "is registered from one; name an entity the store holds instead") continue ids[entry.id] = alias elif name := _slug(entry.name): @@ -186,8 +187,7 @@ def registry_ids(plan: Plan, prefix: str) -> tuple[dict[str, str], list[str]]: def _alias(bundle: Bundle, entry: BundleEntry) -> str | None: - table = {BundleKind.AGENT: AGENT_ALIASES, BundleKind.ENV: ENV_ALIASES}.get(entry.kind) - aliases = bundle.aliases.get(table, {}).get(entry.name, ()) if table else () + aliases = bundle.aliases.get(AGENT_ALIASES, {}).get(entry.name, ()) if entry.kind is BundleKind.AGENT else () return aliases[0] if aliases else None @@ -269,10 +269,12 @@ def _renamed(value: Any, ids: Mapping[str, str]) -> Any: class RegistrationLedger(Ledger): """A registration's rows, in the target store: each names the bundle that wrote it, by its identity, and the commit - and state of the folder it was written from. A task or eval is hashed with the versions it names.""" + and state of the folder it was written from. A task or eval is hashed with the versions it names, and each write + lands on exactly the version its check found next.""" collection = REGISTRATIONS pins_references = True + strict = True def __init__(self, store: DocumentStore, plan: Plan, identity: str, commit: str | None, dirty: bool | None): super().__init__(store, plan) diff --git a/src/agent_env/store/document_store/__init__.py b/src/agent_env/store/document_store/__init__.py index 0ec038e3d..1ca2ed869 100644 --- a/src/agent_env/store/document_store/__init__.py +++ b/src/agent_env/store/document_store/__init__.py @@ -31,6 +31,8 @@ VersionedEntityStore, VersionedEntityStoreCache, compare_and_swap, + expected_version, + expecting_version, reject_reserved_id, rev_precondition, ) @@ -46,6 +48,8 @@ "VersionedEntityStore", "VersionedEntityStoreCache", "compare_and_swap", + "expected_version", + "expecting_version", "reject_reserved_id", "rev_precondition", "Filter", diff --git a/src/agent_env/store/document_store/document_store.py b/src/agent_env/store/document_store/document_store.py index 1849a15ca..9c2f56a11 100644 --- a/src/agent_env/store/document_store/document_store.py +++ b/src/agent_env/store/document_store/document_store.py @@ -5,6 +5,9 @@ import json import logging from abc import ABC, abstractmethod +from collections.abc import Iterator, Mapping +from contextlib import contextmanager +from contextvars import ContextVar from dataclasses import dataclass, field from typing import Any, Callable, Generic, Optional, TypeVar, Union @@ -97,6 +100,26 @@ def where(self, field_path: str, predicate: Predicate) -> "Filter": # --- Sort ------------------------------------------------------------------ + +_expected_versions: ContextVar[Mapping[str, int]] = ContextVar("agent_env_expected_versions", default={}) + + +@contextmanager +def expecting_version(entity_id: str, version: int) -> Iterator[None]: + """Inside the block, a put of ``entity_id`` that names no version writes exactly ``version``, which must be the + id's next, as one that names it does: a write that loses a race to another raises ConcurrentModificationError + rather than land on the version after the other's.""" + token = _expected_versions.set({**_expected_versions.get(), entity_id: version}) + try: + yield + finally: + _expected_versions.reset(token) + + +def expected_version(entity_id: str) -> Optional[int]: + """The version ``expecting_version`` holds a put of ``entity_id`` to, or None.""" + return _expected_versions.get().get(entity_id) + @dataclass(frozen=True) class SortKey: field: str @@ -465,9 +488,10 @@ def put(self, entity: T, max_retries: int = 5) -> int: 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"] + version = doc.get("version") or expected_version(doc["id"]) + if version: + self.insert_next(doc, version) + return version last_error: Optional[DuplicateKeyError] = None for _ in range(max_retries): version = self.next_version(doc["id"]) diff --git a/tst/unit/bundle/register_test.py b/tst/unit/bundle/register_test.py index e68aa109e..5b7fb98e4 100644 --- a/tst/unit/bundle/register_test.py +++ b/tst/unit/bundle/register_test.py @@ -23,7 +23,10 @@ from agent_env.entity.refs import EntityRef, RefRole from agent_env.env.env import Env from agent_env.eval import Eval -from agent_env.store import Filter +from agent_env.bundle.parse import BundleKind +from agent_env.store import ConcurrentModificationError, Filter +from agent_env.store.base import NotFoundError +from agent_env.store.document_store import expecting_version from agent_env.task import Task from agent_env.task_step.task_step import TaskStep from tst.unit.bundle._support import DIGEST, layout @@ -394,3 +397,51 @@ def pin(ref, auth): assert ref.startswith("registry.example.com/triage-crm-env_image:src-") and "__" not in ref assert [platform for _, platform in built] == ["linux/amd64"] and pushed == [ref] assert store.repositories == ["triage-crm-env_image", "triage-crm-env_image"] + + +def test_a_put_held_to_a_version_writes_exactly_it_or_fails(home): + (home / "seed.txt").write_text("seed") + + with expecting_version("held", 1): + held = FileArtifact.put("held", description="first", file_path=str(home / "seed.txt")) + with expecting_version("held", 1), pytest.raises(ConcurrentModificationError): + FileArtifact.put("held", description="late", file_path=str(home / "seed.txt")) + with expecting_version("held-task", 2), pytest.raises(ConcurrentModificationError): + Task.put(id="held-task", steps=[]) + + assert held.version == 1 and Artifact.get("held").description == "first" + + +def test_a_registration_that_loses_a_race_for_a_version_fails_rather_than_write_the_next_one(home, monkeypatch): + """As a registration of the same bundle on another machine, which no local lock waits for, would.""" + root = layout(home / "triage", TRIAGE) + writer = materialize_module._WRITERS[BundleKind.ARTIFACT] + + def raced(plan, write): + if write.id == "triage-docs" and not raced.done: + raced.done = True + (home / "other.txt").write_text("the other machine's") + with expecting_version("triage-docs", 1): + FileArtifact.put("triage-docs", description="won the race", file_path=str(home / "other.txt")) + return writer(plan, write) + + raced.done = False + monkeypatch.setitem(materialize_module._WRITERS, BundleKind.ARTIFACT, raced) + + with pytest.raises(ConcurrentModificationError): + register_bundle(root) + + assert Artifact.get("triage-docs", 1).description == "won the race" + with pytest.raises(NotFoundError): + Artifact.get("triage-docs", 2) + + +def test_an_installed_env_a_task_names_by_its_alias_is_refused_before_anything_is_written(home): + files = {"tasks/t.json": json.dumps([{"id": "env", "type": "deploy_env", "env_id": "mail"}])} + + with pytest.raises(BundleError) as refused: + register_bundle(layout(home / "triage", files)) + + assert "comes with an installed bundle, and only an agent with an alias is registered from one" in \ + refused.value.problems[0] + assert get_config().get_document_store().count("envs", Filter()) == 0 From cf01e9417d436806bfabda620d3575a7d9137e3a Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Tue, 27 Oct 2026 12:08:28 -0700 Subject: [PATCH 03/12] fix(bundle): a write that fails leaves no pending ledger row A pending row marks a write a run started, so the next run keeps the version an interrupted one wrote. One whose write raised wrote no version the run knows of, so the row is removed: a registration that lost a race left one behind in the shared store. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/bundle/ledger.py | 9 +++++++-- tst/unit/bundle/register_test.py | 11 +++++++++++ 2 files changed, 18 insertions(+), 2 deletions(-) diff --git a/src/agent_env/bundle/ledger.py b/src/agent_env/bundle/ledger.py index 7f05d47fc..20a180caf 100644 --- a/src/agent_env/bundle/ledger.py +++ b/src/agent_env/bundle/ledger.py @@ -160,8 +160,13 @@ def record(self, check: Check, write: Callable[[], int]) -> int: "digest": check.digest.value if check.digest else None, "inputs": check.digest.inputs if check.digest else None, "at": _now(), **self.row_fields()} self._store.insert(self.collection, pending) - with expecting_version(check.write.id, check.next_version) if self.strict else nullcontext(): - version = write() + try: + with expecting_version(check.write.id, check.next_version) if self.strict else nullcontext(): + version = write() + except Exception: + # It wrote no version this run knows of, so the row marks nothing for the next run to keep. + self._store.delete(self.collection, Filter.of(**key, status="pending")) + raise # Stamped before the check below, which can take seconds, so a run interrupted during it leaves the version # this attempt wrote, and the next run keeps it rather than writing it again. self._store.replace(self.collection, Filter.of(**key, status="pending"), {**pending, "version": version}) diff --git a/tst/unit/bundle/register_test.py b/tst/unit/bundle/register_test.py index 5b7fb98e4..5b147e29b 100644 --- a/tst/unit/bundle/register_test.py +++ b/tst/unit/bundle/register_test.py @@ -445,3 +445,14 @@ def test_an_installed_env_a_task_names_by_its_alias_is_refused_before_anything_i assert "comes with an installed bundle, and only an agent with an alias is registered from one" in \ refused.value.problems[0] assert get_config().get_document_store().count("envs", Filter()) == 0 + + +def test_a_write_that_fails_leaves_no_pending_row(home, monkeypatch): + def failing(*args, **kwargs): + raise RuntimeError("the store went away") + + monkeypatch.setattr(Task, "put", failing) + with pytest.raises(RuntimeError): + register_bundle(layout(home / "triage", TRIAGE)) + + assert get_config().get_document_store().query(REGISTRATIONS, Filter.of(status="pending")) == [] From 229f927892d1bcbd899448f745b3d725f6b1379f Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Tue, 27 Oct 2026 12:49:36 -0700 Subject: [PATCH 04/12] fix(bundle): pin each reference by its own field, resume an interrupted registration, recheck ownership at each write - A step naming no version took another step's pin of the same id; each now takes the version the bundle wrote, or the store's latest the plan read. - A registration killed after a write but before its ledger row was done couldn't resume: the rerun refused the id as unowned. A pending row of this bundle at the stored version now counts as its own, so the ledger adopts it. - Ownership is checked again as each write is checked, so an id another bundle registers after the first check is refused before it's written, not taken at its next version. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/bundle/ledger.py | 4 -- src/agent_env/bundle/materialize.py | 16 +++++--- src/agent_env/bundle/register.py | 57 ++++++++++++++++---------- tst/unit/bundle/register_test.py | 62 +++++++++++++++++++++++++++++ 4 files changed, 109 insertions(+), 30 deletions(-) diff --git a/src/agent_env/bundle/ledger.py b/src/agent_env/bundle/ledger.py index 20a180caf..d97bcc852 100644 --- a/src/agent_env/bundle/ledger.py +++ b/src/agent_env/bundle/ledger.py @@ -188,10 +188,6 @@ def row_fields(self) -> dict[str, Any]: """What each row records beyond the write: nothing here.""" return {} - def latest_row(self, store: str, id: str) -> dict | None: - """The latest done row for (``store``, ``id``), whichever bundle wrote it.""" - return self._latest(store, id) - def _orphan(self, write: Write, stored: int | None, digest: Digest) -> int | None: """The version an interrupted run of this bundle wrote and didn't record, when it's still the store's latest and was made from what ``write`` is now made from.""" diff --git a/src/agent_env/bundle/materialize.py b/src/agent_env/bundle/materialize.py index 6e9f0c1cc..7f42c5049 100644 --- a/src/agent_env/bundle/materialize.py +++ b/src/agent_env/bundle/materialize.py @@ -378,16 +378,20 @@ def _build(write: Write, plan: Plan, versions: dict[tuple[str, str], int] | None def pinned_steps(plan: Plan, write: Write, versions: dict[tuple[str, str], int]) -> list[dict]: - """A copy of the task ``write``'s steps with each reference it reads, the bundle's own entity or a store's, naming - the version it reads (``_version``), where its field takes a version.""" + """A copy of the task ``write``'s steps with each reference it reads that names no version, where its field takes + one, naming the version it reads: the one written for the bundle's own entity, or the store's latest the plan + read.""" steps = copy.deepcopy(write.source.config) - references = {(ref.kind, ref.id): ref for ref in write.source.references} + own = {(ref.kind, ref.id) for ref in write.source.references if ref.local is not None} registry = get_task_step_registry() for step in steps: for site in ref_sites(registry[step["type"]].entity_refs or (), step): - ref = references.get((site.ref.kind, site.value)) - if site.ref.role is RefRole.INPUT and site.version_key is not None and site.version is None and ref: - site.rewrite(site.value, _version(plan, ref, versions)) + if site.ref.role is not RefRole.INPUT or site.version_key is None or site.version is not None: + continue + key = (site.ref.kind, site.value) + version = versions[_KINDS[site.ref.kind], site.value] if key in own else plan.store_latest.get(key) + if version is not None: + site.rewrite(site.value, version) return steps diff --git a/src/agent_env/bundle/register.py b/src/agent_env/bundle/register.py index da813e22f..b2838d6e0 100644 --- a/src/agent_env/bundle/register.py +++ b/src/agent_env/bundle/register.py @@ -30,6 +30,7 @@ from agent_env.config import get_config from agent_env.entity.refs import RefRole, ref_sites +from agent_env.store import Filter from agent_env.plugins._bundles import AGENT_ALIASES from agent_env.store.document_store import DocumentStore, LocalSqliteDocumentStore from agent_env.store.ids import LOCAL_PREFIX, is_local_id @@ -39,7 +40,7 @@ from ._secrets import credentials_in from .installed import installed_bundles -from .ledger import Ledger +from .ledger import Check, Ledger from .materialize import Materialization, Materialized, Registering, materialize from .parse import Bundle, BundleEntry, BundleError, BundleKind, parse_bundle from .plan import AliasUse, Plan, Write, plan_bundle, plan_entities @@ -94,8 +95,8 @@ def register_bundle( target = configured_store(get_config().get_document_store()) shared = shared_store(target) identity, commit, dirty = bundle_identity(bundle.root) - ledger = RegistrationLedger(target, registered, identity, commit, dirty) - problems = [*_unowned(registered, ledger, frozenset(adopt)), *_unregistrable(registered, shared)] + ledger = RegistrationLedger(target, registered, identity, commit, dirty, frozenset(adopt)) + problems = [*_unowned(registered, ledger), *_unregistrable(registered, shared)] if problems: raise BundleError(problems) say = on_progress or (lambda line: None) @@ -276,12 +277,35 @@ class RegistrationLedger(Ledger): pins_references = True strict = True - def __init__(self, store: DocumentStore, plan: Plan, identity: str, commit: str | None, dirty: bool | None): + def __init__(self, store: DocumentStore, plan: Plan, identity: str, commit: str | None, dirty: bool | None, + adopt: frozenset[str] = frozenset()): super().__init__(store, plan) - self.identity, self.commit, self.dirty = identity, commit, dirty - - def bundle_key(self, write: Write) -> str: - return self._bundle(write) + self.identity, self.commit, self.dirty, self.adopted_ids = identity, commit, dirty, adopt + + def check(self, write: Write, needs: Mapping[tuple[str, str], int]) -> Check: + """``Ledger.check``, refusing a write whose id the store holds that this bundle didn't register, checked again + as it's about to be written, since another bundle can write the id after the check before any write.""" + check = super().check(write, needs) + if (problem := self.unowned(write, check.stored)) is not None: + raise BundleError([problem]) + return check + + def unowned(self, write: Write, stored: int | None) -> str | None: + """Why this bundle may not write ``write``'s id, which the store holds at ``stored``, or None when it may: the + store doesn't hold it, ``adopt`` names it, this bundle registered its latest version, or an interrupted + registration of this bundle wrote it.""" + if stored is None or write.id in self.adopted_ids: + return None + row = self._latest(write.kind.store, write.id) + if row is not None and row.get("bundle") == self._bundle(write): + return None + pending = self._find(self.collection, Filter.of(store=write.kind.store, id=write.id, bundle=self._bundle(write), + status="pending")) + if pending is not None and pending.get("version") == stored: + return None + by = f"registered from {row['bundle']}" if row is not None else "not registered from a bundle" + return (f"{self._plan.path(write.source.entry)}: {write.id!r} is in the store already, at v{stored}, {by}; pass " + f"--adopt {write.id} to write a new version of it, or --prefix to register this bundle under other ids") def stored_version(self, write: Write) -> int | None: return self._stored_version(write) @@ -327,19 +351,12 @@ def _remote_name(remote: str) -> str: return f"{host}/{path.strip('/').removesuffix('.git')}" -def _unowned(plan: Plan, ledger: RegistrationLedger, adopt: frozenset[str]) -> Iterator[str]: - """A problem for each write whose id the store holds already, unless this bundle registered it or ``adopt`` names - it.""" +def _unowned(plan: Plan, ledger: RegistrationLedger) -> Iterator[str]: + """A problem for each write this bundle may not make (``RegistrationLedger.unowned``), every one found before any + write.""" for write in plan.writes: - stored = ledger.stored_version(write) - if stored is None or write.id in adopt: - continue - row = ledger.latest_row(write.kind.store, write.id) - if row is not None and row.get("bundle") == ledger.bundle_key(write): - continue - by = f"registered from {row['bundle']}" if row is not None else "not registered from a bundle" - yield (f"{plan.path(write.source.entry)}: {write.id!r} is in the store already, at v{stored}, {by}; pass " - f"--adopt {write.id} to write a new version of it, or --prefix to register this bundle under other ids") + if (problem := ledger.unowned(write, ledger.stored_version(write))) is not None: + yield problem def _unregistrable(plan: Plan, shared: bool) -> Iterator[str]: diff --git a/tst/unit/bundle/register_test.py b/tst/unit/bundle/register_test.py index 5b147e29b..6ca71b66a 100644 --- a/tst/unit/bundle/register_test.py +++ b/tst/unit/bundle/register_test.py @@ -456,3 +456,65 @@ def failing(*args, **kwargs): register_bundle(layout(home / "triage", TRIAGE)) assert get_config().get_document_store().query(REGISTRATIONS, Filter.of(status="pending")) == [] + + +def test_a_step_naming_no_version_takes_the_latest_even_where_another_step_pins_one(home): + (home / "seed.txt").write_text("seed") + FileArtifact.put("shared-seed", description="seed", file_path=str(home / "seed.txt")) + FileArtifact.put("shared-seed", description="seed", file_path=str(home / "seed.txt")) + files = {"tasks/t.json": json.dumps([BOX, _load("latest", "shared-seed"), + {**_load("pinned", "shared-seed"), "artifact_version": 1}])} + + register_bundle(layout(home / "triage", files)) + + assert [step["artifacts"] for step in _steps("triage-t")[1:]] == [[{"id": "shared-seed", "version": 2}], + [{"id": "shared-seed", "version": 1}]] + + +class _Killed(BaseException): + """Stands in for the process dying.""" + + +def test_a_registration_killed_after_a_write_and_before_its_row_is_done_resumes(home, monkeypatch): + root = layout(home / "triage", TRIAGE) + digest = register_module.Ledger.digest + calls = [] + + def dies_after_writing_docs(self, write, needs): + calls.append(write.id) + if write.id == "triage-docs" and calls.count("triage-docs") == 2: # the check's, then the one after the write + raise _Killed + return digest(self, write, needs) + + monkeypatch.setattr(register_module.Ledger, "digest", dies_after_writing_docs) + with pytest.raises(_Killed): + register_bundle(root) + monkeypatch.setattr(register_module.Ledger, "digest", digest) + + resumed = register_bundle(root) + + assert {done.write.id: (done.version, done.reused) for done in resumed.materialization.writes} == { + "triage-docs": (1, True), "triage-greeting": (1, False), "triage-t": (1, False)} + + +def test_an_id_another_bundle_registers_after_the_check_is_refused_before_it_is_written(home, monkeypatch): + root = layout(home / "triage", TRIAGE) + unowned = register_module._unowned + + def then_another_bundle_writes(plan, ledger): + problems = list(unowned(plan, ledger)) + (home / "theirs.txt").write_text("theirs") + FileArtifact.put("triage-docs", description="theirs", file_path=str(home / "theirs.txt")) + get_config().get_document_store().insert(REGISTRATIONS, { + "store": "artifact", "id": "triage-docs", "bundle": "folder:/elsewhere/triage", "status": "done", + "version": 1, "scheme": 2, "digest": None, "inputs": None}) + return iter(problems) + + monkeypatch.setattr(register_module, "_unowned", then_another_bundle_writes) + + with pytest.raises(BundleError) as refused: + register_bundle(root) + + assert "'triage-docs' is in the store already, at v1, registered from folder:/elsewhere/triage" in \ + refused.value.problems[0] + assert Artifact.get("triage-docs").description == "theirs" From d469357ab4314a8f5bc6177f3b78b809ca0bff3a Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Tue, 27 Oct 2026 13:38:39 -0700 Subject: [PATCH 05/12] fix(bundle): don't register over a version someone else wrote since The bundle's own ledger row let it write an id whatever the store's latest was, so a version another tool or person wrote after the last registration was written over unseen. Now an id whose latest version isn't the one this bundle last registered is refused unless --adopt. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/bundle/register.py | 22 +++++++++++++++------- tst/unit/bundle/register_test.py | 16 ++++++++++++++++ 2 files changed, 31 insertions(+), 7 deletions(-) diff --git a/src/agent_env/bundle/register.py b/src/agent_env/bundle/register.py index b2838d6e0..ed139d72a 100644 --- a/src/agent_env/bundle/register.py +++ b/src/agent_env/bundle/register.py @@ -292,20 +292,28 @@ def check(self, write: Write, needs: Mapping[tuple[str, str], int]) -> Check: def unowned(self, write: Write, stored: int | None) -> str | None: """Why this bundle may not write ``write``'s id, which the store holds at ``stored``, or None when it may: the - store doesn't hold it, ``adopt`` names it, this bundle registered its latest version, or an interrupted - registration of this bundle wrote it.""" + store doesn't hold it, ``adopt`` names it, an interrupted registration of this bundle wrote its latest version, + or that's the version this bundle last registered. One the store has written anew since is refused, so no one + else's version is written over unseen.""" if stored is None or write.id in self.adopted_ids: return None - row = self._latest(write.kind.store, write.id) - if row is not None and row.get("bundle") == self._bundle(write): - return None pending = self._find(self.collection, Filter.of(store=write.kind.store, id=write.id, bundle=self._bundle(write), status="pending")) if pending is not None and pending.get("version") == stored: return None + where = self._plan.path(write.source.entry) + row = self._latest(write.kind.store, write.id) + if row is not None and row.get("bundle") == self._bundle(write): + return None if row.get("version") == stored else self._moved(where, write.id, stored, row["version"], + "this bundle last registered") by = f"registered from {row['bundle']}" if row is not None else "not registered from a bundle" - return (f"{self._plan.path(write.source.entry)}: {write.id!r} is in the store already, at v{stored}, {by}; pass " - f"--adopt {write.id} to write a new version of it, or --prefix to register this bundle under other ids") + return (f"{where}: {write.id!r} is in the store already, at v{stored}, {by}; pass --adopt {write.id} to write a " + "new version of it, or --prefix to register this bundle under other ids") + + @staticmethod + def _moved(where: str, entity_id: str, stored: int, own: int, whose: str) -> str: + return (f"{where}: {entity_id!r} is at v{stored} in the store, newer than the v{own} {whose}; pass --adopt " + f"{entity_id} to write over it") def stored_version(self, write: Write) -> int | None: return self._stored_version(write) diff --git a/tst/unit/bundle/register_test.py b/tst/unit/bundle/register_test.py index 6ca71b66a..37002c5f7 100644 --- a/tst/unit/bundle/register_test.py +++ b/tst/unit/bundle/register_test.py @@ -518,3 +518,19 @@ def then_another_bundle_writes(plan, ledger): assert "'triage-docs' is in the store already, at v1, registered from folder:/elsewhere/triage" in \ refused.value.problems[0] assert Artifact.get("triage-docs").description == "theirs" + + +def test_a_version_someone_else_wrote_since_is_not_written_over_unless_adopted(home): + root = layout(home / "triage", TRIAGE) + register_bundle(root) + (home / "seed.txt").write_text("seed") + FileArtifact.put("triage-docs", description="someone else's", file_path=str(home / "seed.txt")) + (root / "artifacts/docs/a.md").write_text("changed\n") + + with pytest.raises(BundleError) as refused: + register_bundle(root) + adopted = register_bundle(root, adopt=["triage-docs"]) + + assert refused.value.problems == ("artifacts/docs: 'triage-docs' is at v2 in the store, newer than the v1 this " + "bundle last registered; pass --adopt triage-docs to write over it",) + assert {done.write.id: done.version for done in adopted.materialization.writes}["triage-docs"] == 3 From d89086f1343c516553183a9f70ec8d7869f728cc Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Tue, 27 Oct 2026 15:04:07 -0700 Subject: [PATCH 06/12] fix(bundle): refuse an env:/secret: reference in a step's env vars, which nothing resolves The shared-store check let a secret-named env var through when its value was a secret: reference, as if the run would resolve it. Nothing resolves env: or secret: references in a step's env vars, so the run would get the reference itself. Such a value is now refused, whatever the variable's name, with a reason saying so. The other refusals now say what each one exposes. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/bundle/_secrets.py | 22 ++++++++++++++-------- src/agent_env/bundle/register.py | 3 +-- tst/unit/bundle/register_test.py | 13 +++++++++---- 3 files changed, 24 insertions(+), 14 deletions(-) diff --git a/src/agent_env/bundle/_secrets.py b/src/agent_env/bundle/_secrets.py index 1a5a279cf..c79262ea8 100644 --- a/src/agent_env/bundle/_secrets.py +++ b/src/agent_env/bundle/_secrets.py @@ -6,7 +6,7 @@ from typing import Any from agent_env.config.describe import is_secret_key -from agent_env.config.loader import SECRET_REF_PREFIX +from agent_env.config.loader import ENV_REF_PREFIX, SECRET_REF_PREFIX # The fields that map environment variable names to values, where a credential is most often set. _ENV_VAR_MAPS = ("env_vars", "default_env_vars") @@ -62,23 +62,29 @@ def _masked(text: str) -> tuple[str, list[str]]: def credentials_in(value: Any, path: str = "") -> list[tuple[str, str]]: - """Each credential ``value`` holds written out, by its path, and what it is: an env var whose name looks like a - secret's, set to anything but a ``secret:`` reference, and any value holding a credential ``scrub`` masks.""" + """Each secret ``value`` holds that a store would keep as written, by its path, and why: an env var whose name looks + like a secret's set to anything, any value holding a credential ``scrub`` masks, and an env var set to an ``env:`` + or ``secret:`` reference, which nothing resolves in a step's env vars, so the run would get the reference itself.""" found = [] items = value.items() if isinstance(value, dict) else enumerate(value) if isinstance(value, list) else () for key, item in items: at = f"{path}.{key}" if path and isinstance(key, str) else f"{path}[{key}]" if isinstance(key, int) else key if isinstance(value, dict) and key in _ENV_VAR_MAPS and isinstance(item, dict): for name, setting in item.items(): - if not isinstance(setting, str) or not setting or setting.startswith(SECRET_REF_PREFIX): + if not isinstance(setting, str) or not setting: continue - if is_secret_key(name): - found.append((f"{at}.{name}", "its name looks like a secret's, and it's set")) + if setting.startswith((ENV_REF_PREFIX, SECRET_REF_PREFIX)): + found.append((f"{at}.{name}", "it's a reference, and nothing resolves one in a step's env vars, so " + "the run would get the reference itself")) + elif is_secret_key(name): + found.append((f"{at}.{name}", "its name looks like a secret's and it's set, which a shared store " + "shows to anyone who reads it")) elif kinds := _masked(setting)[1]: - found.append((f"{at}.{name}", f"it holds {kinds[0]}")) + found.append((f"{at}.{name}", f"it holds {kinds[0]}, which a shared store shows to anyone who " + "reads it")) elif isinstance(item, str): if kinds := _masked(item)[1]: - found.append((at, f"it holds {', '.join(kinds)}")) + found.append((at, f"it holds {', '.join(kinds)}, which a shared store shows to anyone who reads it")) else: found += credentials_in(item, at) return found diff --git a/src/agent_env/bundle/register.py b/src/agent_env/bundle/register.py index ed139d72a..573151931 100644 --- a/src/agent_env/bundle/register.py +++ b/src/agent_env/bundle/register.py @@ -381,8 +381,7 @@ def _unregistrable(plan: Plan, shared: bool) -> Iterator[str]: yield f"{where}: {path or 'its value'} names {url}, which only this machine reaches" if shared: for path, why in credentials_in(write.source.config): - yield (f"{where}: {path}: {why}, and a shared store shows it to anyone who reads it; leave it out, for " - "whatever runs the task to give it") + yield f"{where}: {path}: {why}; leave it out, for whatever runs the task to give it" def _strings(value: Any, path: str = "") -> Iterator[tuple[str, str]]: diff --git a/tst/unit/bundle/register_test.py b/tst/unit/bundle/register_test.py index 37002c5f7..9d51a0168 100644 --- a/tst/unit/bundle/register_test.py +++ b/tst/unit/bundle/register_test.py @@ -298,18 +298,23 @@ def test_an_at_local_id_left_in_what_would_be_written_is_refused(home): "tasks/t.json: [0].metadata.from names '@local/~/elsewhere/x', which only this machine's stores hold",) -def test_a_shared_store_refuses_a_credential_written_out_and_a_url_only_this_machine_reaches(home, monkeypatch): +def test_a_shared_store_refuses_a_credential_or_reference_in_env_vars_and_a_url_only_this_machine_reaches( + home, monkeypatch): monkeypatch.setattr(register_module, "shared_store", lambda store: True) files = {"tasks/t.json": json.dumps([ - {**BOX, "env_vars": {"API_KEY": "abc", "TOKEN": "secret:team_token", "MODE": "fast"}}, + {**BOX, "env_vars": {"API_KEY": "abc", "TOKEN": "secret:team_token", "MODE": "fast", "DB_URL": "env:DB_URL"}}, {"id": "agent", "type": "deploy_agent", "env_ids": [], "metadata": {"wiki": "http://localhost:8080/x"}}])} with pytest.raises(BundleError) as refused: register_bundle(layout(home / "triage", files)) assert refused.value.problems == ( - "tasks/t.json: [0].env_vars.API_KEY: its name looks like a secret's, and it's set, and a shared store shows " - "it to anyone who reads it; leave it out, for whatever runs the task to give it", + "tasks/t.json: [0].env_vars.API_KEY: its name looks like a secret's and it's set, which a shared store shows " + "to anyone who reads it; leave it out, for whatever runs the task to give it", + "tasks/t.json: [0].env_vars.DB_URL: it's a reference, and nothing resolves one in a step's env vars, so the run " + "would get the reference itself; leave it out, for whatever runs the task to give it", + "tasks/t.json: [0].env_vars.TOKEN: it's a reference, and nothing resolves one in a step's env vars, so the run " + "would get the reference itself; leave it out, for whatever runs the task to give it", "tasks/t.json: [1].metadata.wiki names http://localhost:8080, which only this machine reaches") From 0141fbef12b7dc2180845c2392c30c5db6426a6f Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Tue, 27 Oct 2026 15:06:40 -0700 Subject: [PATCH 07/12] test(bundle): join the remote URLs' fake credentials in at run time The tests that a bundle's identity drops a remote's credentials held URLs with user info written out, which the pre-commit secret scan flags as new secrets. Co-Authored-By: Claude Opus 5.5 --- tst/unit/bundle/register_test.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/tst/unit/bundle/register_test.py b/tst/unit/bundle/register_test.py index 9d51a0168..dd46e2972 100644 --- a/tst/unit/bundle/register_test.py +++ b/tst/unit/bundle/register_test.py @@ -32,6 +32,8 @@ from tst.unit.bundle._support import DIGEST, layout BOX = {"id": "box", "type": "deploy_sandbox", "sandbox_name": "box", "sandbox_mode": "vm"} +# Not a credential: what a remote URL's user info holds in the tests that its name drops, joined in at run time. +_FAKE_TOKEN = "not-a-token" def _load(step_id, artifact): @@ -134,7 +136,8 @@ def test_the_ledger_rows_are_in_the_target_store_and_name_the_bundle_by_its_remo monkeypatch.setattr(register_module, "_git", _GIT) repo = tmp_path / "repo" layout(repo / "bundles" / "triage", TRIAGE) - for args in (["init", "-q"], ["remote", "add", "origin", "https://user:token123@github.com/team/repo.git"], + remote = "https://user:" + _FAKE_TOKEN + "@github.com/team/repo.git" + for args in (["init", "-q"], ["remote", "add", "origin", remote], ["add", "."], ["-c", "user.email=a@b.c", "-c", "user.name=a", "commit", "-q", "-m", "x"]): subprocess.run(["git", "-C", str(repo), *args], check=True) @@ -145,12 +148,12 @@ def test_the_ledger_rows_are_in_the_target_store_and_name_the_bundle_by_its_remo assert {row["id"] for row in rows} == {"triage-docs", "triage-greeting", "triage-t"} assert {row["bundle"] for row in rows} == {"github.com/team/repo:bundles/triage"} assert all(row["commit"] and row["dirty"] is False for row in rows) - assert "token123" not in json.dumps(rows, default=str) + assert _FAKE_TOKEN not in json.dumps(rows, default=str) @pytest.mark.parametrize(("remote", "name"), [ ("git@github.com:team/repo.git", "github.com/team/repo"), - ("https://x-access-token:ghs_abc@github.com/team/repo", "github.com/team/repo"), + ("https://x-access-token:" + _FAKE_TOKEN + "@github.com/team/repo", "github.com/team/repo"), ("ssh://git@gitlab.example.com:2222/group/sub/repo.git", "gitlab.example.com/group/sub/repo"), ]) def test_a_remote_is_named_by_its_host_and_path_alone(remote, name): From f69989cdc74e1c8bcba66f0232a447f6cbf6dc5c Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Tue, 27 Oct 2026 15:58:52 -0700 Subject: [PATCH 08/12] fix(bundle): check duplicate registry ids per store, and name the git timeout A task and an artifact whose names make one id live in different stores, so they no longer refuse each other. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/bundle/register.py | 30 ++++++++++++++++++++---------- tst/unit/bundle/register_test.py | 9 +++++++++ 2 files changed, 29 insertions(+), 10 deletions(-) diff --git a/src/agent_env/bundle/register.py b/src/agent_env/bundle/register.py index 573151931..959671d1b 100644 --- a/src/agent_env/bundle/register.py +++ b/src/agent_env/bundle/register.py @@ -8,12 +8,13 @@ changes meaning. Each image the bundle builds is built here for linux/amd64 and pushed to the configured image store under a tag of its sources, with its build context beside it; one published for its sources is registered as published. -Nothing is written when any problem is found before the first write: an id the rules can't give, two writes given one -id, an id the store holds that this bundle didn't register (unless adopted), an ``@local`` id left in what would be -written, a store reference the store doesn't hold, or, for a store that isn't this machine's, a credential written out -or a URL that only reaches this machine. A dry run checks the same and writes, builds and pushes nothing. Each write -goes through a ledger kept in the target store, so registering again writes only what changed. Two registrations -racing for one id can't both write a version of it: the loser fails, and running it again reuses what the winner wrote. +Nothing is written when any problem is found before the first write: an id the rules can't give, two writes to one store +given one id, an id the store holds that this bundle didn't register (unless adopted), an ``@local`` id left in what +would be written, a store reference the store doesn't hold, or, for a store that isn't this machine's, a credential +written out or a URL that only reaches this machine. A dry run checks the same and writes, builds and pushes nothing. +Each write goes through a ledger kept in the target store, so registering again writes only what changed. Two +registrations racing for one id can't both write a version of it: the loser fails, and running it again reuses what the +winner wrote. """ from __future__ import annotations @@ -52,6 +53,7 @@ _REGISTRY_ID = re.compile(r"[a-z0-9][a-z0-9._-]*") _NOT_IN_ID = re.compile(r"[^a-z0-9._-]+") _URL = re.compile(r"\b(?:file:///?[^\s\"']*|[a-z][a-z0-9+.-]*://[^\s/\"']+)", re.IGNORECASE) +_GIT_TIMEOUT_SECONDS = 10 @dataclass(frozen=True) @@ -176,12 +178,19 @@ def registry_ids(plan: Plan, prefix: str) -> tuple[dict[str, str], list[str]]: for output in source.outputs: name = _slug(output.id.removeprefix(f"{source.entry.id}/").rsplit("/", 1)[-1]) ids[output.id] = f"{ids[source.entry.id]}-{name}" - by_target: dict[str, list[str]] = {} + stores: dict[str, set[str]] = {} + for write in plan.writes: + stores.setdefault(write.id, set()).add(write.kind.store) + if isinstance(write.source, ResolvedEntry): + for output in write.source.outputs: + stores.setdefault(output.id, set()).add(BundleKind.ARTIFACT.store) + by_target: dict[tuple[str, str], list[str]] = {} for local, registry_id in ids.items(): - by_target.setdefault(registry_id, []).append(local) + for store in stores.get(local, ()): + by_target.setdefault((store, registry_id), []).append(local) if not _REGISTRY_ID.fullmatch(registry_id.replace("__", "_")): problems.append(f"{registry_id!r}: isn't an id a registry takes; rename what it's made from") - for registry_id, locals_ in by_target.items(): + for (_, registry_id), locals_ in by_target.items(): if len(locals_) > 1: problems.append(f"{', '.join(sorted(locals_))} would all be registered as {registry_id!r}; rename one") return ids, problems @@ -343,7 +352,8 @@ def bundle_identity(root: Path) -> tuple[str, str | None, bool | None]: def _git(root: Path, *args: str) -> str | None: try: - result = subprocess.run(["git", "-C", str(root), *args], capture_output=True, text=True, timeout=10) + result = subprocess.run(["git", "-C", str(root), *args], capture_output=True, text=True, + timeout=_GIT_TIMEOUT_SECONDS) except (OSError, subprocess.TimeoutExpired): return None return result.stdout.strip() if result.returncode == 0 else None diff --git a/tst/unit/bundle/register_test.py b/tst/unit/bundle/register_test.py index dd46e2972..305cb730a 100644 --- a/tst/unit/bundle/register_test.py +++ b/tst/unit/bundle/register_test.py @@ -202,6 +202,15 @@ def test_a_name_an_id_cant_hold_is_made_one_and_names_that_collide_are_refused(h "registered as 'big-triage-my-docs'; rename one",) +def test_entries_whose_names_make_one_id_register_in_their_own_stores(home): + files = {"artifacts/Echo/a.md": "a\n", "tasks/echo.json": json.dumps([BOX, _load("greet", "Echo")])} + + registered = register_bundle(layout(home / "triage", files)) + + assert {(done.write.kind.store, done.write.id) for done in registered.materialization.writes} == { + ("artifact", "triage-echo"), ("task", "triage-echo")} + + def test_a_bad_prefix_is_refused(home): with pytest.raises(BundleError) as refused: register_bundle(layout(home / "triage", TRIAGE), prefix="Team/") From 8e8eba1af4785671dd155aec6077293bc3113478 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Tue, 27 Oct 2026 16:52:01 -0700 Subject: [PATCH 09/12] refactor(bundle): let the ledger say how a registration writes, and read each pinned version one way - materialize takes the ledger instead of a Registering wrapper: Ledger.pins_references now also decides whether tasks and evals are written pinned, and Ledger.pushes_images whether built images are pushed. - Tasks and evals pin through one helper, reading the versions straight from what was written rather than a copy kept in step by hand. An eval ref whose store version the plan didn't read now floats instead of raising, as a task's does. - The kind-to-store table is gone: an EntityKind's value is its store's key. - The registry-id check no longer rewrites "__", which the id pattern already allows. - translated() builds its fields in _translations, beside _KEPT, the fields that hold no @local id; a test fails when a Plan field is in neither. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/bundle/ledger.py | 3 ++ src/agent_env/bundle/materialize.py | 68 ++++++++++++----------------- src/agent_env/bundle/register.py | 41 ++++++++++------- tst/unit/bundle/register_test.py | 14 +++++- 4 files changed, 69 insertions(+), 57 deletions(-) diff --git a/src/agent_env/bundle/ledger.py b/src/agent_env/bundle/ledger.py index d97bcc852..573b86d8a 100644 --- a/src/agent_env/bundle/ledger.py +++ b/src/agent_env/bundle/ledger.py @@ -85,6 +85,9 @@ class Ledger: # Whether each write lands on exactly the version its check found next, so one that loses a race to another # writer raises ConcurrentModificationError rather than add a version after the other's. strict = False + # Whether each image the bundle builds is built here and pushed to the configured image store, for whatever reads the + # store, rather than written as its build context for each deploy to build. + pushes_images = False def __init__(self, store: DocumentStore, plan: Plan): self._store = store diff --git a/src/agent_env/bundle/materialize.py b/src/agent_env/bundle/materialize.py index 7f42c5049..1382a4b7b 100644 --- a/src/agent_env/bundle/materialize.py +++ b/src/agent_env/bundle/materialize.py @@ -17,7 +17,7 @@ from __future__ import annotations import copy -from collections.abc import Callable, Iterator +from collections.abc import Callable, Iterator, Mapping from contextlib import contextmanager, nullcontext from dataclasses import dataclass from typing import Any @@ -42,14 +42,7 @@ from .ledger import Ledger, locked_ids, materializing from .parse import BundleError, BundleKind from .plan import Plan, Write, artifact_writer, env_writer, unpinned_store_refs -from .resolve import BuiltImage, BundleImage, ReferencedImage, Reference, build_step - -@dataclass(frozen=True) -class Registering: - """How a registration writes a plan: through ``ledger``, which keeps its rows in the target store, each task and - eval naming the version of each entity it references, and each image the bundle builds built here and pushed.""" - - ledger: Ledger +from .resolve import BuiltImage, BundleImage, ReferencedImage, build_step @dataclass(frozen=True) @@ -87,7 +80,7 @@ def materialize( on_build: Callable[[Write], None] | None = None, on_write: Callable[[Materialized], None] | None = None, build_here: frozenset[str] = frozenset(), - registering: Registering | None = None, + ledger: Ledger | None = None, ) -> Materialization: """Write ``plan``'s entities, tasks and evals. The first write that fails stops it, and every earlier one stays: they're in the ledger, so the next run reuses them. ``on_wait`` is called when another run @@ -101,12 +94,14 @@ def materialize( Each write gets the version it would reuse or the store's next, which another run can take first. A step reading what the run would write first isn't preflighted, since the store doesn't hold it yet. - ``registering`` writes as a registration does (``Registering``).""" + ``ledger`` keeps the rows, this run's in the ``@local`` namespace's store unless given; it also says whether each + task and eval names the versions it reads (``Ledger.pins_references``) and whether each image the bundle builds is + built here and pushed (``Ledger.pushes_images``), as a registration's does.""" _refuse_unwritable(plan) if not namespace_routing_enabled(): raise RuntimeError("materializing a bundle needs namespace routing, which the agent-env CLI turns on; " "call it inside agent_env.store.routing.namespace_routing()") - ledger = registering.ledger if registering is not None else Ledger.for_plan(plan) + ledger = ledger if ledger is not None else Ledger.for_plan(plan) entities = [write for write in plan.writes if write.kind not in (BundleKind.TASK, BundleKind.EVAL)] tasks = [write for write in plan.writes if write.kind is BundleKind.TASK] evals = [write for write in plan.writes if write.kind is BundleKind.EVAL] @@ -130,18 +125,18 @@ def through_ledger(write: Write, write_fn: Callable[[], int]) -> None: with nullcontext() if dry_run else materializing(locked_ids(plan), on_wait): for write in entities: here = write.id in build_here and not dry_run - if registering is not None and isinstance(write.source, BuiltImage) and write.id not in plan.published: + if ledger.pushes_images and isinstance(write.source, BuiltImage) and write.id not in plan.published: through_ledger(write, lambda: _write_pushed_image(plan, write, on_build)) else: through_ledger(write, lambda: _write_entity(plan, write, here, on_build)) if here and done[_key(write)].reused: _build_reused(plan, write, done[_key(write)].version, on_build) unwritten = {_key(item.write) for item in done.values() if not item.reused} if dry_run else set() - versions = {key: item.version for key, item in done.items()} if registering is not None else None + written = done if ledger.pins_references else None built, problems, unchecked = {}, [], [] for write in tasks: with _noted(plan, write, "preflighting"): - built[write.id] = _build(write, plan, versions) + built[write.id] = _build(write, plan, written) found, skipped = _preflight(plan, write, built[write.id], unwritten) problems.extend(f"{_path(plan, write)}: {problem}" for problem in found) unchecked.extend((write, step) for step in skipped) @@ -150,9 +145,7 @@ def through_ledger(write: Write, write_fn: Callable[[], int]) -> None: for write in tasks: through_ledger(write, lambda: Task.put(id=write.id, steps=built[write.id].steps).version) for write in evals: - if versions is not None: - versions.update({key: item.version for key, item in done.items()}) - through_ledger(write, lambda: _write_eval(plan, write, versions)) + through_ledger(write, lambda: _write_eval(plan, write, written)) return Materialization(plan, tuple(done[_key(write)] for write in plan.writes), tuple(unchecked)) @@ -310,26 +303,20 @@ def _pinned(plan: Plan, write: Write, refs: tuple[EntityRef, ...]) -> Any: return config -def _write_eval(plan: Plan, write: Write, versions: dict[tuple[str, str], int] | None = None) -> int: +def _write_eval(plan: Plan, write: Write, written: Mapping[tuple[str, str], Materialized] | None = None) -> int: """An unpinned task ref, the bundle's own or a store's, is written without a version, so the eval runs its latest; - given the ``versions`` written, as a registration is, each names the version it reads.""" - tasks = [EvalTask(ref.id, ref.version if versions is None else _version(plan, ref, versions)) + given what was ``written``, as a registration is, each names the version it reads.""" + tasks = [EvalTask(ref.id, ref.version if written is None or ref.version is not None + else _read_version(plan, written, ref.kind, ref.id, own=ref.local is not None)) for ref in write.source.references] return Eval.put(id=write.id, tasks=tasks).version -def _version(plan: Plan, ref: Reference, versions: dict[tuple[str, str], int]) -> int: - """The version ``ref`` reads: its own, the one written for the bundle's entity it names, or the store's latest the - plan read.""" - if ref.version is not None: - return ref.version - if ref.local is not None: - return versions[_KINDS[ref.kind], ref.id] - return plan.store_latest[ref.kind, ref.id] - - -# The store each kind a reference names is written to, by the key ``done`` holds it under. -_KINDS = {EntityKind.ENV: "env", EntityKind.AGENT: "agent", EntityKind.ARTIFACT: "artifact", EntityKind.TASK: "task"} +def _read_version(plan: Plan, written: Mapping[tuple[str, str], Materialized], kind: EntityKind, id: str, *, + own: bool) -> int | None: + """The version a reference naming none reads of ``kind`` ``id``: the one ``written`` for the bundle's ``own`` + entity, else the store's latest the plan read, or None when the plan read none.""" + return written[kind.value, id].version if own else plan.store_latest.get((kind, id)) # The writers this release has, by kind; any other kind is refused before anything is written. Tasks are @@ -370,17 +357,16 @@ def _noted(plan: Plan, write: Write, doing: str) -> Iterator[None]: raise -def _build(write: Write, plan: Plan, versions: dict[tuple[str, str], int] | None) -> Task: - """The task ``write`` puts: as resolved, or given the ``versions`` written, with each reference that takes a version +def _build(write: Write, plan: Plan, written: Mapping[tuple[str, str], Materialized] | None) -> Task: + """The task ``write`` puts: as resolved, or given what was ``written``, with each reference that takes a version naming the one it reads.""" - steps = write.source.config if versions is None else pinned_steps(plan, write, versions) + steps = write.source.config if written is None else pinned_steps(plan, write, written) return Task(id=write.id, version=None, steps=[build_step(step) for step in steps]) -def pinned_steps(plan: Plan, write: Write, versions: dict[tuple[str, str], int]) -> list[dict]: +def pinned_steps(plan: Plan, write: Write, written: Mapping[tuple[str, str], Materialized]) -> list[dict]: """A copy of the task ``write``'s steps with each reference it reads that names no version, where its field takes - one, naming the version it reads: the one written for the bundle's own entity, or the store's latest the plan - read.""" + one, naming the version it reads (``_read_version``).""" steps = copy.deepcopy(write.source.config) own = {(ref.kind, ref.id) for ref in write.source.references if ref.local is not None} registry = get_task_step_registry() @@ -388,8 +374,8 @@ def pinned_steps(plan: Plan, write: Write, versions: dict[tuple[str, str], int]) for site in ref_sites(registry[step["type"]].entity_refs or (), step): if site.ref.role is not RefRole.INPUT or site.version_key is None or site.version is not None: continue - key = (site.ref.kind, site.value) - version = versions[_KINDS[site.ref.kind], site.value] if key in own else plan.store_latest.get(key) + own_entity = (site.ref.kind, site.value) in own + version = _read_version(plan, written, site.ref.kind, site.value, own=own_entity) if version is not None: site.rewrite(site.value, version) return steps diff --git a/src/agent_env/bundle/register.py b/src/agent_env/bundle/register.py index 959671d1b..06a3455ff 100644 --- a/src/agent_env/bundle/register.py +++ b/src/agent_env/bundle/register.py @@ -42,7 +42,7 @@ from ._secrets import credentials_in from .installed import installed_bundles from .ledger import Check, Ledger -from .materialize import Materialization, Materialized, Registering, materialize +from .materialize import Materialization, Materialized, materialize from .parse import Bundle, BundleEntry, BundleError, BundleKind, parse_bundle from .plan import AliasUse, Plan, Write, plan_bundle, plan_entities from .run import _label @@ -103,7 +103,7 @@ def register_bundle( raise BundleError(problems) say = on_progress or (lambda line: None) materialization = materialize( - registered, dry_run=dry_run, registering=Registering(ledger), + registered, dry_run=dry_run, ledger=ledger, on_wait=lambda: say("waiting for another agent-env run to finish writing this bundle's ids"), on_build=lambda write: say(f" {_label(registered, write)}: building with docker and pushing"), on_write=lambda done: say(f" {registered_line(registered, done)}")) @@ -188,7 +188,7 @@ def registry_ids(plan: Plan, prefix: str) -> tuple[dict[str, str], list[str]]: for local, registry_id in ids.items(): for store in stores.get(local, ()): by_target.setdefault((store, registry_id), []).append(local) - if not _REGISTRY_ID.fullmatch(registry_id.replace("__", "_")): + if not _REGISTRY_ID.fullmatch(registry_id): problems.append(f"{registry_id!r}: isn't an id a registry takes; rename what it's made from") for (_, registry_id), locals_ in by_target.items(): if len(locals_) > 1: @@ -214,8 +214,18 @@ def _slug(name: str) -> str: return _NOT_IN_ID.sub("-", ascii_name.replace("__", "-")).strip("-._") +# The fields of a Plan that hold no ``@local`` id, so a translated plan keeps them: the store's latest versions, by store +# id, and the agent ``--agent`` names, which a registration plans without. +_KEPT = frozenset({"store_latest", "agent"}) + + def translated(plan: Plan, ids: Mapping[str, str]) -> Plan: """``plan`` writing each entity under its registry id, and naming it so wherever the plan's documents do.""" + return replace(plan, **_translations(plan, ids)) + + +def _translations(plan: Plan, ids: Mapping[str, str]) -> dict[str, Any]: + """Each field of ``plan`` but those it keeps (``_KEPT``), with each ``@local`` id in it its registry id.""" def entity(entry: BundleEntry) -> BundleEntry: if entry not in entries: @@ -250,16 +260,16 @@ def source_of(source: ResolvedEntry | BundleImage) -> ResolvedEntry | BundleImag resolved_bundle = ResolvedBundle(bundle, tuple(sources.get(entity(entry.entry)) or resolved(entry) for entry in plan.bundle.entries), tuple(image(source) for source in plan.bundle.images)) - return replace( - plan, bundle=resolved_bundle, writes=writes, - tasks=tuple(sources[entity(task.entry)] for task in plan.tasks), - evals=tuple(sources[entity(item.entry)] for item in plan.evals), - store_refs=tuple(reference(ref) for ref in plan.store_refs), - origins={entity(entry): origin for entry, origin in plan.origins.items()}, - published={ids.get(id, id): publication for id, publication in plan.published.items()}, - aliases={alias: _aliased(use, ids) for alias, use in plan.aliases.items()}, - env_aliases={alias: _aliased(use, ids) for alias, use in plan.env_aliases.items()}, - ) + return { + "bundle": resolved_bundle, "writes": writes, + "tasks": tuple(sources[entity(task.entry)] for task in plan.tasks), + "evals": tuple(sources[entity(item.entry)] for item in plan.evals), + "store_refs": tuple(reference(ref) for ref in plan.store_refs), + "origins": {entity(entry): origin for entry, origin in plan.origins.items()}, + "published": {ids.get(id, id): publication for id, publication in plan.published.items()}, + "aliases": {alias: _aliased(use, ids) for alias, use in plan.aliases.items()}, + "env_aliases": {alias: _aliased(use, ids) for alias, use in plan.env_aliases.items()}, + } def _aliased(use: AliasUse, ids: Mapping[str, str]) -> AliasUse: @@ -279,12 +289,13 @@ def _renamed(value: Any, ids: Mapping[str, str]) -> Any: class RegistrationLedger(Ledger): """A registration's rows, in the target store: each names the bundle that wrote it, by its identity, and the commit - and state of the folder it was written from. A task or eval is hashed with the versions it names, and each write - lands on exactly the version its check found next.""" + and state of the folder it was written from. A task or eval is hashed with the versions it names, each write lands + on exactly the version its check found next, and each image the bundle builds is built here and pushed.""" collection = REGISTRATIONS pins_references = True strict = True + pushes_images = True def __init__(self, store: DocumentStore, plan: Plan, identity: str, commit: str | None, dirty: bool | None, adopt: frozenset[str] = frozenset()): diff --git a/tst/unit/bundle/register_test.py b/tst/unit/bundle/register_test.py index 305cb730a..136647b27 100644 --- a/tst/unit/bundle/register_test.py +++ b/tst/unit/bundle/register_test.py @@ -4,6 +4,7 @@ import json import subprocess +from dataclasses import fields from pathlib import Path import pytest @@ -23,7 +24,9 @@ from agent_env.entity.refs import EntityRef, RefRole from agent_env.env.env import Env from agent_env.eval import Eval -from agent_env.bundle.parse import BundleKind +from agent_env.bundle.parse import BundleKind, parse_bundle +from agent_env.bundle.plan import Plan, plan_bundle +from agent_env.bundle.resolve import resolve_bundle from agent_env.store import ConcurrentModificationError, Filter from agent_env.store.base import NotFoundError from agent_env.store.document_store import expecting_version @@ -211,6 +214,15 @@ def test_entries_whose_names_make_one_id_register_in_their_own_stores(home): ("artifact", "triage-echo"), ("task", "triage-echo")} +def test_a_translated_plan_renames_every_field_of_a_plan_but_those_that_hold_no_local_id(home): + plan = plan_bundle(resolve_bundle(parse_bundle(layout(home / "triage", TRIAGE))), tasks=["t"]) + + translations = register_module._translations(plan, {}) + + assert translations.keys() | register_module._KEPT == {field.name for field in fields(Plan)} + assert not translations.keys() & register_module._KEPT + + def test_a_bad_prefix_is_refused(home): with pytest.raises(BundleError) as refused: register_bundle(layout(home / "triage", TRIAGE), prefix="Team/") From 8df98e605e8ad070ef455651364891fa385a70d9 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Tue, 27 Oct 2026 22:31:18 -0700 Subject: [PATCH 10/12] test(bundle): register the echo-mcp example against a real registry, and through one that's down Builds and pushes the env's image to the local registry, runs the registered task by its id to a passing score, and checks registering again writes nothing. A registry that can't be reached fails the registration before the image is built or named, and a rerun once it's back finishes. Co-Authored-By: Claude Opus 5.5 --- tst/integration/cli/bundle_register_test.py | 119 ++++++++++++++++++++ 1 file changed, 119 insertions(+) create mode 100644 tst/integration/cli/bundle_register_test.py diff --git a/tst/integration/cli/bundle_register_test.py b/tst/integration/cli/bundle_register_test.py new file mode 100644 index 000000000..e0239a9ac --- /dev/null +++ b/tst/integration/cli/bundle_register_test.py @@ -0,0 +1,119 @@ +"""``agent-env bundle register`` on the echo-mcp example agent-env ships, into the local stores and this machine's local +registry: its env's image is built with Docker and pushed, the registered task runs by its id and passes, and +registering again writes nothing. A registry it can't reach fails the registration before it names the image anywhere, +and once the registry is back a rerun finishes. Needs a Docker daemon and the local registry.""" + +import json +import logging +import shutil +import uuid +from pathlib import Path + +import pytest +from click.testing import CliRunner + +import agent_env +from agent_env.artifact import DockerImageArtifact +from agent_env.artifact.store import reset_artifact_store +from agent_env.cli import cli +from agent_env.config import configure, reset_config +from agent_env.env.bootstrap import GATEWAY, SERVICE_DB, ensure_default_envs +from agent_env.env.env import Env +from agent_env.store.base import NotFoundError +from agent_env.task import Task + +pytestmark = [pytest.mark.integration, pytest.mark.int_test_slow] + +ECHO = Path(agent_env.__file__).parent / "examples" / "echo-mcp" +IMAGE = "envs/echo (Dockerfile image)" +# A registry nothing serves: the plain OCI store, since the local registry's would start one on its port. +UNREACHABLE = """\ +[stores.image] +impl = "agent_env.store.image_store:OciRegistryImageStore" +[stores.image.config] +registry_host = "127.0.0.1:9" +""" + + +@pytest.fixture +def bundle(monkeypatch, tmp_path): + """A copy of the echo-mcp example, with the local stores under this test's folder, and a prefix of its own, so + the registry, which outlives the test, holds nothing under its ids. Whatever ran, no sandbox work folder may be + left behind.""" + sandboxes = tmp_path / "sandboxes" + # HOME stays: docker's credential helper (the macOS keychain, for one) can hang a build's base-image lookup + # when HOME moves. + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("XDG_STATE_HOME", str(tmp_path / "state")) + monkeypatch.setenv("AGENT_ENV_DOCUMENT_STORE", "local") + monkeypatch.setenv("AGENT_ENV_OBJECT_STORE", "local") + monkeypatch.setenv("AGENT_ENV_LOCAL_SANDBOX_DIR", str(sandboxes)) + configure() + reset_artifact_store() + logging.disable(logging.CRITICAL) # pytest's live logging would take CliRunner's stdout + root = tmp_path / "echo" + shutil.copytree(ECHO, root) + try: + yield root, f"it-{uuid.uuid4().hex[:8]}-" + assert not sandboxes.exists() or not any(sandboxes.iterdir()), "a run left a sandbox work folder" + finally: + logging.disable(logging.NOTSET) + reset_artifact_store() + reset_config() + + +def _register(root, prefix): + return CliRunner().invoke(cli, ["bundle", "register", str(root), "--prefix", prefix]) + + +def test_a_registered_bundle_s_image_is_pushed_its_task_runs_by_id_and_registering_again_writes_nothing(bundle, + tmp_path): + root, prefix = bundle + + registered = _register(root, prefix) + + assert registered.exit_code == 0, registered.output + assert registered.output.count("building with docker and pushing") == 1, registered.output + assert f"{IMAGE}: building with docker and pushing" in registered.output, registered.output + image = DockerImageArtifact.get(f"{prefix}echo__env_image") + assert image.image_digest.startswith("sha256:") and image.build_context_object_url, image + assert Env.get(f"{prefix}echo").docker_image_artifact.id == image.id + + # The gateway and service-db envs a stored task's deploy needs, which `agent-env up` sets up in a store this new. + ensure_default_envs((SERVICE_DB, GATEWAY)) + ran = CliRunner().invoke(cli, ["task", "run", "--id", f"{prefix}smoke", "--output-dir", str(tmp_path / "out")]) + again = _register(root, prefix) + + assert ran.exit_code == 0, ran.output + [context] = (tmp_path / "out").iterdir() + assert json.loads(context.read_text())["metadata"]["verifications"]["echoes"]["score"] == 1.0, ran.output + assert again.exit_code == 0, again.output + assert "Nothing to write: every version is registered already." in again.output, again.output + assert "building with docker" not in again.output, again.output + + +def test_a_registry_it_cant_reach_fails_it_before_the_image_is_named_and_a_rerun_once_it_is_back_finishes( + bundle, tmp_path, monkeypatch): + root, prefix = bundle + config = tmp_path / "unreachable.toml" + config.write_text(UNREACHABLE) + monkeypatch.delenv("AGENT_ENV_IMAGE_STORE", raising=False) + with monkeypatch.context() as unreachable: + unreachable.setenv("AGENT_ENV_CONFIG", str(config)) + reset_config() + + failed = _register(root, prefix) + + assert failed.exit_code != 0, failed.output + assert "couldn't read its manifest from 127.0.0.1:9" in failed.output, failed.output + assert "building with docker" not in failed.output, failed.output + for read in (lambda: Env.get(f"{prefix}echo"), lambda: Task.get(f"{prefix}smoke")): + with pytest.raises(NotFoundError): + read() + reset_config() + + finished = _register(root, prefix) + + assert finished.exit_code == 0, finished.output + assert f"{IMAGE}: building with docker and pushing" in finished.output, finished.output + assert Task.get(f"{prefix}smoke").version == 1 From 5bc8d33115a39c76bb9ffc0f4daab9e837a21e13 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Tue, 27 Oct 2026 23:17:21 -0700 Subject: [PATCH 11/12] fix(bundle): translate main's alias_uses, and keep the register integration tests on the local registry main folded Plan's aliases and env_aliases into alias_uses, by kind, so a registered plan renames each of them there. The integration tests now use the local registry whatever AGENT_ENV_CONFIG names, so they never push to a hosted one. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/bundle/register.py | 4 ++-- tst/integration/cli/bundle_register_test.py | 9 +++++---- 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/src/agent_env/bundle/register.py b/src/agent_env/bundle/register.py index 06a3455ff..2cdeb4672 100644 --- a/src/agent_env/bundle/register.py +++ b/src/agent_env/bundle/register.py @@ -267,8 +267,8 @@ def source_of(source: ResolvedEntry | BundleImage) -> ResolvedEntry | BundleImag "store_refs": tuple(reference(ref) for ref in plan.store_refs), "origins": {entity(entry): origin for entry, origin in plan.origins.items()}, "published": {ids.get(id, id): publication for id, publication in plan.published.items()}, - "aliases": {alias: _aliased(use, ids) for alias, use in plan.aliases.items()}, - "env_aliases": {alias: _aliased(use, ids) for alias, use in plan.env_aliases.items()}, + "alias_uses": {kind: {alias: _aliased(use, ids) for alias, use in uses.items()} + for kind, uses in plan.alias_uses.items()}, } diff --git a/tst/integration/cli/bundle_register_test.py b/tst/integration/cli/bundle_register_test.py index e0239a9ac..2e1ddb9c4 100644 --- a/tst/integration/cli/bundle_register_test.py +++ b/tst/integration/cli/bundle_register_test.py @@ -37,9 +37,9 @@ @pytest.fixture def bundle(monkeypatch, tmp_path): - """A copy of the echo-mcp example, with the local stores under this test's folder, and a prefix of its own, so - the registry, which outlives the test, holds nothing under its ids. Whatever ran, no sandbox work folder may be - left behind.""" + """A copy of the echo-mcp example, with the local stores under this test's folder and the local registry, whatever + the config names, and a prefix of its own, so the registry, which outlives the test, holds nothing under its ids. + Whatever ran, no sandbox work folder may be left behind.""" sandboxes = tmp_path / "sandboxes" # HOME stays: docker's credential helper (the macOS keychain, for one) can hang a build's base-image lookup # when HOME moves. @@ -47,6 +47,7 @@ def bundle(monkeypatch, tmp_path): monkeypatch.setenv("XDG_STATE_HOME", str(tmp_path / "state")) monkeypatch.setenv("AGENT_ENV_DOCUMENT_STORE", "local") monkeypatch.setenv("AGENT_ENV_OBJECT_STORE", "local") + monkeypatch.setenv("AGENT_ENV_IMAGE_STORE", "local") monkeypatch.setenv("AGENT_ENV_LOCAL_SANDBOX_DIR", str(sandboxes)) configure() reset_artifact_store() @@ -97,9 +98,9 @@ def test_a_registry_it_cant_reach_fails_it_before_the_image_is_named_and_a_rerun root, prefix = bundle config = tmp_path / "unreachable.toml" config.write_text(UNREACHABLE) - monkeypatch.delenv("AGENT_ENV_IMAGE_STORE", raising=False) with monkeypatch.context() as unreachable: unreachable.setenv("AGENT_ENV_CONFIG", str(config)) + unreachable.delenv("AGENT_ENV_IMAGE_STORE") reset_config() failed = _register(root, prefix) From 530bb76c8d47ec1874d57eee3a2102fa68e5858a Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Wed, 28 Oct 2026 09:35:14 -0700 Subject: [PATCH 12/12] refactor(bundle): drop the ownership check and --adopt from bundle register An id the store holds already now takes the bundle's next version, whoever wrote the store's latest. Two bundles taking one id is unlikely, and the check cost a refusal, its recheck at each write, and --adopt to get past it. The strict version each write lands on stays, so two registrations racing for one id still can't both write it. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/bundle/register.py | 73 ++++++-------------------------- src/agent_env/cli/bundle.py | 14 +++--- tst/unit/bundle/register_test.py | 66 ++++------------------------- 3 files changed, 27 insertions(+), 126 deletions(-) diff --git a/src/agent_env/bundle/register.py b/src/agent_env/bundle/register.py index 2cdeb4672..cb2477bde 100644 --- a/src/agent_env/bundle/register.py +++ b/src/agent_env/bundle/register.py @@ -9,12 +9,11 @@ under a tag of its sources, with its build context beside it; one published for its sources is registered as published. Nothing is written when any problem is found before the first write: an id the rules can't give, two writes to one store -given one id, an id the store holds that this bundle didn't register (unless adopted), an ``@local`` id left in what -would be written, a store reference the store doesn't hold, or, for a store that isn't this machine's, a credential -written out or a URL that only reaches this machine. A dry run checks the same and writes, builds and pushes nothing. -Each write goes through a ledger kept in the target store, so registering again writes only what changed. Two -registrations racing for one id can't both write a version of it: the loser fails, and running it again reuses what the -winner wrote. +given one id, an ``@local`` id left in what would be written, a store reference the store doesn't hold, or, for a store +that isn't this machine's, a credential written out or a URL that only reaches this machine. A dry run checks the same +and writes, builds and pushes nothing. Each write goes through a ledger kept in the target store, so registering again +writes only what changed. Two registrations racing for one id can't both write a version of it: the loser fails, and +running it again reuses what the winner wrote. """ from __future__ import annotations @@ -31,7 +30,6 @@ from agent_env.config import get_config from agent_env.entity.refs import RefRole, ref_sites -from agent_env.store import Filter from agent_env.plugins._bundles import AGENT_ALIASES from agent_env.store.document_store import DocumentStore, LocalSqliteDocumentStore from agent_env.store.ids import LOCAL_PREFIX, is_local_id @@ -41,7 +39,7 @@ from ._secrets import credentials_in from .installed import installed_bundles -from .ledger import Check, Ledger +from .ledger import Ledger from .materialize import Materialization, Materialized, materialize from .parse import Bundle, BundleEntry, BundleError, BundleKind, parse_bundle from .plan import AliasUse, Plan, Write, plan_bundle, plan_entities @@ -74,15 +72,13 @@ def register_bundle( tasks: Sequence[str] = (), evals: Sequence[str] = (), prefix: str | None = None, - adopt: Sequence[str] = (), dry_run: bool = False, on_progress: Callable[[str], None] | None = None, ) -> Registration: """Register the bundle at ``root``: ``tasks`` and ``evals`` by name, or with neither every task and eval it has, or every entity when it has none, with what they reference. ``prefix`` opens each entry's registry id, the folder's - name and a dash unless given. ``adopt`` names ids the store holds that this bundle may write new versions of though - it didn't register them. ``on_progress`` is given a line for each write, as it's checked or written, and before - each image is built. + name and a dash unless given. An id the store holds already takes a new version, whoever wrote it. ``on_progress`` + is given a line for each write, as it's checked or written, and before each image is built. It turns namespace routing on while it runs. Raises BundleError listing every problem found before anything is written; a write that fails stops it, and every earlier one stays, in the ledger, for the next run to reuse.""" @@ -97,8 +93,8 @@ def register_bundle( target = configured_store(get_config().get_document_store()) shared = shared_store(target) identity, commit, dirty = bundle_identity(bundle.root) - ledger = RegistrationLedger(target, registered, identity, commit, dirty, frozenset(adopt)) - problems = [*_unowned(registered, ledger), *_unregistrable(registered, shared)] + ledger = RegistrationLedger(target, registered, identity, commit, dirty) + problems = list(_unregistrable(registered, shared)) if problems: raise BundleError(problems) say = on_progress or (lambda line: None) @@ -297,46 +293,9 @@ class RegistrationLedger(Ledger): strict = True pushes_images = True - def __init__(self, store: DocumentStore, plan: Plan, identity: str, commit: str | None, dirty: bool | None, - adopt: frozenset[str] = frozenset()): + def __init__(self, store: DocumentStore, plan: Plan, identity: str, commit: str | None, dirty: bool | None): super().__init__(store, plan) - self.identity, self.commit, self.dirty, self.adopted_ids = identity, commit, dirty, adopt - - def check(self, write: Write, needs: Mapping[tuple[str, str], int]) -> Check: - """``Ledger.check``, refusing a write whose id the store holds that this bundle didn't register, checked again - as it's about to be written, since another bundle can write the id after the check before any write.""" - check = super().check(write, needs) - if (problem := self.unowned(write, check.stored)) is not None: - raise BundleError([problem]) - return check - - def unowned(self, write: Write, stored: int | None) -> str | None: - """Why this bundle may not write ``write``'s id, which the store holds at ``stored``, or None when it may: the - store doesn't hold it, ``adopt`` names it, an interrupted registration of this bundle wrote its latest version, - or that's the version this bundle last registered. One the store has written anew since is refused, so no one - else's version is written over unseen.""" - if stored is None or write.id in self.adopted_ids: - return None - pending = self._find(self.collection, Filter.of(store=write.kind.store, id=write.id, bundle=self._bundle(write), - status="pending")) - if pending is not None and pending.get("version") == stored: - return None - where = self._plan.path(write.source.entry) - row = self._latest(write.kind.store, write.id) - if row is not None and row.get("bundle") == self._bundle(write): - return None if row.get("version") == stored else self._moved(where, write.id, stored, row["version"], - "this bundle last registered") - by = f"registered from {row['bundle']}" if row is not None else "not registered from a bundle" - return (f"{where}: {write.id!r} is in the store already, at v{stored}, {by}; pass --adopt {write.id} to write a " - "new version of it, or --prefix to register this bundle under other ids") - - @staticmethod - def _moved(where: str, entity_id: str, stored: int, own: int, whose: str) -> str: - return (f"{where}: {entity_id!r} is at v{stored} in the store, newer than the v{own} {whose}; pass --adopt " - f"{entity_id} to write over it") - - def stored_version(self, write: Write) -> int | None: - return self._stored_version(write) + self.identity, self.commit, self.dirty = identity, commit, dirty def _bundle(self, write: Write) -> str: """The bundle's identity, or for an entity an installed bundle comes with, that bundle's, which every bundle @@ -380,14 +339,6 @@ def _remote_name(remote: str) -> str: return f"{host}/{path.strip('/').removesuffix('.git')}" -def _unowned(plan: Plan, ledger: RegistrationLedger) -> Iterator[str]: - """A problem for each write this bundle may not make (``RegistrationLedger.unowned``), every one found before any - write.""" - for write in plan.writes: - if (problem := ledger.unowned(write, ledger.stored_version(write))) is not None: - yield problem - - def _unregistrable(plan: Plan, shared: bool) -> Iterator[str]: """A problem for each ``@local`` id left in what the plan writes and, for a ``shared`` store, each credential written out and each URL that reaches only this machine.""" diff --git a/src/agent_env/cli/bundle.py b/src/agent_env/cli/bundle.py index ec4998c26..79e6adaad 100644 --- a/src/agent_env/cli/bundle.py +++ b/src/agent_env/cli/bundle.py @@ -71,12 +71,10 @@ def _report(imported: BundleImport, folder: Path) -> None: @click.option("--eval", "evals", multiple=True, help="Register this eval, by name, and the tasks it names. Repeatable.") @click.option("--prefix", default=None, help="What each entry's registry id opens with; by default the folder's name and a dash.") -@click.option("--adopt", multiple=True, - help="An id the store holds that this bundle didn't register, to write a new version of. Repeatable.") @click.option("--dry-run", is_flag=True, help="Check everything and say what would be written; write nothing.") @click.option("--yes", is_flag=True, help="Write to a store other than this machine's own without asking.") def register(folder: Path, tasks: tuple[str, ...], evals: tuple[str, ...], prefix: str | None, - adopt: tuple[str, ...], dry_run: bool, yes: bool): + dry_run: bool, yes: bool): """Write the bundle at FOLDER into the configured stores under registry ids, so whatever runs the stores' tasks runs what the folder runs. @@ -87,11 +85,11 @@ def register(folder: Path, tasks: tuple[str, ...], evals: tuple[str, ...], prefi image the bundle builds is built here for linux/amd64 with Docker and pushed to the configured image store, with its build context beside it. - It checks everything first and says what it would write: an id the store holds that this bundle didn't register, - an @local id left in what it would write, or, for a store other than this machine's, a credential written out or a - URL only this machine reaches, refuses the whole registration. A store other than this machine's is written only - after you confirm, or with --yes. Registering again writes only what changed since.""" - options = {"tasks": tasks, "evals": evals, "prefix": prefix, "adopt": adopt} + It checks everything first and says what it would write: an @local id left in what it would write or, for a store + other than this machine's, a credential written out or a URL only this machine reaches refuses the whole + registration. An id the store holds already takes a new version, whoever wrote it. A store other than this + machine's is written only after you confirm, or with --yes. Registering again writes only what changed since.""" + options = {"tasks": tasks, "evals": evals, "prefix": prefix} click.echo(f"Registering {folder} into:") for line in describe_target(): click.echo(f" {line}") diff --git a/tst/unit/bundle/register_test.py b/tst/unit/bundle/register_test.py index 136647b27..702f7b806 100644 --- a/tst/unit/bundle/register_test.py +++ b/tst/unit/bundle/register_test.py @@ -163,29 +163,20 @@ def test_a_remote_is_named_by_its_host_and_path_alone(remote, name): assert _remote_name(remote) == name -def test_an_id_the_store_holds_that_the_bundle_didnt_register_is_refused_unless_adopted(home): +def test_an_id_the_store_holds_takes_the_bundle_s_next_version_whoever_wrote_it(home): (home / "seed.txt").write_text("seed") FileArtifact.put("triage-docs", description="someone else's", file_path=str(home / "seed.txt")) root = layout(home / "triage", TRIAGE) - with pytest.raises(BundleError) as refused: - register_bundle(root) - adopted = register_bundle(root, adopt=["triage-docs"]) - - assert refused.value.problems == ( - "artifacts/docs: 'triage-docs' is in the store already, at v1, not registered from a bundle; pass --adopt " - "triage-docs to write a new version of it, or --prefix to register this bundle under other ids",) - assert {done.write.id: done.version for done in adopted.materialization.writes}["triage-docs"] == 2 - - -def test_another_folder_of_the_same_name_is_refused_the_ids_this_one_registered(home): - register_bundle(layout(home / "a" / "triage", TRIAGE)) - - with pytest.raises(BundleError) as refused: - register_bundle(layout(home / "b" / "triage", TRIAGE)) + first = register_bundle(root) + FileArtifact.put("triage-docs", description="someone else's again", file_path=str(home / "seed.txt")) + (root / "artifacts/docs/a.md").write_text("changed\n") + second = register_bundle(root) + elsewhere = register_bundle(layout(home / "b" / "triage", TRIAGE)) - assert len(refused.value.problems) == 3 - assert all(f"registered from folder:{home / 'a' / 'triage'}" in problem for problem in refused.value.problems) + assert {done.write.id: done.version for done in first.materialization.writes}["triage-docs"] == 2 + assert {done.write.id: done.version for done in second.materialization.writes}["triage-docs"] == 4 + assert {done.write.id: done.version for done in elsewhere.materialization.writes}["triage-docs"] == 5 def test_a_prefix_registers_under_other_ids(home): @@ -524,42 +515,3 @@ def dies_after_writing_docs(self, write, needs): assert {done.write.id: (done.version, done.reused) for done in resumed.materialization.writes} == { "triage-docs": (1, True), "triage-greeting": (1, False), "triage-t": (1, False)} - - -def test_an_id_another_bundle_registers_after_the_check_is_refused_before_it_is_written(home, monkeypatch): - root = layout(home / "triage", TRIAGE) - unowned = register_module._unowned - - def then_another_bundle_writes(plan, ledger): - problems = list(unowned(plan, ledger)) - (home / "theirs.txt").write_text("theirs") - FileArtifact.put("triage-docs", description="theirs", file_path=str(home / "theirs.txt")) - get_config().get_document_store().insert(REGISTRATIONS, { - "store": "artifact", "id": "triage-docs", "bundle": "folder:/elsewhere/triage", "status": "done", - "version": 1, "scheme": 2, "digest": None, "inputs": None}) - return iter(problems) - - monkeypatch.setattr(register_module, "_unowned", then_another_bundle_writes) - - with pytest.raises(BundleError) as refused: - register_bundle(root) - - assert "'triage-docs' is in the store already, at v1, registered from folder:/elsewhere/triage" in \ - refused.value.problems[0] - assert Artifact.get("triage-docs").description == "theirs" - - -def test_a_version_someone_else_wrote_since_is_not_written_over_unless_adopted(home): - root = layout(home / "triage", TRIAGE) - register_bundle(root) - (home / "seed.txt").write_text("seed") - FileArtifact.put("triage-docs", description="someone else's", file_path=str(home / "seed.txt")) - (root / "artifacts/docs/a.md").write_text("changed\n") - - with pytest.raises(BundleError) as refused: - register_bundle(root) - adopted = register_bundle(root, adopt=["triage-docs"]) - - assert refused.value.problems == ("artifacts/docs: 'triage-docs' is at v2 in the store, newer than the v1 this " - "bundle last registered; pass --adopt triage-docs to write over it",) - assert {done.write.id: done.version for done in adopted.materialization.writes}["triage-docs"] == 3