feat(task): a registered task records the plugin packages it requires, and task run checks them first

AGENTS.md

@@ -189,6 +189,15 @@ Out-of-tree code plugs in either through `config.toml`, as a `module:Class` poin
   bundle requiring nothing. A type only configuration registers gets a hint naming the installed package its class
   comes from, never a requirement. `bundle init
   --from-task` writes the same file.
+- `bundle register` records on each task it writes the plugin packages behind its own step types and the types of the
+  bundle entities it references (`Task.requires`, saved only when there are some), in the range `bundle.toml` gives a
+  package it requires, else at least the installed version; the ledger hashes them, so a new range writes the task
+  anew. Before a stored task's first step, `agent-env task run` and `run-batch` refuse one whose requirements aren't
+  met, or whose steps read an env or artifact of a type nothing here registers, offering in a terminal to install what
+  it requires (`cli/task/_checks.py`), then set up in local stores the infra envs its gateway deploys run on
+  (`preflight.stored_infra`). A step type nothing registers names the task's unmet requirements in the load error
+  (`UnknownStepTypeError`). A task stored without `requires` loads and runs as before; a hub edit drops it until the
+  hub passes it through.
 - Environment providers have no table: an installed package registers one as an `agent_env.env_providers`
   entry point named its `type`, which records carry as `env_provider_type`.
 - `[plugins.<distribution name>]` is reserved for an installed plugin's own settings, read with

src/agent_env/bundle/importer.py

