feat(bundle): run --dry-run, another bundle's entities by @local id, bare depends_on ids

src/agent_env/bundle/__init__.py

@@ -1,4 +1,4 @@
 """Bundles: folders of entities that run without a registry."""
 
 from .parse import Bundle, BundleEntry, BundleError, BundleKind, is_ignored, parse_bundle
-from .run import BundleRun, EvalRun, Outcome, RunInterrupted, TaskRun, run_bundle
+from .run import BundleRun, DryRun, EvalRun, Outcome, RunInterrupted, TaskRun, dry_run_bundle, run_bundle

src/agent_env/bundle/ledger.py

@@ -66,11 +66,18 @@ class Check:
     digest: Digest | None  # None when the write's inputs aren't tracked, so it is always written
     version: int | None  # the version to reuse; None when the write needs a new one
     reasons: tuple[str, ...]
+    stored: int | None  # the store's latest version of the id when checked
+    needs: Mapping[tuple[str, str], int]  # the version hashed for each earlier write it needs, by (store, id)
 
     @property
     def unchanged(self) -> bool:
         return self.version is not None
 
+    @property
+    def next_version(self) -> int:
+        """The version a new write lands on, unless another run writes the id first."""
+        return (self.stored or 0) + 1
+
 
 class Ledger:
     """One plan's rows in the ``@local`` namespace's store."""
@@ -84,15 +91,17 @@ def __init__(self, store: LocalSqliteDocumentStore, plan: Plan):
     def for_plan(cls, plan: Plan) -> Ledger:
         return cls(get_config().local_namespace_document_store(), plan)
 
-    def check(self, write: Write) -> Check:
-        """Compare ``write`` with the latest version the ledger records for it. Checks run in the plan's
-        order, after each earlier write is recorded, so the versions a dependent records are current."""
-        digest = self.digest(write)
+    def check(self, write: Write, needs: Mapping[tuple[str, str], int]) -> Check:
+        """Compare ``write`` with the latest version the ledger records for it. ``needs`` holds, by (store, id),
+        the version each earlier write it needs left or, in a dry run, would leave; an env or agent is hashed with
+        them, so one written anew rewrites it too."""
+        stored = self._stored_version(write)
+        digest = self.digest(write, needs)
         if digest is None:
-            return Check(write, None, None, ("its inputs aren't tracked yet, so it is written every run",))
+            untracked = ("its inputs aren't tracked yet, so it is written every run",)
+            return Check(write, None, None, untracked, stored, needs)
         row = self._latest(write.kind.store, write.id)
         recorded = row["version"] if row else None
-        stored = self._stored_version(write)
         reasons = []
         if row is not None and row["bundle"] != self._bundle:
             reasons.append(f"last written by the bundle {row['bundle']}")
@@ -102,10 +111,11 @@ def check(self, write: Write) -> Check:
             reasons.append(f"the store's latest, v{stored}, wasn't recorded by this bundle")
         if row is not None:
             reasons.extend(_changes(row, digest))
-        return Check(write, digest, None if reasons else recorded, tuple(reasons))
+        return Check(write, digest, None if reasons else recorded, tuple(reasons), stored, needs)
 
-    def digest(self, write: Write) -> Digest | None:
-        """What ``write`` is made from, or None when the ledger can't tell it all, so it is written every run."""
+    def digest(self, write: Write, needs: Mapping[tuple[str, str], int]) -> Digest | None:
+        """What ``write`` is made from, given the version of each write it needs, or None when the ledger can't
+        tell it all, so it is written every run."""
         if not _tracked(write):
             return None
         inputs = {"type": _type(write), "config": _sha256(_canonical(write.source.config)),
@@ -118,8 +128,7 @@ def digest(self, write: Write) -> Digest | None:
             # a store entity it names without a version getting a new one, means it must be written again. A
             # task or eval names its references without a version.
             for store, id in write.needs:
-                row = self._latest(store, id)
-                inputs["needs"][f"{store} {id}"] = str(row["version"]) if row else "unwritten"
+                inputs["needs"][f"{store} {id}"] = str(needs[store, id])
             for ref in write.source.references:
                 if ref.local is None and ref.version is None:
                     inputs["store_refs"][f"{ref.kind} {ref.id}"] = str(self._plan.store_latest[ref.kind, ref.id])
@@ -139,7 +148,7 @@ def record(self, check: Check, write: Callable[[], int]) -> int:
         self._store.insert(LEDGER_COLLECTION, pending)
         version = write()
         done = {**pending, "status": "done", "version": version, "at": _now()}
-        if check.digest is not None and self.digest(check.write) != check.digest:
+        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)
@@ -228,7 +237,7 @@ def _changes(row: Mapping[str, Any], digest: Digest) -> list[str]:
             reasons.append(f"{label}: {shown}")
     needs_before, needs_after = before["needs"], after["needs"]
     added, removed, changed = _diff(needs_before, needs_after)
