diff --git a/AGENTS.md b/AGENTS.md index 5219364fa..18e66c586 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/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/artifact/store.py b/src/agent_env/artifact/store.py index 042efdaa0..8ee8271fd 100644 --- a/src/agent_env/artifact/store.py +++ b/src/agent_env/artifact/store.py @@ -24,6 +24,7 @@ UpdateSpec, VersionedEntityStore, VersionedEntityStoreCache, + expected_version, ) from agent_env.store.ids import key_segment from agent_env.store.query import QueryBuilder, to_document_query @@ -195,6 +196,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/_secrets.py b/src/agent_env/bundle/_secrets.py index 22f1590c0..c79262ea8 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 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") @@ -58,3 +59,32 @@ 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 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: + continue + 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]}, 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)}, 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/ledger.py b/src/agent_env/bundle/ledger.py index 219c79d97..573b86d8a 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 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 @@ -79,7 +79,17 @@ 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 + # 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 self._plan = plan @@ -130,10 +140,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 +156,46 @@ 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) - version = write() + "inputs": check.digest.inputs if check.digest else None, "at": _now(), **self.row_fields()} + self._store.insert(self.collection, pending) + 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(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 _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 +205,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..1382a4b7b 100644 --- a/src/agent_env/bundle/materialize.py +++ b/src/agent_env/bundle/materialize.py @@ -17,16 +17,16 @@ 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 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 @@ -44,6 +44,7 @@ from .plan import Plan, Write, artifact_writer, env_writer, unpinned_store_refs from .resolve import BuiltImage, BundleImage, ReferencedImage, build_step + @dataclass(frozen=True) class Materialized: """One write: the version it left in the store, or in a dry run would leave, and whether that is the one the @@ -79,6 +80,7 @@ def materialize( on_build: Callable[[Write], None] | None = None, on_write: Callable[[Materialized], None] | None = None, build_here: frozenset[str] = frozenset(), + 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 @@ -90,12 +92,16 @@ 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. + + ``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 = 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] @@ -119,14 +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 - through_ledger(write, lambda: _write_entity(plan, write, here, on_build)) + 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() + 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) + 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) @@ -135,7 +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: - through_ledger(write, lambda: _WRITERS[write.kind](plan, write)) + through_ledger(write, lambda: _write_eval(plan, write, written)) return Materialization(plan, tuple(done[_key(write)] for write in plan.writes), tuple(unchecked)) @@ -170,6 +180,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 +303,22 @@ 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, 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 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 _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 # 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 +357,28 @@ 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, 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 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, 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 (``_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() + for step in steps: + 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 + 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 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 fcb8f5e70..84fb3f8e9 100644 --- a/src/agent_env/bundle/plan.py +++ b/src/agent_env/bundle/plan.py @@ -191,13 +191,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: @@ -224,8 +225,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] = {} @@ -274,7 +276,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 @@ -324,7 +326,7 @@ def _resolve_aliases(self) -> None: 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(EntityKind.AGENT, get_config().get_default_a2a_agent_id(), None) if isinstance(default, str): diff --git a/src/agent_env/bundle/register.py b/src/agent_env/bundle/register.py new file mode 100644 index 000000000..cb2477bde --- /dev/null +++ b/src/agent_env/bundle/register.py @@ -0,0 +1,404 @@ +"""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 to one store +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 + +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 +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, 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) +_GIT_TIMEOUT_SECONDS = 10 + + +@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, + 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. 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.""" + 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 = list(_unregistrable(registered, shared)) + if problems: + raise BundleError(problems) + say = on_progress or (lambda line: None) + materialization = materialize( + 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)}")) + 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 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] = [] + 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 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): + 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}" + 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(): + for store in stores.get(local, ()): + by_target.setdefault((store, registry_id), []).append(local) + 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: + 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: + aliases = bundle.aliases.get(AGENT_ALIASES, {}).get(entry.name, ()) if entry.kind is BundleKind.AGENT 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("-._") + + +# 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: + 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 { + "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()}, + "alias_uses": {kind: {alias: _aliased(use, ids) for alias, use in uses.items()} + for kind, uses in plan.alias_uses.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, 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): + super().__init__(store, plan) + 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 + 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=_GIT_TIMEOUT_SECONDS) + 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 _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}; 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..79e6adaad 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,55 @@ 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("--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, + 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 @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}") + 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/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/integration/cli/bundle_register_test.py b/tst/integration/cli/bundle_register_test.py new file mode 100644 index 000000000..2e1ddb9c4 --- /dev/null +++ b/tst/integration/cli/bundle_register_test.py @@ -0,0 +1,120 @@ +"""``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 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. + 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_IMAGE_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) + with monkeypatch.context() as unreachable: + unreachable.setenv("AGENT_ENV_CONFIG", str(config)) + unreachable.delenv("AGENT_ENV_IMAGE_STORE") + 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 diff --git a/tst/unit/bundle/register_test.py b/tst/unit/bundle/register_test.py new file mode 100644 index 000000000..702f7b806 --- /dev/null +++ b/tst/unit/bundle/register_test.py @@ -0,0 +1,517 @@ +"""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 dataclasses import fields +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.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 +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"} +# 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): + 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) + 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) + + 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 _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:" + _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): + assert _remote_name(remote) == name + + +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) + + 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 {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): + 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_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_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/") + + 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_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", "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, 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") + + +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"] + + +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 + + +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")) == [] + + +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)}