@@ -170,6 +170,7 @@ def __init__(self, dest: Path, include_scopes: Collection[FieldScope] = (),
         self.floating: list[Floating] = []
         self.dropped: list[Dropped] = []
         self.needs: list[Use] = []
+        self.declared: list[str] = []  # what the imported tasks record they require (Task.requires)
         self.undeclared: list[str] = []
         self.duplicates: list[str] = []
         self.problems: list[str] = []
@@ -216,6 +217,7 @@ def add_task(self, task: Task) -> str:
             self._default_agents(step, cls.entity_refs or (), where)
         self.files[path] = json.dumps(steps, indent=2, ensure_ascii=False, allow_nan=False) + "\n"
         self.tasks.append(ImportedTask(path, task.id, task.version, len(steps)))
+        self.declared += task.requires or ()
         return name
 
     def add_eval(self, stored: Eval, names: list[str]) -> None:
@@ -286,7 +288,7 @@ def write(self, description: str, source: dict) -> BundleImport:
             if self.problems:
                 raise BundleError(self.problems)
             shutil.rmtree(scratch)
-            if requires := detected(self.needs):
+            if requires := detected(self.needs, self.declared):
                 self.files[BUNDLE_TOML] = _toml.dumps({"requires": requires})
             files = sorted([*self.files, PROVENANCE, "README.md"])
             imported = BundleImport(

src/agent_env/bundle/ledger.py

@@ -140,6 +140,8 @@ 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.requires:
+            inputs["requires"] = list(write.requires)
         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 a run writes names
@@ -280,6 +282,8 @@ def _changes(row: Mapping[str, Any], digest: Digest) -> list[str]:
         reasons.append(f"type changed: {before['type']} → {after['type']}")
     if before.get("config") != after["config"]:
         reasons.append("config changed")
+    if before.get("requires") != after.get("requires"):
+        reasons.append(f"requires {', '.join(after['requires'])}" if after.get("requires") else "requires nothing")
     if before.get("published") != after.get("published"):
         reasons.append(f"published as {after['published'].partition('@')[0]}" if after.get("published")
                        else "not found published")

src/agent_env/bundle/materialize.py

@@ -143,7 +143,8 @@ def through_ledger(write: Write, write_fn: Callable[[], int]) -> None:
         if problems:
             raise BundleError(problems)
         for write in tasks:
-            through_ledger(write, lambda: Task.put(id=write.id, steps=built[write.id].steps).version)
+            through_ledger(write, lambda: Task.put(id=write.id, steps=built[write.id].steps,
+                                                   requires=write.requires).version)
         for write in evals:
             through_ledger(write, lambda: _write_eval(plan, write, written))
     return Materialization(plan, tuple(done[_key(write)] for write in plan.writes), tuple(unchecked))

src/agent_env/bundle/plan.py

@@ -164,6 +164,7 @@ class Write:
     id: str
     source: ResolvedEntry | BundleImage
     needs: tuple[_Key, ...]  # (kind.store, id) of the earlier writes it references
+    requires: tuple[str, ...] | None = None  # a task's plugin packages, which a registration records (Task.requires)
 
 
 @dataclass(frozen=True)

src/agent_env/bundle/preflight.py

@@ -26,8 +26,9 @@
 from __future__ import annotations
 
 import functools
-from collections.abc import Callable, Iterable
+from collections.abc import Callable, Iterable, Mapping
 from dataclasses import dataclass
+from typing import Any
 
 from agent_env.a2a_agent import A2AAgent
 from agent_env.artifact.artifacts.docker_image import PUT_AGAIN, DockerImageArtifact, pull_problem
@@ -69,6 +70,7 @@
     get_sandbox_provider,
 )
 from agent_env.store.base import NotFoundError
+from agent_env.task_step.task_step import TaskStep
 from agent_env.task_step.task_steps.deploy_agent import DeployAgentTaskStep, linked_sandbox_problem
 from agent_env.task_step.task_steps.install_agent import InstallAgentTaskStep
 from agent_env.task_step.task_steps.deploy_env import DeployEnvTaskStep
@@ -204,15 +206,11 @@ def finish(self) -> Preflight:
 
     def _env(self, where: str, step: DeployEnvTaskStep) -> None:
         provider = _provider(self.sandbox or step.sandbox_type, get_env_sandbox_provider)
-        deployment = self._planned(step.env_id) if step.env_id in self.envs else self._stored(step)
+        deployment = self._planned(step.env_id) if step.env_id in self.envs else _stored(step)
         if deployment is None:
             return
         if issubclass(deployment.provider, EnvironmentGatewayProvider):
-            kinds = {GATEWAY}
-            if _state_type(step) == LOCAL_POSTGRES_STATE_TYPE:
-                kinds.add(SERVICE_DB)
-            if deployment.websites:
-                kinds.add(WEBSITE_BROWSER)
+            kinds = _gateway_infra(step, deployment)
         elif issubclass(deployment.provider, EnvironmentServerProvider) and deployment.one_server:
             kinds = set()
         else:  # a plugin's provider, or one deploy_env's own preflight refuses for this env
@@ -237,23 +235,6 @@ def _gateway_elsewhere(self, env_id: str, env_provider: type, kinds: set[str], l
                          "provider, such as --sandbox local")
         return found + self._remote_infra(link, kinds if env_provider.in_vm(link) else kinds - {WEBSITE_BROWSER})
 
-    def _stored(self, step: DeployEnvTaskStep) -> _Deployment | None:
-        """How a store env deploys, or None when the walk leaves it to deploy_env's own preflight."""
-        try:
-            env = Env.get(step.env_id, step.env_version)
-        except NotFoundError:
-            return None  # the plan reports a store env that isn't there
-        if not isinstance(env, (MCPServerEnv, WebsiteEnv, MultiEnv)):
-            return None  # an env type that deploys itself, as deploy_env's own preflight leaves it
-        try:
-            provider_class = provider_or_class(env)
-        except (ValueError, KeyError):
-            return None  # deploy_env's own preflight reports a provider type this process can't load
-        if not isinstance(provider_class, type):
-            provider_class = type(provider_class)
-        images, websites = _stored_images(env)
-        return _Deployment(provider_class, images, websites, isinstance(env, MCPServerEnv))
-
     def _planned(self, env_id: str) -> _Deployment | None:
         """How an env the bundle writes deploys, from its planned config: what deploys it, and its images and its
         children's. None for a type the walk doesn't model, or a provider type the plan has refused."""
@@ -524,19 +505,68 @@ def _problem(self, where: str, problem: str) -> None:
         self.problems.setdefault(problem, {})[where] = None
 
 
+def stored_infra(steps: Iterable[TaskStep], user_overrides: Mapping[str, Any] | None = None) -> frozenset[str]:
+    """The infra env kinds a stored task's ``steps`` deploy on through the gateway, in a run with ``user_overrides``,
+    reading the envs they deploy from the store: the gateway, the service-db its local Postgres state runs in, and the
+    website browser for its sites."""
+    kinds: set[str] = set()
+    for step in steps:
+        if not isinstance(step, DeployEnvTaskStep) or not step.env_id:
+            continue
+        try:
+            deployment = _stored(step)
+        except ValueError:  # an id the run fills in, or an env this process can't load, which its deploy reports
+            continue
+        if deployment is not None and issubclass(deployment.provider, EnvironmentGatewayProvider):
+            kinds |= _gateway_infra(step, deployment, user_overrides or {})
+    return frozenset(kinds)
+
+
+def _gateway_infra(step: DeployEnvTaskStep, deployment: _Deployment,
+                   user_overrides: Mapping[str, Any] | None = None) -> set[str]:
+    """The infra env kinds ``step``'s deploy of ``deployment`` through the gateway runs on, in a run with
+    ``user_overrides``."""
+    kinds = {GATEWAY}
+    if _state_type(step, user_overrides or {}) == LOCAL_POSTGRES_STATE_TYPE:
+        kinds.add(SERVICE_DB)
+    if deployment.websites:
+        kinds.add(WEBSITE_BROWSER)
+    return kinds
+
+
+def _stored(step: DeployEnvTaskStep) -> _Deployment | None:
+    """How a store env deploys, or None when the walk leaves it to deploy_env's own preflight."""
+    try:
+        env = Env.get(step.env_id, step.env_version)
+    except NotFoundError:
+        return None  # the plan reports a store env that isn't there
+    if not isinstance(env, (MCPServerEnv, WebsiteEnv, MultiEnv)):
+        return None  # an env type that deploys itself, as deploy_env's own preflight leaves it
+    try:
+        provider_class = provider_or_class(env)
+    except (ValueError, KeyError):
+        return None  # deploy_env's own preflight reports a provider type this process can't load
+    if not isinstance(provider_class, type):
+        provider_class = type(provider_class)
+    images, websites = _stored_images(env)
+    return _Deployment(provider_class, images, websites, isinstance(env, MCPServerEnv))
+
+
 def _stored_images(env: Env) -> tuple[list[_Image], bool]:
     """The images a store env runs through a gateway, and whether it has websites."""
     topology = _gateway_topology(env)
     images = [*topology.mcp_server_images, *(topology.website_images or [])]
     return [_Image(f"env {env.id!r}'s image {image.id!r}", image) for image in images], bool(topology.website_configs)
 
 
-def _state_type(step: DeployEnvTaskStep) -> str | None:
-    """The type of the env state store the deploy runs on: an attached instance's own, else the one it creates."""
-    if step.env_state_instance_id is None:
-        return step.env_state_type or LOCAL_POSTGRES_STATE_TYPE
+def _state_type(step: DeployEnvTaskStep, user_overrides: Mapping[str, Any]) -> str | None:
+    """The type of the env state store the deploy runs on in a run with ``user_overrides``: an attached instance's
+    own, else the one it creates."""
+    state_type, instance_id = step.run_state(user_overrides)
+    if instance_id is None:
+        return state_type or LOCAL_POSTGRES_STATE_TYPE
     try:
-        return get_env_state_instance_store().get(step.env_state_instance_id).state_type
+        return get_env_state_instance_store().get(instance_id).state_type
     except NotFoundError:
         return None  # the deploy refuses an instance that isn't there
 

src/agent_env/bundle/register.py

@@ -43,7 +43,7 @@
 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 .requires import Declaring, declare, declaring_for
+from .requires import Declaring, Providers, declare, declaring_for, detected, entry_uses
 from .run import _label
 from .resolve import BundleImage, Reference, ResolvedBundle, ResolvedEntry, resolve_bundle
 
@@ -94,7 +94,7 @@ def register_bundle(
         ids, problems = registry_ids(plan, _prefix(bundle, prefix))
         if problems:
             raise BundleError(problems)
-        registered = translated(plan, ids)
+        registered = with_task_requires(translated(plan, ids))
         target = configured_store(get_config().get_document_store())
         shared = shared_store(target)
         identity, commit, dirty = bundle_identity(bundle.root)
@@ -382,6 +382,32 @@ def _local_url(text: str) -> str | None:
     return None
 
 
+def with_task_requires(plan: Plan) -> Plan:
+    """``plan`` with each task it writes requiring the plugin packages behind its own step types and the types of the
+    bundle entities it references (``Write.requires``): as bundle.toml writes a package when it requires it, else at
+    least the version installed here. A task using only what core ships requires nothing."""
+    providers = Providers()
+    writes = {(write.kind.store, write.id): write for write in plan.writes}
+
+    def reached(task: Write) -> Iterator[Write]:
+        seen, queue = set(), [task]
+        while queue:
+            write = queue.pop()
+            if (key := (write.kind.store, write.id)) in seen:
+                continue
+            seen.add(key)
+            yield write
+            queue.extend(writes[need] for need in write.needs if need in writes)
+
+    def requires(task: Write) -> tuple[str, ...] | None:
+        found = [use for write in reached(task) if isinstance(write.source, ResolvedEntry)
+                 for use in entry_uses(plan.bundle_of(write.source.entry).root, write.source, providers)]
+        return tuple(detected(found, plan.bundle.bundle.requires or ())) or None
+
+    return replace(plan, writes=tuple(replace(write, requires=requires(write)) if write.kind is BundleKind.TASK
+                                      else write for write in plan.writes))
+
+
 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()