-    reasons.extend(f"{name} was written anew (v{needs_before[name]} → v{needs_after[name]})" for name in changed)
+    reasons.extend(f"{name} is written anew (v{needs_before[name]} → v{needs_after[name]})" for name in changed)
     reasons.extend(f"now needs {name}" for name in added)
     reasons.extend(f"no longer needs {name}" for name in removed)
     refs_before, refs_after = before["store_refs"], after["store_refs"]

src/agent_env/bundle/materialize.py

@@ -6,14 +6,15 @@
 builds and preflights every task before writing any of them, and writes the evals last, since they name the
 tasks. Each write goes through the ledger, so one whose inputs haven't changed reuses the version the bundle
 last wrote. A reused task keeps whatever its steps took from config when it was first written, such as a
-rubrics verifier's default judge model.
+rubrics verifier's default judge model. A dry run takes the same path, refusals, checks and preflights
+included, without the lock or a single write.
 """
 
 from __future__ import annotations
 
 import copy
 from collections.abc import Callable, Iterator
-from contextlib import contextmanager
+from contextlib import contextmanager, nullcontext
 from dataclasses import dataclass
 from typing import Any
 
@@ -24,6 +25,7 @@
 from agent_env.store.routing import namespace_routing_enabled
 from agent_env.task import Task
 from agent_env.task_step.registry import get_task_step_registry
+from agent_env.task_step.task_step import TaskStep
 
 from ._fs import relative, with_article
 from .authoring import AuthoringContext
@@ -34,7 +36,8 @@
 
 @dataclass(frozen=True)
 class Materialized:
-    """One write: the version it left in the store, and whether that is the one the bundle last wrote."""
+    """One write: the version it left in the store, or in a dry run would leave, and whether that is the one the
+    bundle last wrote."""
 
     write: Write
     version: int
@@ -44,10 +47,12 @@ class Materialized:
 
 @dataclass(frozen=True)
 class Materialization:
-    """What materializing a plan left in the store, one ``Materialized`` per write."""
+    """What materializing a plan left in the store, or in a dry run would leave, one ``Materialized`` per write."""
 
     plan: Plan
     writes: tuple[Materialized, ...]  # in the plan's order
+    # A dry run's steps with a preflight that read what it would write first, so the store can't check them yet.
+    not_preflighted: tuple[tuple[Write, TaskStep], ...] = ()
 
     def version_of(self, store: str, id: str) -> int:
         for done in self.writes:
@@ -59,12 +64,17 @@ def version_of(self, store: str, id: str) -> int:
 def materialize(
     plan: Plan,
     *,
+    dry_run: bool = False,
     on_wait: Callable[[], None] | None = None,
     on_write: Callable[[Materialized], None] | 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
-    holds a lock this one needs, and ``on_write`` after each write, reused or not."""
+    holds a lock this one needs, and ``on_write`` after each write, reused or not.
+
+    ``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."""
     _refuse_unwritable(plan)
     if not namespace_routing_enabled():
         raise RuntimeError("materializing a bundle needs namespace routing, which the agent-env CLI turns on; "
@@ -76,25 +86,36 @@ def materialize(
     done: dict[tuple[str, str], Materialized] = {}
 
     def through_ledger(write: Write, write_fn: Callable[[], int]) -> None:
-        done[_key(write)] = _through_ledger(plan, ledger, write, write_fn)
+        with _noted(plan, write, "checking" if dry_run else "writing"):
+            check = ledger.check(write, {need: done[need].version for need in write.needs})
+            if check.unchanged:
+                version = check.version
+            elif dry_run:
+                version = check.next_version
+            else:
+                version = ledger.record(check, write_fn)
+        done[_key(write)] = Materialized(write, version, check.unchanged, check.reasons)
         if on_write is not None:
             on_write(done[_key(write)])
 
-    with materializing(plan.bundle.bundle, on_wait):
+    with nullcontext() if dry_run else materializing(plan.bundle.bundle, on_wait):
         for write in entities:
             through_ledger(write, lambda: _WRITERS[write.kind](plan, write))
-        built, problems = {}, []
+        unwritten = {_key(item.write) for item in done.values() if not item.reused} if dry_run else set()
+        built, problems, unchecked = {}, [], []
         for write in tasks:
             with _noted(plan, write, "preflighting"):
                 built[write.id] = _build(write)
-                problems.extend(f"{_path(plan, write)}: {problem}" for problem in _preflight(write, built[write.id]))
+                found, skipped = _preflight(write, built[write.id], unwritten)
+            problems.extend(f"{_path(plan, write)}: {problem}" for problem in found)
+            unchecked.extend((write, step) for step in skipped)
         if problems:
             raise BundleError(problems)
         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))
-    return Materialization(plan, tuple(done[_key(write)] for write in plan.writes))
+    return Materialization(plan, tuple(done[_key(write)] for write in plan.writes), tuple(unchecked))
 
 
 def _write_artifact(plan: Plan, write: Write) -> int:
@@ -160,13 +181,6 @@ def _unwritable(write: Write) -> str | None:
     return None
 
 
-def _through_ledger(plan: Plan, ledger: Ledger, write: Write, write_fn: Callable[[], int]) -> Materialized:
-    with _noted(plan, write, "writing"):
-        check = ledger.check(write)
-        version = check.version if check.unchanged else ledger.record(check, write_fn)
-    return Materialized(write, version, check.unchanged, check.reasons)
-
-
 @contextmanager
 def _noted(plan: Plan, write: Write, doing: str) -> Iterator[None]:
     try:
@@ -180,20 +194,29 @@ def _build(write: Write) -> Task:
     return Task(id=write.id, version=None, steps=[build_step(step) for step in write.source.config])
 
 
-def _preflight(write: Write, task: Task) -> list[str]:
-    """``task``'s preflight problems. A step reading one of the task's own outputs is skipped: the output
-    only exists once the task runs."""
+def _preflight(write: Write, task: Task, unwritten: set[tuple[str, str]]) -> tuple[list[str], list[TaskStep]]:
+    """``task``'s preflight problems, and the steps with a preflight left unchecked because they read a
+    (store, id) in ``unwritten``, which a dry run would write first; an artifact and an agent can share an id.
+    A step reading one of the task's own outputs is skipped too, and not listed: the output only exists once the
+    task runs."""
     outputs = {output.id for output in write.source.outputs}
-    problems = []
+    problems, unchecked = [], []
     for config, step in zip(write.source.config, task.steps):
-        if not _reads_any(config, outputs):
-            problems.extend(step.preflight())
-    return problems
+        reads = _reads(config)
+        if any(id in outputs for _, id in reads):
+            continue
+        if reads & unwritten:
+            if type(step).preflight is not TaskStep.preflight:
+                unchecked.append(step)
+            continue
+        problems.extend(step.preflight())
+    return problems, unchecked
 
 
-def _reads_any(step: dict, ids: set[str]) -> bool:
+def _reads(step: dict) -> set[tuple[str, str]]:
+    """The (store, id) of each entity ``step`` reads."""
     refs = get_task_step_registry()[step["type"]].entity_refs or ()
-    return any(site.value in ids for site in ref_sites(refs, step) if site.ref.role is RefRole.INPUT)
+    return {(site.ref.kind.value, site.value) for site in ref_sites(refs, step) if site.ref.role is RefRole.INPUT}
 
 
 def _key(write: Write) -> tuple[str, str]:

src/agent_env/bundle/plan.py

@@ -22,6 +22,7 @@
 from agent_env.env.env import Env
 from agent_env.store.base import NotFoundError
 from agent_env.store.ids import LOCAL_PREFIX
+from agent_env.store.routing import namespace_routing_enabled
 from agent_env.task import Task
 from agent_env.task_step.registry import get_task_step_registry
 
@@ -197,8 +198,7 @@ def _closure(self, roots: list[ResolvedEntry]) -> set[_Key]:
     def _check_store_refs(
         self, closure: set[_Key], *, stores: bool = True
     ) -> tuple[tuple[Reference, ...], dict[tuple[EntityKind, str], int]]:
-        """Check each store id the closure references; without ``stores``, only the ``@local`` ones, which need no
-        store read."""
+        """Check each store id the closure references; without ``stores``, only what needs no store read."""
         read: dict[tuple[EntityKind, str, int | None], Any] = {}
         checked: dict[tuple[EntityKind, str, int | None], Reference] = {}
         latest: dict[tuple[EntityKind, str], int] = {}
@@ -209,31 +209,41 @@ def _check_store_refs(
             for ref in source.references:
                 if ref.local is not None:
                     continue
-                if not stores and not ref.id.startswith(LOCAL_PREFIX):
-                    continue
-                problem = self._store_ref_problem(ref, read)
+                problem = self._bundle_ref_problem(ref)
+                if problem is None and stores:
+                    problem = self._store_ref_problem(ref, read)
                 if problem:
                     self.problems.append(f"{self._path(source.entry)}: {ref.where}: {problem}")
-                else:
+                elif stores:
                     checked.setdefault((ref.kind, ref.id, ref.version), ref)
                     if ref.version is None:
                         latest[ref.kind, ref.id] = read[ref.kind, ref.id, None].version
         return tuple(checked.values()), latest
 
-    def _store_ref_problem(self, ref: Reference, read: dict) -> str | None:
+    def _bundle_ref_problem(self, ref: Reference) -> str | None:
+        """What makes an ``@local`` id wrong without reading a store: it names one of this bundle's outputs by id,
+        one of its entities of another kind, or an id under its root that it doesn't write. Another bundle's id is
+        read like any store id."""
+        if not ref.id.startswith(LOCAL_PREFIX):
+            return None
         kind, key = ref.kind.value, _nfc(ref.id)
-        if ref.id.startswith(LOCAL_PREFIX):
-            if (kind, key) in self.outputs:
-                entry, output = self.outputs[kind, key]
-                return (f"{ref.id!r} is written by step {output.step_id!r} of {self._path(entry)}; refer to an output "
-                        "by its name, in the task that writes it")
-            other = self.identified.get(key)
-            if other is not None:
-                return (f"{ref.id!r} is this bundle's {other.kind.store} {other.name!r}, but this field takes "
-                        f"{with_article(kind)}")
-            if key.startswith(f"{self.resolved.bundle.id_root}/"):
-                return f"{ref.id!r} isn't in this bundle"
-            return f"{ref.id!r} isn't in this bundle, and references to another bundle's entities aren't supported yet"
+        if (kind, key) in self.outputs:
+            entry, output = self.outputs[kind, key]
+            return (f"{ref.id!r} is written by step {output.step_id!r} of {self._path(entry)}; refer to an output "
+                    "by its name, in the task that writes it")
+        other = self.identified.get(key)
+        if other is not None:
+            return (f"{ref.id!r} is this bundle's {other.kind.store} {other.name!r}, but this field takes "
+                    f"{with_article(kind)}")
+        if key.startswith(f"{self.resolved.bundle.id_root}/"):
+            return f"{ref.id!r} isn't in this bundle"
+        return None
+
+    def _store_ref_problem(self, ref: Reference, read: dict) -> str | None:
+        if ref.id.startswith(LOCAL_PREFIX) and not namespace_routing_enabled():
+            raise RuntimeError("reading an @local id from the store needs namespace routing, which the agent-env CLI "
+                               "turns on; call plan_bundle inside agent_env.store.routing.namespace_routing()")
+        kind = ref.kind.value
         lookup = (ref.kind, ref.id, ref.version)
         if lookup not in read:
             try: