From 53d92eeb58effe105d55e11ce9c2def94ac2fc31 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Wed, 28 Oct 2026 21:56:15 -0700 Subject: [PATCH 1/6] feat(task): a registered task records the plugin packages it requires, and task run checks them first bundle register records on each task the plugin packages behind its own step types and the types of the bundle entities it references, in the range bundle.toml gives a package it requires, else at least the installed version (Task.requires, saved only when there are some). The ledger hashes them, so a new range writes the task anew. Before a stored task's first step, 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 with agent-env plugin add. In local stores they then set up the infra envs the task's gateway deploys run on, as agent-env run does. A step type nothing registers names the task's unmet requirements and the command that installs them wherever a stored task is loaded. A task stored without requires loads and runs as before. Co-Authored-By: Claude Opus 5.5 --- AGENTS.md | 9 ++ src/agent_env/bundle/ledger.py | 4 + src/agent_env/bundle/materialize.py | 3 +- src/agent_env/bundle/plan.py | 1 + src/agent_env/bundle/preflight.py | 70 +++++++---- src/agent_env/bundle/register.py | 30 ++++- src/agent_env/bundle/requires.py | 83 ++++++------- src/agent_env/bundle/resolve.py | 7 +- src/agent_env/cli/task/_checks.py | 67 ++++++++++ src/agent_env/cli/task/run.py | 23 +++- src/agent_env/plugins/_requirements.py | 27 ++++- src/agent_env/task/task.py | 28 ++++- tst/unit/bundle/preflight_test.py | 18 +++ tst/unit/bundle/requires_test.py | 81 ++++++++++++- tst/unit/cli/run_role_models_test.py | 2 +- tst/unit/cli/submit_runs_test.py | 2 +- tst/unit/cli/task_run_requires_test.py | 161 +++++++++++++++++++++++++ tst/unit/cli/task_run_teardown_test.py | 2 +- tst/unit/store/derived_ids_test.py | 4 +- tst/unit/task/task_document_test.py | 39 ++++++ 20 files changed, 563 insertions(+), 98 deletions(-) create mode 100644 src/agent_env/cli/task/_checks.py create mode 100644 tst/unit/cli/task_run_requires_test.py diff --git a/AGENTS.md b/AGENTS.md index b75bb9c41..dcf400aa5 100644 --- a/AGENTS.md +++ b/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.]` is reserved for an installed plugin's own settings, read with diff --git a/src/agent_env/bundle/ledger.py b/src/agent_env/bundle/ledger.py index 573b86d8a..715a3f098 100644 --- a/src/agent_env/bundle/ledger.py +++ b/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") diff --git a/src/agent_env/bundle/materialize.py b/src/agent_env/bundle/materialize.py index f85afab91..4366fa053 100644 --- a/src/agent_env/bundle/materialize.py +++ b/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)) diff --git a/src/agent_env/bundle/plan.py b/src/agent_env/bundle/plan.py index fd6d09a1a..f25855f03 100644 --- a/src/agent_env/bundle/plan.py +++ b/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) diff --git a/src/agent_env/bundle/preflight.py b/src/agent_env/bundle/preflight.py index baa408db0..c05f96c6f 100644 --- a/src/agent_env/bundle/preflight.py +++ b/src/agent_env/bundle/preflight.py @@ -69,6 +69,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 +205,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 +234,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,6 +504,50 @@ def _problem(self, where: str, problem: str) -> None: self.problems.setdefault(problem, {})[where] = None +def stored_infra(steps: Iterable[TaskStep]) -> frozenset[str]: + """The infra env kinds a stored task's ``steps`` deploy on through the gateway, 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) + return frozenset(kinds) + + +def _gateway_infra(step: DeployEnvTaskStep, deployment: _Deployment) -> set[str]: + """The infra env kinds ``step``'s deploy of ``deployment`` through the gateway runs on.""" + kinds = {GATEWAY} + if _state_type(step) == 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) diff --git a/src/agent_env/bundle/register.py b/src/agent_env/bundle/register.py index 3f63165b9..9c66d58c7 100644 --- a/src/agent_env/bundle/register.py +++ b/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() diff --git a/src/agent_env/bundle/requires.py b/src/agent_env/bundle/requires.py index 7a207d691..ebdeb6717 100644 --- a/src/agent_env/bundle/requires.py +++ b/src/agent_env/bundle/requires.py @@ -7,7 +7,6 @@ import importlib import json import os -import shlex import uuid from collections.abc import Iterable, Mapping, Sequence from dataclasses import dataclass @@ -32,7 +31,7 @@ from .parse import BUNDLE_TOML, BundleError, BundleKind if TYPE_CHECKING: - from .resolve import ResolvedBundle + from .resolve import ResolvedBundle, ResolvedEntry # What core ships, per plugin group: anything else a bundle uses is provided by a plugin, or by configuration. _BUILTINS = {TASK_STEPS: task_step_registry._builtin_registry, ENVS: env_registry._builtin_registry, @@ -40,8 +39,8 @@ class UnmetRequirements(BundleError): - """A bundle whose ``bundle.toml`` requires packages that aren't installed, or not at a version it allows, with - every other problem resolving it found; ``requirements`` are the unmet ones, as it writes them.""" + """A bundle whose ``bundle.toml``, or a stored task whose ``requires``, names packages that aren't installed, or + not at a version it allows, with every other problem found; ``requirements`` are the unmet ones, as written.""" def __init__(self, problems: Iterable[str], requirements: Sequence[str]): super().__init__(problems) @@ -139,57 +138,45 @@ def _shipping(self, cls: type | None) -> str | None: return f"{names[0]} {version}" if version else names[0] -def unmet_problems(requires: Iterable[str]) -> tuple[list[str], list[str]]: - """The problems ``bundle.toml``'s unmet ``requires`` raise, naming the one command that installs every plugin - package among them, and those requirements. One on agent-env itself, which ``plugin add`` doesn't change, says to - upgrade it instead.""" - found = _requirements.unmet(requires) - core = [(line, why) for line, why in found if canonicalize_name(Requirement(line).name) in _requirements.CORE] - plugins = [(line, why) for line, why in found if (line, why) not in core] - problems = [f"{BUNDLE_TOML}: requires {line!r}, {why}; upgrade agent-env with the installer that installed it" - for line, why in core] - lines = [line for line, _ in plugins] - if len(plugins) == 1: - (line, why), = plugins - problems.append(f"{BUNDLE_TOML}: requires {line!r}, {why}; install it with: {install_command(lines)}") - elif plugins: - problems += [f"{BUNDLE_TOML}: requires {line!r}, {why}" for line, why in plugins] - problems.append(f"{BUNDLE_TOML}: install what it requires with: {install_command(lines)}") - return problems, lines - - -def install_command(requirements: Sequence[str]) -> str: - """The command that installs ``requirements`` into this environment, through the installer that owns it.""" - return shlex.join(["agent-env", "plugin", "add", *requirements]) - - def uses(resolved: ResolvedBundle, providers: Providers | None = None) -> list[Use]: """Each type the bundle's own tasks, envs and artifacts use that core doesn't ship.""" providers = providers or Providers() - steps, envs, artifacts = get_task_step_registry(), get_env_registry(), get_artifact_registry() + return [use for entry in resolved.entries for use in entry_uses(resolved.bundle.root, entry, providers)] + + +def entry_uses(root: Path, entry: ResolvedEntry, providers: Providers) -> list[Use]: + """Each type ``entry``, of the bundle at ``root``, uses that core doesn't ship.""" + where = relative(root, entry.entry.path) + kind = entry.entry.kind found = [] - for entry in resolved.entries: - where = relative(resolved.bundle.root, entry.entry.path) - kind = entry.entry.kind - if kind is BundleKind.TASK: - for step in entry.config: - step_type = step.get("type") - if isinstance(step_type, str) and (provider := providers.of(TASK_STEPS, step_type, - steps.get(step_type))) is not None: - found.append(Use(f"{where}: step {step.get('id')!r}", "step type", step_type, provider)) - elif kind is BundleKind.ENV and (provider := providers.of(ENVS, entry.entry.type, - envs.get(entry.entry.type))) is not None: - found.append(Use(where, "env type", entry.entry.type, provider)) - elif kind in (BundleKind.ARTIFACT, BundleKind.SKILL): - name = canonical_type(entry.entry.type) - if (provider := providers.of(ARTIFACTS, name, artifacts.get(name))) is not None: - found.append(Use(where, "artifact type", name, provider)) + if kind is BundleKind.TASK: + steps = get_task_step_registry() + for step in entry.config: + step_type = step.get("type") + if isinstance(step_type, str) and (provider := providers.of(TASK_STEPS, step_type, + steps.get(step_type))) is not None: + found.append(Use(f"{where}: step {step.get('id')!r}", "step type", step_type, provider)) + elif kind is BundleKind.ENV and (provider := providers.of(ENVS, entry.entry.type, + get_env_registry().get(entry.entry.type))) is not None: + found.append(Use(where, "env type", entry.entry.type, provider)) + elif kind in (BundleKind.ARTIFACT, BundleKind.SKILL): + name = canonical_type(entry.entry.type) + if (provider := providers.of(ARTIFACTS, name, get_artifact_registry().get(name))) is not None: + found.append(Use(where, "artifact type", name, provider)) return found -def detected(found: Iterable[Use]) -> list[str]: - """The requirement for each distinct plugin distribution ``found`` names, at least the version installed here.""" - return sorted({requirement for use in found if (requirement := use.provider.requirement()) is not None}) +def detected(found: Iterable[Use], declared: Iterable[str] = ()) -> list[str]: + """The requirement for each distinct plugin distribution ``found`` names: as ``declared`` writes it when it names + that distribution, else at least the version installed here.""" + by_name: dict[str, list[str]] = {} + for line in declared: + by_name.setdefault(canonicalize_name(Requirement(line).name), []).append(line) + requires: set[str] = set() + for use in found: + if (requirement := use.provider.requirement()) is not None: + requires.update(by_name.get(canonicalize_name(use.provider.distribution), [requirement])) + return sorted(requires) def declare(resolved: ResolvedBundle, declaring: Declaring) -> list[str]: diff --git a/src/agent_env/bundle/resolve.py b/src/agent_env/bundle/resolve.py index fd0cce0e5..5b7e59e1b 100644 --- a/src/agent_env/bundle/resolve.py +++ b/src/agent_env/bundle/resolve.py @@ -25,6 +25,7 @@ from agent_env.env.registry import get_env_registry from agent_env.eval.eval import Eval from agent_env.plugins import _registration +from agent_env.plugins._requirements import unmet_problems from agent_env.store.ids import LOCAL_PREFIX, derive_id, validate_local_id from agent_env.store.image_store.oci_registry_credentials import names_registry from agent_env.task_step.registry import get_task_step_registry @@ -32,8 +33,8 @@ from ._fs import fold, relative, show from .authoring import AuthoringContext, default_dockerfile -from .parse import DOCKERFILE, NAMED_BY, Bundle, BundleEntry, BundleError, BundleKind -from .requires import UnmetRequirements, unmet_problems +from .parse import BUNDLE_TOML, DOCKERFILE, NAMED_BY, Bundle, BundleEntry, BundleError, BundleKind +from .requires import UnmetRequirements # Step keys that hold step ids, never entity ids. _STEP_ID_KEYS = ("id", "type", "depends_on") @@ -171,7 +172,7 @@ def __init__(self, bundle: Bundle): self.ignored[kind, fold(name)] = path def resolve(self) -> ResolvedBundle: - problems, unmet = unmet_problems(self.bundle.requires or ()) + problems, unmet = unmet_problems(self.bundle.requires or (), BUNDLE_TOML) self.problems.extend(problems) resolved = [] for entry in self.bundle.entries: diff --git a/src/agent_env/cli/task/_checks.py b/src/agent_env/cli/task/_checks.py new file mode 100644 index 000000000..093a6db17 --- /dev/null +++ b/src/agent_env/cli/task/_checks.py @@ -0,0 +1,67 @@ +"""What `agent-env task run` checks of a stored task before its first step: that the plugin packages it requires are +installed and every env and artifact type its steps read resolves. With local stores it then sets up the infra envs +the task deploys on, as `agent-env run` does.""" + +from __future__ import annotations + +from collections.abc import Callable + +import click + +from agent_env.artifact.artifact import Artifact +from agent_env.bundle.preflight import stored_infra +from agent_env.bundle.requires import UnmetRequirements +from agent_env.entity.refs import EntityKind, RefRole, ref_sites +from agent_env.env.bootstrap import GATEWAY, SERVICE_DB, ensure_default_envs, stores_are_local +from agent_env.env.env import Env +from agent_env.plugins._requirements import unmet_problems +from agent_env.store.base import NotFoundError +from agent_env.store.document_store import UnknownDocumentTypeError +from agent_env.task import Task +from agent_env.task.task import UnknownStepTypeError + +# The entities a step reads whose type a plugin may register. +_GETTERS = {EntityKind.ENV: Env.get, EntityKind.ARTIFACT: Artifact.get} + + +def runnable_task(task_id: str, version: int | None, *, say: Callable[[str], None], gateway_env_id: str | None = None, + service_db_env_id: str | None = None) -> Task: + """The stored task ``task_id`` at ``version``, or its latest, once what it requires is installed and every env and + artifact type its steps read loads here. With local stores, the infra envs it deploys on are set up, but for one + the run is given its own of (``gateway_env_id``, ``service_db_env_id``). Raises UnmetRequirements naming the + command that installs what it requires, or a ClickException for a type nothing here registers.""" + try: + task = Task.get(task_id, version=version) + except UnknownStepTypeError as e: + if e.requirements: + raise UnmetRequirements([str(e)], e.requirements) from e + raise + owner = f"task {task.id} v{task.version}" + problems, unmet = unmet_problems(task.requires or (), owner) + problems += [f"{owner}: {problem}" for problem in _unknown_types(task)] + if unmet: + raise UnmetRequirements(problems, unmet) + if problems: + raise click.ClickException("\n".join(problems)) + overridden = {kind for kind, given in ((GATEWAY, gateway_env_id), (SERVICE_DB, service_db_env_id)) if given} + if stores_are_local() and (kinds := stored_infra(task.steps) - overridden): + ensure_default_envs(kinds, say=say) + return task + + +def _unknown_types(task: Task) -> list[str]: + """Each env and artifact a step of ``task`` reads, as the store holds it, whose type nothing here registers. One + the store doesn't hold, or whose id the run fills in, is left to its step.""" + found = [] + for step in task.steps: + for site in ref_sites(type(step).entity_refs or (), step.to_dict()): + get = _GETTERS.get(site.ref.kind) + if site.ref.role is RefRole.OUTPUT or get is None or not isinstance(site.value, str): + continue + try: + get(site.value, site.version if isinstance(site.version, int) else None) + except UnknownDocumentTypeError as e: + found.append(f"step {step.id!r}: {site.path} {site.value!r}: {e}") + except (NotFoundError, ValueError): + continue + return found diff --git a/src/agent_env/cli/task/run.py b/src/agent_env/cli/task/run.py index 70310ad17..67f62040a 100644 --- a/src/agent_env/cli/task/run.py +++ b/src/agent_env/cli/task/run.py @@ -11,7 +11,10 @@ import click +from agent_env.bundle.requires import UnmetRequirements +from agent_env.cli._install_prompt import offer_to_install from agent_env.cli.banner import print_banner +from agent_env.cli.task._checks import runnable_task from agent_env.cli import submitted_runs from agent_env.cli.identity import get_agent_env_client_id from agent_env.cli.teardown_output import echo_teardown @@ -377,8 +380,6 @@ def run(task_id: str, task_version: int | None, output_dir: str | None, k: int, if k < 1: raise click.BadParameter("must be at least 1", param_hint="'--k'") - from agent_env.task import Task - if output_dir is None: output_dir = f"/tmp/taskid={fs_safe(task_id)}-run{uuid4().hex[:6]}" if context_json: @@ -419,7 +420,7 @@ def run(task_id: str, task_version: int | None, output_dir: str | None, k: int, initial_context.metadata["task_id"] = task_id click.echo(f"Fetching task: id={task_id} version={task_version or 'latest'}...") - task = Task.get(task_id, version=task_version) + task = _runnable(task_id, task_version, gateway_env_id, service_db_env_id) click.echo(f"Found task: id={task.id} version={task.version} steps={len(task.steps)}") print_banner() @@ -464,6 +465,18 @@ def run(task_id: str, task_version: int | None, output_dir: str | None, k: int, raise SystemExit(1) +def _runnable(task_id: str, task_version: int | None, gateway_env_id: str | None = None, + service_db_env_id: str | None = None): + """The stored task, checked before its first step (``runnable_task``); a refusal for what it requires offers, in a + terminal, to install it and run the command again.""" + try: + return runnable_task(task_id, task_version, say=click.echo, gateway_env_id=gateway_env_id, + service_db_env_id=service_db_env_id) + except UnmetRequirements as e: + offer_to_install(e) + raise + + def _seed_universe_id(task, seed: dict) -> str | None: """The universe a seed's runs collect into, so runs of one seed share it: derived from the task version and the seed's id, else its name. None leaves collect_artifacts to name it after the run.""" @@ -506,8 +519,6 @@ def run_batch(task_id: str, task_version: int | None, seeds: str, concurrency: i context.metadata["seed"]. Prompt templates with variables are rendered using the seed values before being sent to the agent. """ - from agent_env.task import Task - if output_dir is None: output_dir = f"/tmp/taskid={fs_safe(task_id)}-batch{uuid4().hex[:6]}" @@ -521,7 +532,7 @@ def run_batch(task_id: str, task_version: int | None, seeds: str, concurrency: i raise click.ClickException("No seeds found in CSV") click.echo(f"Fetching task: id={task_id} version={task_version or 'latest'}...") - task = Task.get(task_id, version=task_version) + task = _runnable(task_id, task_version) click.echo(f"Found task: id={task.id} version={task.version} steps={len(task.steps)}") suffixes = {step.universe_id_suffix or "" for step in task.steps if isinstance(step, CollectArtifactsTaskStep)} if is_local_id(task.id) and suffixes: diff --git a/src/agent_env/plugins/_requirements.py b/src/agent_env/plugins/_requirements.py index f7f5b1f91..5c2bb88ab 100644 --- a/src/agent_env/plugins/_requirements.py +++ b/src/agent_env/plugins/_requirements.py @@ -8,7 +8,8 @@ from __future__ import annotations import functools -from collections.abc import Iterable +import shlex +from collections.abc import Iterable, Sequence from importlib.metadata import PackageNotFoundError, distribution, version from typing import Any, Optional @@ -71,6 +72,30 @@ def unmet(requires: Iterable[str]) -> list[tuple[str, str]]: return found +def unmet_problems(requires: Iterable[str], owner: str) -> tuple[list[str], list[str]]: + """The problems ``owner``'s unmet ``requires`` raise, a bundle's ``bundle.toml`` or a stored task, naming the one + command that installs every plugin package among them, and those requirements. One on agent-env itself, which + ``plugin add`` doesn't change, says to upgrade it instead.""" + found = unmet(requires) + core = [(line, why) for line, why in found if canonicalize_name(Requirement(line).name) in CORE] + plugins = [(line, why) for line, why in found if (line, why) not in core] + problems = [f"{owner}: requires {line!r}, {why}; upgrade agent-env with the installer that installed it" + for line, why in core] + lines = [line for line, _ in plugins] + if len(plugins) == 1: + (line, why), = plugins + problems.append(f"{owner}: requires {line!r}, {why}; install it with: {install_command(lines)}") + elif plugins: + problems += [f"{owner}: requires {line!r}, {why}" for line, why in plugins] + problems.append(f"{owner}: install what it requires with: {install_command(lines)}") + return problems, lines + + +def install_command(requirements: Sequence[str]) -> str: + """The command that installs ``requirements`` into this environment, through the installer that owns it.""" + return shlex.join(["agent-env", "plugin", "add", *requirements]) + + def _allows(requirement: Requirement, version: str) -> bool: """Whether ``requirement`` admits the installed ``version``, prereleases included as pip check counts them; one pip can't compare it admits only when it names no versions.""" diff --git a/src/agent_env/task/task.py b/src/agent_env/task/task.py index 3f2d07a5d..b6de1064e 100644 --- a/src/agent_env/task/task.py +++ b/src/agent_env/task/task.py @@ -13,10 +13,12 @@ import uuid from dataclasses import dataclass from datetime import datetime, timezone +from collections.abc import Sequence from typing import TYPE_CHECKING, Callable, ClassVar, Optional from agent_env.config import run_env_ids from agent_env.plugins import _registration +from agent_env.plugins._requirements import unmet_problems from agent_env.store.document_store import UnknownDocumentTypeError from agent_env.store.routing import run_scope from agent_env.task_step.context import TaskStepContext, regraft_redacted_keys @@ -311,11 +313,22 @@ class _SchedulerState: epoch: int = 0 +class UnknownStepTypeError(UnknownDocumentTypeError): + """A stored task naming a step type nothing here registers; ``requirements`` are the ones it records + (``Task.requires``) that aren't met here, which the message names with the command that installs them.""" + + def __init__(self, message: str, requirements: Sequence[str] = ()): + super().__init__(message) + self.requirements = tuple(requirements) + + class Task: """A task is an ordered sequence of task steps. Steps are stored inline and executed by run() according to the DAG encoded in each step's depends_on. Independent steps run concurrently. + ``requires`` names the plugin packages its steps and what they reference use, as pip requirements, which a + registration records (``agent-env bundle register``); None when nothing recorded it. """ type: ClassVar[str] = "task" @@ -325,10 +338,12 @@ def __init__( id: str, version: Optional[int], steps: list[TaskStep] | None = None, + requires: Sequence[str] | None = None, ): self.id = id self.version = version self.steps: list[TaskStep] = steps or [] + self.requires: list[str] | None = list(requires) if requires is not None else None self._validate_dag(self.steps) @classmethod @@ -377,12 +392,15 @@ def preflight(self) -> list[str]: return problems def to_dict(self) -> dict: - return { + data = { "id": self.id, "type": self.type, "version": self.version, "steps": [step.to_dict() for step in self.steps], } + if self.requires is not None: + data["requires"] = list(self.requires) + return data @classmethod def from_dict(cls, data: dict) -> "Task": @@ -394,8 +412,11 @@ def from_dict(cls, data: dict) -> "Task": step_type = step_dict["type"] step_cls = registry.get(step_type) if step_cls is None: - raise UnknownDocumentTypeError( - f"Unknown task step type: {step_type}{_registration.failure_note(_registration.TASK_STEPS, step_type)}" + problems, unmet = unmet_problems(data.get("requires") or (), f"task {data['id']} v{data.get('version')}") + raise UnknownStepTypeError( + "; ".join([f"Unknown task step type: {step_type}" + f"{_registration.failure_note(_registration.TASK_STEPS, step_type)}", *problems]), + unmet, ) from agent_env.task_step.task_step import attach_retry_config @@ -404,6 +425,7 @@ def from_dict(cls, data: dict) -> "Task": id=data["id"], version=data.get("version"), steps=steps, + requires=data.get("requires"), ) async def run( diff --git a/tst/unit/bundle/preflight_test.py b/tst/unit/bundle/preflight_test.py index 394c2ab75..1248b24f2 100644 --- a/tst/unit/bundle/preflight_test.py +++ b/tst/unit/bundle/preflight_test.py @@ -19,6 +19,7 @@ from agent_env.artifact.store import get_artifact_store from agent_env.bundle import BundleError, BundleKind, dry_run_bundle, run_bundle from agent_env.bundle.run import _write_bundle +from agent_env.bundle.resolve import build_step from agent_env.cli import cli from agent_env.config import get_config from agent_env.config.runtime import Config @@ -753,6 +754,23 @@ def test_an_attached_state_instance_needs_the_service_db_by_its_own_type(bundle_ assert [build.kind for build in dry_run_bundle(bundle_dir).infra] == kinds +@pytest.mark.parametrize("step, kinds", [ + ({"env_id": "crm"}, {"gateway", "service-db"}), + ({"env_id": "shop", "env_state_type": "external_db"}, {"gateway", "website-browser"}), + ({"env_id": "{{seed.env}}"}, set()), + ({"env_id": "absent"}, set()), +], ids=["local-postgres", "website", "filled-in-by-the-run", "not-in-the-store"]) +def test_a_stored_tasks_gateway_deploys_name_their_infra_as_a_run_of_its_folder_does(bundle_dir, step, kinds): + _env("crm") + with namespace_routing(): + WebsiteEnv.put(id="shop", backend_docker_image_artifact=_image("shop-back"), + frontend_docker_image_artifact=_image("shop-front"), environment_name="shop") + steps = [build_step({"id": "env", "type": "deploy_env", **step}), + build_step({"id": "box", "type": "deploy_sandbox", "sandbox_name": "box", "sandbox_mode": "vm"})] + + assert preflight_module.stored_infra(steps) == kinds + + def test_infra_built_from_other_inputs_is_rebuilt_and_infra_built_from_this_releases_isnt(bundle_dir, monkeypatch): _env("crm") _task(bundle_dir, [{"id": "env", "type": "deploy_env", "env_id": "crm"}]) diff --git a/tst/unit/bundle/requires_test.py b/tst/unit/bundle/requires_test.py index ec296e9c1..89ff53859 100644 --- a/tst/unit/bundle/requires_test.py +++ b/tst/unit/bundle/requires_test.py @@ -9,26 +9,29 @@ import tomllib from pathlib import PurePosixPath from types import SimpleNamespace +from typing import ClassVar, Literal import pytest from click.testing import CliRunner from packaging.requirements import Requirement +from agent_env.artifact import Artifact from agent_env.bundle import BundleError, dry_run_bundle, parse_bundle, prepare_bundle, run_bundle from agent_env.bundle import register as register_module from agent_env.bundle import requires as requires_module from agent_env.bundle.register import register_bundle from agent_env.bundle.requires import ( - Declaring, Provider, Providers, UnmetRequirements, _create, declare, declaring_for, install_command, unmet_problems, + Declaring, Provider, Providers, UnmetRequirements, _create, declare, declaring_for, ) from agent_env.bundle.resolve import resolve_bundle from agent_env.bundle.run import _write_bundle from agent_env.cli import _install_prompt, cli from agent_env.cli._installers import InstallerError from agent_env.config.runtime import Config -from agent_env.plugins import TASK_STEPS, Contribution, Distribution -from agent_env.plugins._requirements import unmet +from agent_env.plugins import ARTIFACTS, TASK_STEPS, Contribution, Distribution +from agent_env.plugins._requirements import install_command, unmet, unmet_problems from agent_env.store.routing import namespace_routing +from agent_env.task import Task from agent_env.task_step.task_step import TaskStep from tst.unit.bundle._support import fake_plugins, layout @@ -98,6 +101,10 @@ def _plugins(monkeypatch, *, versions=None, **more): **more) +def _bundle_problems(requires): + return unmet_problems(requires, "bundle.toml") + + def _bundle(root, *types, toml=None): files = {"tasks/t.json": json.dumps([{"id": f"s{i}", "type": step_type} for i, step_type in enumerate(types)])} if toml is not None: @@ -182,12 +189,12 @@ def test_an_extra_whose_needs_are_installed_or_that_needs_nothing_is_met(self, m assert unmet([line]) == [] def test_one_unmet_requirement_names_the_command_that_installs_it(self): - assert unmet_problems(["agentenv-missing>=1", "agentenv-acme"]) == ([MISSING], ["agentenv-missing>=1"]) + assert _bundle_problems(["agentenv-missing>=1", "agentenv-acme"]) == ([MISSING], ["agentenv-missing>=1"]) def test_agent_env_itself_is_upgraded_not_added(self, monkeypatch): _plugins(monkeypatch, versions={"agentenv-framework": "0.9.1375"}) - assert unmet_problems(["agentenv-framework>=0.9.1400", "agentenv-missing>=1"]) == ([ + assert _bundle_problems(["agentenv-framework>=0.9.1400", "agentenv-missing>=1"]) == ([ "bundle.toml: requires 'agentenv-framework>=0.9.1400', and agentenv-framework 0.9.1375 is installed; " "upgrade agent-env with the installer that installed it", MISSING, @@ -204,7 +211,7 @@ def test_a_bundle_needing_a_newer_agent_env_is_refused_without_an_install_to_off assert caught.value.problems[0].endswith("upgrade agent-env with the installer that installed it") def test_several_name_one_command_that_installs_them_all(self): - assert unmet_problems(["agentenv-missing>=1", "agentenv-acme>=2"]) == ([ + assert _bundle_problems(["agentenv-missing>=1", "agentenv-acme>=2"]) == ([ "bundle.toml: requires 'agentenv-missing>=1', which isn't installed", "bundle.toml: requires 'agentenv-acme>=2', and agentenv-acme 1.2.0 is installed", "bundle.toml: install what it requires with: agent-env plugin add 'agentenv-missing>=1' 'agentenv-acme>=2'", @@ -593,5 +600,67 @@ def test_a_registration_refused_for_it_asks_too(self, terminal): assert len(calls["exec"]) == 1 +class _AcmeSetting(Artifact): + """An artifact type the plugin agentenv-acme-settings ships, made only of its artifact.toml's values.""" + + type: Literal["acme_setting_requires_test"] = "acme_setting_requires_test" + toml_keys: ClassVar[dict[str, type]] = {"value": str} + value: str = "" + + @classmethod + def from_toml(cls, data, ctx): + return cls.put(id=ctx.id, value=ctx.accept(data, **cls.toml_keys).get("value", "")) + + +SETTINGS = "agentenv-acme-settings" +SANDBOX = {"id": "box", "type": "deploy_sandbox", "sandbox_name": "box", "sandbox_mode": "vm"} +REGISTERED = { + "tasks/plugin.json": json.dumps([{"id": "s0", "type": _Acme.type}]), + "tasks/loads.json": json.dumps([SANDBOX, {"id": "load", "type": "load_artifact", "sandbox_name": "box", + "artifact_id": "cfg", "depends_on": ["box"]}]), + "tasks/core.json": json.dumps([SANDBOX]), + "artifacts/cfg/artifact.toml": 'type = "acme_setting_requires_test"\nvalue = "x"\n', +} + + +class TestRegisteredTasks: + @pytest.fixture(autouse=True) + def two_plugins(self, stores, monkeypatch): + """agentenv-acme's step type, and agentenv-acme-settings's artifact type, which a task references.""" + monkeypatch.setattr(register_module, "_git", lambda root, *args: None) + artifacts = {**Config().artifact_registry(), _AcmeSetting.model_fields["type"].default: _AcmeSetting} + monkeypatch.setattr(Config, "artifact_registry", lambda self: artifacts) + fake_plugins(monkeypatch, [ + Distribution(ACME, "1.2.0", (Contribution(TASK_STEPS, _Acme.type, _value(_Acme), "active"),)), + Distribution(SETTINGS, "2.0", (Contribution(ARTIFACTS, _AcmeSetting.model_fields["type"].default, + _value(_AcmeSetting), "active"),)), + ], versions={ACME: "1.2.0", SETTINGS: "2.0"}) + + def test_each_task_requires_what_it_and_the_bundle_entities_it_references_use(self, stores): + root = layout(stores / "b", REGISTERED) + + register_bundle(root) + + assert {name: Task.get(f"b-{name}").requires for name in ("plugin", "loads", "core")} == { + "plugin": ["agentenv-acme>=1.2.0"], "loads": ["agentenv-acme-settings>=2.0"], "core": None} + assert "requires" not in Task.get("b-core").to_dict() + + def test_a_package_bundle_toml_requires_keeps_its_range_and_a_new_range_writes_the_task_anew(self, stores): + root = layout(stores / "b", {**REGISTERED, + "bundle.toml": 'requires = ["agentenv-acme>=1.0,<2", "agentenv-acme-settings"]\n'}) + + register_bundle(root) + again = register_bundle(root) + (root / "bundle.toml").write_text('requires = ["agentenv-acme>=1.1,<2", "agentenv-acme-settings"]\n') + changed = [done for done in register_bundle(root).materialization.writes if not done.reused] + + assert (Task.get("b-plugin", 1).requires, Task.get("b-loads").requires) == ( + ["agentenv-acme>=1.0,<2"], ["agentenv-acme-settings"]) + assert all(done.reused for done in again.materialization.writes) + assert [(done.write.id, done.version, done.reasons) for done in changed] == [ + ("b-plugin", 2, ("requires agentenv-acme>=1.1,<2",))] + assert Task.get("b-plugin").requires == ["agentenv-acme>=1.1,<2"] + + def test_unmet_requirements_are_a_bundle_error(): assert issubclass(UnmetRequirements, BundleError) diff --git a/tst/unit/cli/run_role_models_test.py b/tst/unit/cli/run_role_models_test.py index d266d2fe5..b4837aefe 100644 --- a/tst/unit/cli/run_role_models_test.py +++ b/tst/unit/cli/run_role_models_test.py @@ -21,7 +21,7 @@ class _Recorded: """Records each context it runs with.""" - id, version, steps = "t", 1, [] + id, version, steps, requires = "t", 1, [], None def __init__(self): self.contexts = [] diff --git a/tst/unit/cli/submit_runs_test.py b/tst/unit/cli/submit_runs_test.py index 115f37025..da5d83159 100644 --- a/tst/unit/cli/submit_runs_test.py +++ b/tst/unit/cli/submit_runs_test.py @@ -28,7 +28,7 @@ class _Task: - id, version, steps = "t", 1, [] + id, version, steps, requires = "t", 1, [], None def __init__(self, nap=0.0, failed_steps=None): self.nap, self.failed_steps, self.contexts, self.calls = nap, failed_steps, [], [] diff --git a/tst/unit/cli/task_run_requires_test.py b/tst/unit/cli/task_run_requires_test.py new file mode 100644 index 000000000..fac65a999 --- /dev/null +++ b/tst/unit/cli/task_run_requires_test.py @@ -0,0 +1,161 @@ +"""What `agent-env task run` checks of a stored task before its first step: the plugin packages it requires, the types +of the envs and artifacts its steps read, and, in local stores, the infra envs its deploys run on.""" + +import logging + +import pytest +from click.testing import CliRunner + +from agent_env.cli import _install_prompt, cli, submitted_runs +from agent_env.cli.task import _checks +from agent_env.config import get_config +from agent_env.env.bootstrap import GATEWAY, SERVICE_DB +from agent_env.env.store import ENVS_COLLECTION +from agent_env.plugins import _requirements +from agent_env.task import Task +from agent_env.task.store import TASKS_COLLECTION + +MISSING = "agentenv-e2e-missing>=1" +REFUSED = f"task t v1: requires '{MISSING}', which isn't installed; install it with: agent-env plugin add '{MISSING}'" + + +@pytest.fixture +def stores(local_stores, monkeypatch, tmp_path): + """Local stores holding no infra envs, only agentenv-e2e-acme 1.2 installed beside agent-env, and a run that + fails the test.""" + monkeypatch.setattr(_requirements, "installed", {"agentenv-framework": "0.9.1380", "agentenv-e2e-acme": "1.2"}.get) + + def ran(*args, **kwargs): + raise AssertionError("the task ran") + + monkeypatch.setattr(submitted_runs, "follow", ran) + logging.disable(logging.CRITICAL) + yield tmp_path + logging.disable(logging.NOTSET) + + +def _stored(steps, **fields): + """Store task ``t`` v1 as its document, steps and all, so one naming a type nothing registers can be stored.""" + get_config().get_document_store().insert(TASKS_COLLECTION, {"id": "t", "type": "task", "version": 1, + "steps": steps, **fields}) + + +def _run(*args, **kwargs): + return CliRunner().invoke(cli, ["task", "run", "--id", "t", *args], **kwargs) + + +def test_a_task_requiring_what_isnt_installed_is_refused_before_it_runs(stores): + Task.put(id="t", steps=[], requires=[MISSING]) + + result = _run() + + assert result.exit_code == 1 + assert f"Error: {REFUSED}" in result.output + + +def test_a_step_type_nothing_registers_names_the_package_the_task_requires(stores): + _stored([{"id": "s", "type": "acme_task_run_test"}], requires=[MISSING]) + (stores / "seeds.csv").write_text("n\n1\n") + + ran = _run() + batch = CliRunner().invoke(cli, ["task", "run-batch", "--id", "t", "--seeds", str(stores / "seeds.csv")]) + + for result in (ran, batch): + assert result.exit_code == 1 + assert f"Error: Unknown task step type: acme_task_run_test; {REFUSED}" in result.output + + +def test_without_a_record_a_step_type_nothing_registers_says_only_that(stores): + _stored([{"id": "s", "type": "acme_task_run_test"}]) + + result = _run() + + assert result.exit_code == 1 + assert "Unknown task step type: acme_task_run_test" in result.output + assert "plugin add" not in result.output + + +def test_an_env_a_step_reads_whose_type_nothing_registers_is_refused_before_it_runs(stores): + get_config().get_document_store().insert(ENVS_COLLECTION, {"id": "crm", "type": "acme_env_task_run_test", + "version": 1}) + _stored([{"id": "env", "type": "deploy_env", "env_id": "crm"}], requires=[MISSING]) + + result = _run() + + assert result.exit_code == 1 + assert f"Error: {REFUSED}" in result.output + assert "task t v1: step 'env': env_id 'crm': Unknown env type: acme_env_task_run_test" in result.output + + +def test_one_the_store_doesnt_hold_or_the_run_fills_in_is_left_to_its_step(stores, monkeypatch): + _stored([{"id": "env", "type": "deploy_env", "env_id": "absent"}, + {"id": "seeded", "type": "deploy_env", "env_id": "{{seed.env}}"}]) + monkeypatch.setattr(_checks, "stores_are_local", lambda: False) + + assert _checks.runnable_task("t", None, say=print).id == "t" + + +def test_a_task_whose_requirements_are_met_is_run(stores, monkeypatch): + Task.put(id="t", steps=[], requires=["agentenv-e2e-acme>=1"]) + + assert _checks.runnable_task("t", None, say=print).requires == ["agentenv-e2e-acme>=1"] + + +@pytest.fixture +def terminal(stores, monkeypatch): + """A person at a terminal, an install that records what it's asked for, and a re-exec that stops the test.""" + calls = {"added": [], "exec": []} + + def execv(path, argv): + calls["exec"].append(argv) + raise SystemExit(0) + + monkeypatch.setattr(_install_prompt, "_interactive", lambda: True) + monkeypatch.setattr(_install_prompt._plugin_changes, "add", lambda specs, **kwargs: calls["added"].append(specs)) + monkeypatch.setattr(_install_prompt.os, "execv", execv) + monkeypatch.setenv(_install_prompt._ASKED, "") + return calls + + +@pytest.mark.parametrize("store", [ + lambda: Task.put(id="t", steps=[], requires=[MISSING]), + lambda: _stored([{"id": "s", "type": "acme_task_run_test"}], requires=[MISSING]), +], ids=["requirement", "step-type"]) +def test_at_a_terminal_it_offers_to_install_what_the_task_requires_and_runs_it_again(terminal, store): + store() + + result = _run(input="y\n") + + assert result.exit_code == 0, result.output + assert f"Install {MISSING} with agent-env plugin add? [y/N]: y" in result.output + assert terminal["added"] == [[MISSING]] and len(terminal["exec"]) == 1 + + +def test_a_task_needing_a_newer_agent_env_says_to_upgrade_it_and_never_asks(terminal): + Task.put(id="t", steps=[], requires=["agentenv-framework>=99"]) + + result = _run(input="y\n") + + assert result.exit_code == 1 + assert ("task t v1: requires 'agentenv-framework>=99', and agentenv-framework 0.9.1380 is installed; upgrade " + "agent-env with the installer that installed it") in result.output + assert "Install" not in result.output and terminal["added"] == [] + + +@pytest.mark.parametrize("local, overrides, kinds", [ + (True, {}, {GATEWAY, SERVICE_DB}), + (True, {"gateway_env_id": "mine"}, {SERVICE_DB}), + (True, {"gateway_env_id": "mine", "service_db_env_id": "mine-db"}, None), + (False, {}, None), +], ids=["local", "own-gateway", "own-both", "shared"]) +def test_the_infra_envs_its_deploys_run_on_are_set_up_in_local_stores_but_for_those_it_is_given( + stores, monkeypatch, local, overrides, kinds): + Task.put(id="t", steps=[]) + set_up = [] + monkeypatch.setattr(_checks, "stores_are_local", lambda: local) + monkeypatch.setattr(_checks, "stored_infra", lambda steps: frozenset({GATEWAY, SERVICE_DB})) + monkeypatch.setattr(_checks, "ensure_default_envs", lambda kinds, say: set_up.append(set(kinds))) + + _checks.runnable_task("t", None, say=print, **overrides) + + assert set_up == ([kinds] if kinds else []) diff --git a/tst/unit/cli/task_run_teardown_test.py b/tst/unit/cli/task_run_teardown_test.py index b0d52aa63..a1bb4b505 100644 --- a/tst/unit/cli/task_run_teardown_test.py +++ b/tst/unit/cli/task_run_teardown_test.py @@ -29,7 +29,7 @@ class _Deploys: """Records a sandbox on a fake backend, then naps, then ends as told.""" - id, version, steps = "t", 1, [] + id, version, steps, requires = "t", 1, [], None def __init__(self, nap=0.0, ending=None): self.nap, self.ending, self.contexts = nap, ending, [] diff --git a/tst/unit/store/derived_ids_test.py b/tst/unit/store/derived_ids_test.py index 20c6c97e6..3b0108ddb 100644 --- a/tst/unit/store/derived_ids_test.py +++ b/tst/unit/store/derived_ids_test.py @@ -335,7 +335,7 @@ def test_a_seeded_runs_universe_is_named_after_its_task_and_seed(tmp_path, monke seen: list[str] = [] class _Task: - id, version, steps = task_id, 2, [] + id, version, steps, requires = task_id, 2, [], None async def run(self, context, **kwargs): seen.append(context.metadata["universe_id"]) @@ -367,7 +367,7 @@ def test_a_seed_that_cant_name_an_local_universe_stops_a_collecting_run_batch_be ran: list[TaskStepContext] = [] class _Task: - id, version, steps = LOCAL_TASK, 2, task_steps + id, version, steps, requires = LOCAL_TASK, 2, task_steps, None async def run(self, context, **kwargs): ran.append(context) diff --git a/tst/unit/task/task_document_test.py b/tst/unit/task/task_document_test.py index 7272dcd4c..80e779caa 100644 --- a/tst/unit/task/task_document_test.py +++ b/tst/unit/task/task_document_test.py @@ -1,4 +1,8 @@ +import pytest + +from agent_env.plugins import _requirements from agent_env.task import Task +from agent_env.task.task import UnknownStepTypeError def test_a_stored_task_document_with_an_unread_project_id_still_loads(): @@ -7,3 +11,38 @@ def test_a_stored_task_document_with_an_unread_project_id_still_loads(): assert (task.id, task.version, task.steps) == ("t", 3, []) assert "project_id" not in task.to_dict() + + +def test_a_task_records_what_it_requires_only_when_given_it(): + assert "requires" not in Task("t", 1).to_dict() + task = Task.from_dict({"id": "t", "type": "task", "version": 2, "steps": [], "requires": ["agentenv-acme>=1.2"]}) + + assert task.requires == ["agentenv-acme>=1.2"] + assert task.to_dict()["requires"] == ["agentenv-acme>=1.2"] + + +_UNKNOWN = {"id": "t", "type": "task", "version": 2, "steps": [{"id": "s", "type": "acme_task_document_test"}]} + + +def test_a_step_type_nothing_registers_names_the_package_the_task_requires(monkeypatch): + monkeypatch.setattr(_requirements, "installed", {}.get) + + with pytest.raises(UnknownStepTypeError) as caught: + Task.from_dict({**_UNKNOWN, "requires": ["agentenv-acme>=1.2", "agentenv-other>=1"]}) + + assert str(caught.value) == ( + "Unknown task step type: acme_task_document_test; task t v2: requires 'agentenv-acme>=1.2', which isn't " + "installed; task t v2: requires 'agentenv-other>=1', which isn't installed; task t v2: install what it requires " + "with: agent-env plugin add 'agentenv-acme>=1.2' 'agentenv-other>=1'") + assert caught.value.requirements == ("agentenv-acme>=1.2", "agentenv-other>=1") + + +@pytest.mark.parametrize("record", [{}, {"requires": ["agentenv-acme>=1.2"]}], ids=["no-record", "met"]) +def test_without_an_unmet_requirement_it_names_only_the_type(monkeypatch, record): + monkeypatch.setattr(_requirements, "installed", {"agentenv-acme": "1.3"}.get) + + with pytest.raises(UnknownStepTypeError) as caught: + Task.from_dict({**_UNKNOWN, **record}) + + assert str(caught.value) == "Unknown task step type: acme_task_document_test" + assert caught.value.requirements == () From d19e6f95b51a1d86751bdda3798b9cfceb57616f Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Wed, 28 Oct 2026 22:15:03 -0700 Subject: [PATCH 2/6] fix(task): task run checks only the steps a run runs, as the run gives them A resumed run (--start-step) no longer reads what the steps it skips reference, and a step the run gives step_params, which it reads in place of what it names, isn't checked against its stored references. The infra envs follow the run's env state overrides, so an external state needs no service-db and an attached instance's own type decides. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/cli/task/_checks.py | 41 ++++++++++++++------- src/agent_env/cli/task/run.py | 15 ++++---- tst/unit/cli/task_run_requires_test.py | 51 +++++++++++++++++++++++--- 3 files changed, 80 insertions(+), 27 deletions(-) diff --git a/src/agent_env/cli/task/_checks.py b/src/agent_env/cli/task/_checks.py index 093a6db17..a5a76e87f 100644 --- a/src/agent_env/cli/task/_checks.py +++ b/src/agent_env/cli/task/_checks.py @@ -4,7 +4,8 @@ from __future__ import annotations -from collections.abc import Callable +from collections.abc import Callable, Mapping +from typing import Any import click @@ -19,17 +20,20 @@ from agent_env.store.document_store import UnknownDocumentTypeError from agent_env.task import Task from agent_env.task.task import UnknownStepTypeError +from agent_env.task_step.task_step import TaskStep +from agent_env.task_step.task_steps.deploy_env import DeployEnvTaskStep # The entities a step reads whose type a plugin may register. _GETTERS = {EntityKind.ENV: Env.get, EntityKind.ARTIFACT: Artifact.get} -def runnable_task(task_id: str, version: int | None, *, say: Callable[[str], None], gateway_env_id: str | None = None, - service_db_env_id: str | None = None) -> Task: +def runnable_task(task_id: str, version: int | None, metadata: Mapping[str, Any], *, + say: Callable[[str], None]) -> Task: """The stored task ``task_id`` at ``version``, or its latest, once what it requires is installed and every env and - artifact type its steps read loads here. With local stores, the infra envs it deploys on are set up, but for one - the run is given its own of (``gateway_env_id``, ``service_db_env_id``). Raises UnmetRequirements naming the - command that installs what it requires, or a ClickException for a type nothing here registers.""" + artifact type the steps a run with ``metadata`` runs read loads here. With local stores, the infra envs those steps + deploy on are set up, but for one the run is given its own of (``gateway_env_id``, ``service_db_env_id``). Raises + UnmetRequirements naming the command that installs what it requires, or a ClickException for a type nothing here + registers.""" try: task = Task.get(task_id, version=version) except UnknownStepTypeError as e: @@ -37,23 +41,34 @@ def runnable_task(task_id: str, version: int | None, *, say: Callable[[str], Non raise UnmetRequirements([str(e)], e.requirements) from e raise owner = f"task {task.id} v{task.version}" + overrides = metadata.get("user_overrides") or {} + steps = task.steps[metadata.get("start_step") or 0:] # the steps the run runs + params = overrides.get("step_params") or {} problems, unmet = unmet_problems(task.requires or (), owner) - problems += [f"{owner}: {problem}" for problem in _unknown_types(task)] + # A step the run gives step_params reads them its own way, in place of what it names. + problems += [f"{owner}: {problem}" for problem in _unknown_types([step for step in steps if step.id not in params])] if unmet: raise UnmetRequirements(problems, unmet) if problems: raise click.ClickException("\n".join(problems)) - overridden = {kind for kind, given in ((GATEWAY, gateway_env_id), (SERVICE_DB, service_db_env_id)) if given} - if stores_are_local() and (kinds := stored_infra(task.steps) - overridden): + given = {kind for kind, key in ((GATEWAY, "gateway_env_id"), (SERVICE_DB, "service_db_env_id")) if overrides.get(key)} + if stores_are_local() and (kinds := stored_infra(_with_state(steps, overrides)) - given): ensure_default_envs(kinds, say=say) return task -def _unknown_types(task: Task) -> list[str]: - """Each env and artifact a step of ``task`` reads, as the store holds it, whose type nothing here registers. One - the store doesn't hold, or whose id the run fills in, is left to its step.""" +def _with_state(steps: list[TaskStep], overrides: Mapping[str, Any]) -> list[TaskStep]: + """``steps``, each deploy_env step with the run's env state overrides, which it reads over its own.""" + state = {key: overrides[key] for key in ("env_state_type", "env_state_instance_id") if overrides.get(key)} + return [type(step).from_dict({**step.to_dict(), **state}) if state and isinstance(step, DeployEnvTaskStep) + else step for step in steps] + + +def _unknown_types(steps: list[TaskStep]) -> list[str]: + """Each env and artifact one of ``steps`` reads, as the store holds it, whose type nothing here registers. One the + store doesn't hold, or whose id the run fills in, is left to its step.""" found = [] - for step in task.steps: + for step in steps: for site in ref_sites(type(step).entity_refs or (), step.to_dict()): get = _GETTERS.get(site.ref.kind) if site.ref.role is RefRole.OUTPUT or get is None or not isinstance(site.value, str): diff --git a/src/agent_env/cli/task/run.py b/src/agent_env/cli/task/run.py index 67f62040a..4c8ba6186 100644 --- a/src/agent_env/cli/task/run.py +++ b/src/agent_env/cli/task/run.py @@ -420,7 +420,7 @@ def run(task_id: str, task_version: int | None, output_dir: str | None, k: int, initial_context.metadata["task_id"] = task_id click.echo(f"Fetching task: id={task_id} version={task_version or 'latest'}...") - task = _runnable(task_id, task_version, gateway_env_id, service_db_env_id) + task = _runnable(task_id, task_version, initial_context.metadata) click.echo(f"Found task: id={task.id} version={task.version} steps={len(task.steps)}") print_banner() @@ -465,13 +465,11 @@ def run(task_id: str, task_version: int | None, output_dir: str | None, k: int, raise SystemExit(1) -def _runnable(task_id: str, task_version: int | None, gateway_env_id: str | None = None, - service_db_env_id: str | None = None): - """The stored task, checked before its first step (``runnable_task``); a refusal for what it requires offers, in a - terminal, to install it and run the command again.""" +def _runnable(task_id: str, task_version: int | None, metadata: dict): + """The stored task, checked before the first step a run with ``metadata`` runs (``runnable_task``); a refusal for + what it requires offers, in a terminal, to install it and run the command again.""" try: - return runnable_task(task_id, task_version, say=click.echo, gateway_env_id=gateway_env_id, - service_db_env_id=service_db_env_id) + return runnable_task(task_id, task_version, metadata, say=click.echo) except UnmetRequirements as e: offer_to_install(e) raise @@ -532,7 +530,8 @@ def run_batch(task_id: str, task_version: int | None, seeds: str, concurrency: i raise click.ClickException("No seeds found in CSV") click.echo(f"Fetching task: id={task_id} version={task_version or 'latest'}...") - task = _runnable(task_id, task_version) + task = _runnable(task_id, task_version, {"user_overrides": {"env_state_type": env_state_type, + "env_state_instance_id": env_state_instance_id}}) click.echo(f"Found task: id={task.id} version={task.version} steps={len(task.steps)}") suffixes = {step.universe_id_suffix or "" for step in task.steps if isinstance(step, CollectArtifactsTaskStep)} if is_local_id(task.id) and suffixes: diff --git a/tst/unit/cli/task_run_requires_test.py b/tst/unit/cli/task_run_requires_test.py index fac65a999..6475856db 100644 --- a/tst/unit/cli/task_run_requires_test.py +++ b/tst/unit/cli/task_run_requires_test.py @@ -3,15 +3,20 @@ import logging +import click import pytest from click.testing import CliRunner +from agent_env.artifact.artifacts.docker_image import DockerImageArtifact +from agent_env.artifact.store import get_artifact_store from agent_env.cli import _install_prompt, cli, submitted_runs from agent_env.cli.task import _checks from agent_env.config import get_config +from agent_env.env import MCPServerEnv from agent_env.env.bootstrap import GATEWAY, SERVICE_DB from agent_env.env.store import ENVS_COLLECTION from agent_env.plugins import _requirements +from agent_env.store.routing import namespace_routing from agent_env.task import Task from agent_env.task.store import TASKS_COLLECTION @@ -92,13 +97,13 @@ def test_one_the_store_doesnt_hold_or_the_run_fills_in_is_left_to_its_step(store {"id": "seeded", "type": "deploy_env", "env_id": "{{seed.env}}"}]) monkeypatch.setattr(_checks, "stores_are_local", lambda: False) - assert _checks.runnable_task("t", None, say=print).id == "t" + assert _checks.runnable_task("t", None, {}, say=print).id == "t" def test_a_task_whose_requirements_are_met_is_run(stores, monkeypatch): Task.put(id="t", steps=[], requires=["agentenv-e2e-acme>=1"]) - assert _checks.runnable_task("t", None, say=print).requires == ["agentenv-e2e-acme>=1"] + assert _checks.runnable_task("t", None, {}, say=print).requires == ["agentenv-e2e-acme>=1"] @pytest.fixture @@ -146,16 +151,50 @@ def test_a_task_needing_a_newer_agent_env_says_to_upgrade_it_and_never_asks(term (True, {}, {GATEWAY, SERVICE_DB}), (True, {"gateway_env_id": "mine"}, {SERVICE_DB}), (True, {"gateway_env_id": "mine", "service_db_env_id": "mine-db"}, None), + (True, {"env_state_type": "external_db"}, {GATEWAY}), (False, {}, None), -], ids=["local", "own-gateway", "own-both", "shared"]) +], ids=["local", "own-gateway", "own-both", "external-state", "shared"]) def test_the_infra_envs_its_deploys_run_on_are_set_up_in_local_stores_but_for_those_it_is_given( stores, monkeypatch, local, overrides, kinds): - Task.put(id="t", steps=[]) + _mcp_env("crm") + _stored([{"id": "env", "type": "deploy_env", "env_id": "crm"}]) set_up = [] monkeypatch.setattr(_checks, "stores_are_local", lambda: local) - monkeypatch.setattr(_checks, "stored_infra", lambda steps: frozenset({GATEWAY, SERVICE_DB})) monkeypatch.setattr(_checks, "ensure_default_envs", lambda kinds, say: set_up.append(set(kinds))) - _checks.runnable_task("t", None, say=print, **overrides) + _checks.runnable_task("t", None, {"user_overrides": overrides}, say=print) assert set_up == ([kinds] if kinds else []) + + +def _mcp_env(id): + with namespace_routing(): + image = get_artifact_store().put_document(DockerImageArtifact( + id=f"{id}-image", description=id, image_name=f"{id}:latest", tar_gz_s3_url=None, platform=None)) + MCPServerEnv.put(id=id, docker_image_artifact=image, environment_name=id) + + +_UNKNOWN_ENV = [{"id": "env", "type": "deploy_env", "env_id": "acme-env"}, + {"id": "box", "type": "deploy_sandbox", "sandbox_name": "box", "sandbox_mode": "vm"}, + {"id": "load", "type": "load_artifact", "sandbox_name": "box", "artifact_id": "acme-data", + "depends_on": [{"task_step_id": "box"}]}] + + +@pytest.mark.parametrize("metadata", [ + {"start_step": 3}, + {"start_step": 1, "user_overrides": {"step_params": {"load": {"artifact_id": "absent"}}}}, +], ids=["resumed-past-them", "overridden"]) +def test_what_only_the_steps_a_run_skips_or_overrides_read_is_left_alone(stores, monkeypatch, metadata): + store = get_config().get_document_store() + store.insert(ENVS_COLLECTION, {"id": "acme-env", "type": "acme_env_task_run_test", "version": 1}) + store.insert("artifacts", {"id": "acme-data", "type": "acme_artifact_task_run_test", "version": 1}) + _stored(_UNKNOWN_ENV) + monkeypatch.setattr(_checks, "stores_are_local", lambda: False) + + with pytest.raises(click.ClickException) as caught: + _checks.runnable_task("t", None, {}, say=print) + assert _checks.runnable_task("t", None, metadata, say=print).id == "t" + env, load = caught.value.message.splitlines() + assert env == "task t v1: step 'env': env_id 'acme-env': Unknown env type: acme_env_task_run_test" + assert load.startswith("task t v1: step 'load': artifacts[0].id 'acme-data': Unknown artifact type " + "'acme_artifact_task_run_test'") From 29bc743c98dfcf004fe77f3fbb5c543bebc67d19 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Wed, 28 Oct 2026 22:28:47 -0700 Subject: [PATCH 3/6] fix(task): step_params skip only the references they replace A step's step_params that name one of its fields referencing a kind replace what it reads of that kind, so only those references are left to it; overrides that leave its references as stored, such as a load_artifact destination_path, still have them checked before the run. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/cli/task/_checks.py | 19 ++++++++++++------- tst/unit/cli/task_run_requires_test.py | 19 +++++++++++++++++-- 2 files changed, 29 insertions(+), 9 deletions(-) diff --git a/src/agent_env/cli/task/_checks.py b/src/agent_env/cli/task/_checks.py index a5a76e87f..ca4b9855e 100644 --- a/src/agent_env/cli/task/_checks.py +++ b/src/agent_env/cli/task/_checks.py @@ -45,8 +45,7 @@ def runnable_task(task_id: str, version: int | None, metadata: Mapping[str, Any] steps = task.steps[metadata.get("start_step") or 0:] # the steps the run runs params = overrides.get("step_params") or {} problems, unmet = unmet_problems(task.requires or (), owner) - # A step the run gives step_params reads them its own way, in place of what it names. - problems += [f"{owner}: {problem}" for problem in _unknown_types([step for step in steps if step.id not in params])] + problems += [f"{owner}: {problem}" for problem in _unknown_types(steps, params)] if unmet: raise UnmetRequirements(problems, unmet) if problems: @@ -64,14 +63,20 @@ def _with_state(steps: list[TaskStep], overrides: Mapping[str, Any]) -> list[Tas else step for step in steps] -def _unknown_types(steps: list[TaskStep]) -> list[str]: - """Each env and artifact one of ``steps`` reads, as the store holds it, whose type nothing here registers. One the - store doesn't hold, or whose id the run fills in, is left to its step.""" +def _unknown_types(steps: list[TaskStep], params: Mapping[str, Mapping[str, Any]]) -> list[str]: + """Each env and artifact one of ``steps`` reads, as the store holds it, whose type nothing here registers. A step's + ``params`` (the run's step_params) naming one of its fields that references a kind replace what it reads of that + kind, so those references are left to it, as is one the store doesn't hold or whose id the run fills in.""" found = [] for step in steps: - for site in ref_sites(type(step).entity_refs or (), step.to_dict()): + refs = type(step).entity_refs or () + given = set(params.get(step.id) or {}) + replaced = {ref.kind for ref in refs + if {ref.path.split(".")[0].removesuffix("[]"), ref.version_field} & given} + for site in ref_sites(refs, step.to_dict()): get = _GETTERS.get(site.ref.kind) - if site.ref.role is RefRole.OUTPUT or get is None or not isinstance(site.value, str): + if site.ref.role is RefRole.OUTPUT or site.ref.kind in replaced or get is None \ + or not isinstance(site.value, str): continue try: get(site.value, site.version if isinstance(site.version, int) else None) diff --git a/tst/unit/cli/task_run_requires_test.py b/tst/unit/cli/task_run_requires_test.py index 6475856db..9b293b839 100644 --- a/tst/unit/cli/task_run_requires_test.py +++ b/tst/unit/cli/task_run_requires_test.py @@ -183,8 +183,8 @@ def _mcp_env(id): @pytest.mark.parametrize("metadata", [ {"start_step": 3}, {"start_step": 1, "user_overrides": {"step_params": {"load": {"artifact_id": "absent"}}}}, -], ids=["resumed-past-them", "overridden"]) -def test_what_only_the_steps_a_run_skips_or_overrides_read_is_left_alone(stores, monkeypatch, metadata): +], ids=["resumed-past-them", "replaced"]) +def test_what_only_the_steps_a_run_skips_or_its_step_params_replace_read_is_left_alone(stores, monkeypatch, metadata): store = get_config().get_document_store() store.insert(ENVS_COLLECTION, {"id": "acme-env", "type": "acme_env_task_run_test", "version": 1}) store.insert("artifacts", {"id": "acme-data", "type": "acme_artifact_task_run_test", "version": 1}) @@ -198,3 +198,18 @@ def test_what_only_the_steps_a_run_skips_or_overrides_read_is_left_alone(stores, assert env == "task t v1: step 'env': env_id 'acme-env': Unknown env type: acme_env_task_run_test" assert load.startswith("task t v1: step 'load': artifacts[0].id 'acme-data': Unknown artifact type " "'acme_artifact_task_run_test'") + + +@pytest.mark.parametrize("params", [{"load": {"destination_path": "/tmp/data"}}, {"env": {"memory_mb": 4096}}], + ids=["load-elsewhere", "other-params"]) +def test_step_params_that_leave_what_a_step_references_as_stored_check_it(stores, monkeypatch, params): + store = get_config().get_document_store() + store.insert(ENVS_COLLECTION, {"id": "acme-env", "type": "acme_env_task_run_test", "version": 1}) + store.insert("artifacts", {"id": "acme-data", "type": "acme_artifact_task_run_test", "version": 1}) + _stored(_UNKNOWN_ENV) + monkeypatch.setattr(_checks, "stores_are_local", lambda: False) + + with pytest.raises(click.ClickException) as caught: + _checks.runnable_task("t", None, {"user_overrides": {"step_params": params}}, say=print) + + assert [line.split(":")[1].strip() for line in caught.value.message.splitlines()] == ["step 'env'", "step 'load'"] From badb9a52ea7249bb62fa93ce60622bd995e806c3 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Wed, 28 Oct 2026 23:05:38 -0700 Subject: [PATCH 4/6] fix(task): stored requires are read defensively, infra is set up only when missing, and imports keep recorded ranges A stored requires entry that isn't a pip requirement, or holds a character a terminal wouldn't show as it is, is named as unmet and never offered for install, where it used to crash the load. task run sets up only the infra envs a local store doesn't hold, leaving one it holds for agent-env up or run to rebuild. Each entity the steps reference is read once however many name it. bundle init --from-task writes the ranges the stored task records. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/bundle/importer.py | 4 ++- src/agent_env/bundle/requires.py | 3 +- src/agent_env/cli/task/_checks.py | 44 +++++++++++++++++------ src/agent_env/plugins/_requirements.py | 30 ++++++++++++---- src/agent_env/task/task.py | 2 +- tst/unit/bundle/requires_test.py | 9 +++++ tst/unit/cli/task_run_requires_test.py | 48 +++++++++++++++++++++++++- tst/unit/task/task_document_test.py | 18 ++++++++++ 8 files changed, 136 insertions(+), 22 deletions(-) diff --git a/src/agent_env/bundle/importer.py b/src/agent_env/bundle/importer.py index 4b7bd74da..d00824c3f 100644 --- a/src/agent_env/bundle/importer.py +++ b/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( diff --git a/src/agent_env/bundle/requires.py b/src/agent_env/bundle/requires.py index ebdeb6717..8d5d403ab 100644 --- a/src/agent_env/bundle/requires.py +++ b/src/agent_env/bundle/requires.py @@ -171,7 +171,8 @@ def detected(found: Iterable[Use], declared: Iterable[str] = ()) -> list[str]: that distribution, else at least the version installed here.""" by_name: dict[str, list[str]] = {} for line in declared: - by_name.setdefault(canonicalize_name(Requirement(line).name), []).append(line) + if (requirement := _requirements.parsed(line)) is not None: + by_name.setdefault(canonicalize_name(requirement.name), []).append(line) requires: set[str] = set() for use in found: if (requirement := use.provider.requirement()) is not None: diff --git a/src/agent_env/cli/task/_checks.py b/src/agent_env/cli/task/_checks.py index ca4b9855e..8249cab48 100644 --- a/src/agent_env/cli/task/_checks.py +++ b/src/agent_env/cli/task/_checks.py @@ -1,6 +1,6 @@ """What `agent-env task run` checks of a stored task before its first step: that the plugin packages it requires are installed and every env and artifact type its steps read resolves. With local stores it then sets up the infra envs -the task deploys on, as `agent-env run` does.""" +the task deploys on that the store doesn't hold.""" from __future__ import annotations @@ -13,7 +13,7 @@ from agent_env.bundle.preflight import stored_infra from agent_env.bundle.requires import UnmetRequirements from agent_env.entity.refs import EntityKind, RefRole, ref_sites -from agent_env.env.bootstrap import GATEWAY, SERVICE_DB, ensure_default_envs, stores_are_local +from agent_env.env.bootstrap import GATEWAY, SERVICE_DB, default_env_id, ensure_default_envs, stores_are_local from agent_env.env.env import Env from agent_env.plugins._requirements import unmet_problems from agent_env.store.base import NotFoundError @@ -31,7 +31,8 @@ def runnable_task(task_id: str, version: int | None, metadata: Mapping[str, Any] say: Callable[[str], None]) -> Task: """The stored task ``task_id`` at ``version``, or its latest, once what it requires is installed and every env and artifact type the steps a run with ``metadata`` runs read loads here. With local stores, the infra envs those steps - deploy on are set up, but for one the run is given its own of (``gateway_env_id``, ``service_db_env_id``). Raises + deploy on that the store doesn't hold are set up, but for one the run is given its own of (``gateway_env_id``, + ``service_db_env_id``); one it holds is left as it is, for ``agent-env up`` or ``run`` to rebuild. Raises UnmetRequirements naming the command that installs what it requires, or a ClickException for a type nothing here registers.""" try: @@ -51,11 +52,22 @@ def runnable_task(task_id: str, version: int | None, metadata: Mapping[str, Any] if problems: raise click.ClickException("\n".join(problems)) given = {kind for kind, key in ((GATEWAY, "gateway_env_id"), (SERVICE_DB, "service_db_env_id")) if overrides.get(key)} - if stores_are_local() and (kinds := stored_infra(_with_state(steps, overrides)) - given): + if stores_are_local() and (kinds := _missing(stored_infra(_with_state(steps, overrides)) - given)): ensure_default_envs(kinds, say=say) return task +def _missing(kinds: set[str] | frozenset[str]) -> set[str]: + """Those of ``kinds`` whose default infra env the store doesn't hold.""" + missing = set() + for kind in kinds: + try: + Env.get(default_env_id(kind)) + except NotFoundError: + missing.add(kind) + return missing + + def _with_state(steps: list[TaskStep], overrides: Mapping[str, Any]) -> list[TaskStep]: """``steps``, each deploy_env step with the run's env state overrides, which it reads over its own.""" state = {key: overrides[key] for key in ("env_state_type", "env_state_instance_id") if overrides.get(key)} @@ -67,7 +79,7 @@ def _unknown_types(steps: list[TaskStep], params: Mapping[str, Mapping[str, Any] """Each env and artifact one of ``steps`` reads, as the store holds it, whose type nothing here registers. A step's ``params`` (the run's step_params) naming one of its fields that references a kind replace what it reads of that kind, so those references are left to it, as is one the store doesn't hold or whose id the run fills in.""" - found = [] + found, read = [], {} # each entity's problem, read once however many steps name it for step in steps: refs = type(step).entity_refs or () given = set(params.get(step.id) or {}) @@ -78,10 +90,20 @@ def _unknown_types(steps: list[TaskStep], params: Mapping[str, Mapping[str, Any] if site.ref.role is RefRole.OUTPUT or site.ref.kind in replaced or get is None \ or not isinstance(site.value, str): continue - try: - get(site.value, site.version if isinstance(site.version, int) else None) - except UnknownDocumentTypeError as e: - found.append(f"step {step.id!r}: {site.path} {site.value!r}: {e}") - except (NotFoundError, ValueError): - continue + key = (site.ref.kind, site.value, site.version if isinstance(site.version, int) else None) + if key not in read: + read[key] = _unloadable(get, key[1], key[2]) + if read[key] is not None: + found.append(f"step {step.id!r}: {site.path} {site.value!r}: {read[key]}") return found + + +def _unloadable(get: Callable[..., Any], id: str, version: int | None) -> str | None: + """Why the entity ``id`` at ``version`` can't be loaded here: a type nothing registers; None otherwise.""" + try: + get(id, version) + except UnknownDocumentTypeError as e: + return str(e) + except (NotFoundError, ValueError): + pass + return None diff --git a/src/agent_env/plugins/_requirements.py b/src/agent_env/plugins/_requirements.py index 5c2bb88ab..d298340fd 100644 --- a/src/agent_env/plugins/_requirements.py +++ b/src/agent_env/plugins/_requirements.py @@ -54,11 +54,14 @@ def incompatibility(dist: Any) -> Optional[str]: def unmet(requires: Iterable[str]) -> list[tuple[str, str]]: - """Each of ``requires`` that isn't met here, with why: its package isn't installed, is at a version it excludes, or - is without what the extras it names need. One whose marker is false here doesn't apply.""" + """Each of ``requires`` that isn't met here, with why: it isn't a pip requirement this can show as written, its + package isn't installed, is at a version it excludes, or is without what the extras it names need. One whose marker + is false here doesn't apply.""" found = [] for line in requires: - requirement = Requirement(line) + if (requirement := parsed(line)) is None: + found.append((line, "which isn't a pip requirement that can be shown as written")) + continue if requirement.marker is not None and not requirement.marker.evaluate({"extra": ""}): continue have = installed(canonicalize_name(requirement.name)) @@ -77,10 +80,12 @@ def unmet_problems(requires: Iterable[str], owner: str) -> tuple[list[str], list command that installs every plugin package among them, and those requirements. One on agent-env itself, which ``plugin add`` doesn't change, says to upgrade it instead.""" found = unmet(requires) - core = [(line, why) for line, why in found if canonicalize_name(Requirement(line).name) in CORE] - plugins = [(line, why) for line, why in found if (line, why) not in core] - problems = [f"{owner}: requires {line!r}, {why}; upgrade agent-env with the installer that installed it" - for line, why in core] + names = {line: canonicalize_name(requirement.name) if (requirement := parsed(line)) else None for line, _ in found} + core = [(line, why) for line, why in found if names[line] in CORE] + plugins = [(line, why) for line, why in found if names[line] is not None and names[line] not in CORE] + problems = [f"{owner}: requires {line!r}, {why}" for line, why in found if names[line] is None] + problems += [f"{owner}: requires {line!r}, {why}; upgrade agent-env with the installer that installed it" + for line, why in core] lines = [line for line, _ in plugins] if len(plugins) == 1: (line, why), = plugins @@ -91,6 +96,17 @@ def unmet_problems(requires: Iterable[str], owner: str) -> tuple[list[str], list return problems, lines +def parsed(line: object) -> Requirement | None: + """``line`` as a pip requirement, or None when it isn't one, or holds a character a terminal wouldn't show as it + is, which could make the install it names look like another.""" + if not isinstance(line, str) or not line.isprintable(): + return None + try: + return Requirement(line) + except InvalidRequirement: + return None + + def install_command(requirements: Sequence[str]) -> str: """The command that installs ``requirements`` into this environment, through the installer that owns it.""" return shlex.join(["agent-env", "plugin", "add", *requirements]) diff --git a/src/agent_env/task/task.py b/src/agent_env/task/task.py index b6de1064e..605947ba7 100644 --- a/src/agent_env/task/task.py +++ b/src/agent_env/task/task.py @@ -425,7 +425,7 @@ def from_dict(cls, data: dict) -> "Task": id=data["id"], version=data.get("version"), steps=steps, - requires=data.get("requires"), + requires=data["requires"] if isinstance(data.get("requires"), list) else None, ) async def run( diff --git a/tst/unit/bundle/requires_test.py b/tst/unit/bundle/requires_test.py index 89ff53859..b085d64b4 100644 --- a/tst/unit/bundle/requires_test.py +++ b/tst/unit/bundle/requires_test.py @@ -19,6 +19,7 @@ from agent_env.bundle import BundleError, dry_run_bundle, parse_bundle, prepare_bundle, run_bundle from agent_env.bundle import register as register_module from agent_env.bundle import requires as requires_module +from agent_env.bundle.importer import import_task from agent_env.bundle.register import register_bundle from agent_env.bundle.requires import ( Declaring, Provider, Providers, UnmetRequirements, _create, declare, declaring_for, @@ -661,6 +662,14 @@ def test_a_package_bundle_toml_requires_keeps_its_range_and_a_new_range_writes_t ("b-plugin", 2, ("requires agentenv-acme>=1.1,<2",))] assert Task.get("b-plugin").requires == ["agentenv-acme>=1.1,<2"] + def test_an_import_of_a_registered_task_keeps_the_range_it_records(self, stores): + root = layout(stores / "b", {**REGISTERED, "bundle.toml": 'requires = ["agentenv-acme>=1.0,<2"]\n'}) + register_bundle(root) + + imported = import_task("b-plugin", stores / "imported") + + assert parse_bundle(imported.root).requires == ("agentenv-acme>=1.0,<2",) + def test_unmet_requirements_are_a_bundle_error(): assert issubclass(UnmetRequirements, BundleError) diff --git a/tst/unit/cli/task_run_requires_test.py b/tst/unit/cli/task_run_requires_test.py index 9b293b839..7f39b765e 100644 --- a/tst/unit/cli/task_run_requires_test.py +++ b/tst/unit/cli/task_run_requires_test.py @@ -13,9 +13,11 @@ from agent_env.cli.task import _checks from agent_env.config import get_config from agent_env.env import MCPServerEnv -from agent_env.env.bootstrap import GATEWAY, SERVICE_DB +from agent_env.entity.refs import EntityKind +from agent_env.env.bootstrap import GATEWAY, SERVICE_DB, default_env_id from agent_env.env.store import ENVS_COLLECTION from agent_env.plugins import _requirements +from agent_env.store.document_store import UnknownDocumentTypeError from agent_env.store.routing import namespace_routing from agent_env.task import Task from agent_env.task.store import TASKS_COLLECTION @@ -213,3 +215,47 @@ def test_step_params_that_leave_what_a_step_references_as_stored_check_it(stores _checks.runnable_task("t", None, {"user_overrides": {"step_params": params}}, say=print) assert [line.split(":")[1].strip() for line in caught.value.message.splitlines()] == ["step 'env'", "step 'load'"] + + +def test_a_reference_many_steps_make_is_read_once_and_named_at_each(stores, monkeypatch): + reads = [] + + def get(id, version): + reads.append((id, version)) + raise UnknownDocumentTypeError("Unknown artifact type 'acme_artifact_task_run_test'") + + monkeypatch.setitem(_checks._GETTERS, EntityKind.ARTIFACT, get) + _stored([{"id": "box", "type": "deploy_sandbox", "sandbox_name": "box", "sandbox_mode": "vm"}, + *({"id": f"load{i}", "type": "load_artifact", "sandbox_name": "box", "artifact_id": "acme-data", + "depends_on": [{"task_step_id": "box"}]} for i in range(3))]) + + with pytest.raises(click.ClickException) as caught: + _checks.runnable_task("t", None, {}, say=print) + + assert reads == [("acme-data", None)] + assert [line.split(":")[1].strip() for line in caught.value.message.splitlines()] == [ + "step 'load0'", "step 'load1'", "step 'load2'"] + + +def test_an_infra_env_the_store_holds_is_left_as_it_is(stores, monkeypatch): + _mcp_env("crm") + _mcp_env(default_env_id(GATEWAY)) # whatever it was built from + _stored([{"id": "env", "type": "deploy_env", "env_id": "crm"}]) + set_up = [] + monkeypatch.setattr(_checks, "stores_are_local", lambda: True) + monkeypatch.setattr(_checks, "ensure_default_envs", lambda kinds, say: set_up.append(set(kinds))) + + _checks.runnable_task("t", None, {}, say=print) + + assert set_up == [{SERVICE_DB}] + + +def test_a_record_that_isnt_a_requirement_is_refused_and_never_offered(terminal): + Task.put(id="t", steps=[], requires=["not a requirement!", MISSING]) + + result = _run(input="y\n") + + assert result.exit_code == 0, result.output + assert "task t v1: requires 'not a requirement!', which isn't a pip requirement that can be shown as written" in \ + result.output + assert terminal["added"] == [[MISSING]] diff --git a/tst/unit/task/task_document_test.py b/tst/unit/task/task_document_test.py index 80e779caa..605aaa770 100644 --- a/tst/unit/task/task_document_test.py +++ b/tst/unit/task/task_document_test.py @@ -46,3 +46,21 @@ def test_without_an_unmet_requirement_it_names_only_the_type(monkeypatch, record assert str(caught.value) == "Unknown task step type: acme_task_document_test" assert caught.value.requirements == () + + +@pytest.mark.parametrize("record", [["not a requirement!"], ["evil; os_name != \"\x1b[2K\ragentenv-browser\""], [7]], + ids=["unparseable", "terminal-escape", "not-a-string"]) +def test_a_record_that_isnt_a_requirement_is_named_and_never_offered(monkeypatch, record): + monkeypatch.setattr(_requirements, "installed", {}.get) + + with pytest.raises(UnknownStepTypeError) as caught: + Task.from_dict({**_UNKNOWN, "requires": [*record, "agentenv-acme>=1.2"]}) + + assert "isn't a pip requirement that can be shown as written" in str(caught.value) + assert caught.value.requirements == ("agentenv-acme>=1.2",) + assert "plugin add 'agentenv-acme>=1.2'" in str(caught.value) and repr(record[0]) in str(caught.value) + + +def test_a_record_that_isnt_a_list_is_read_as_none(): + assert Task.from_dict({"id": "t", "type": "task", "version": 1, "steps": [], "requires": "agentenv-acme"}).requires \ + is None From 75334f0012e70604d8b3af1ec940e4bd0d378cab Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Wed, 28 Oct 2026 23:46:19 -0700 Subject: [PATCH 5/6] refactor(task): read run overrides and infra policy where they're defined DeployEnvTaskStep.run_state is the one place a deploy's env state comes from, for its run, preflight and task run's checks; the gateway and service-db override keys are constants beside run_env_ids. Setting up only missing infra is ensure_default_envs(missing_only=True), decided with its other rules under the same lock. run-batch checks the task with the metadata it runs with. The checks use EntityRef.field, and the bundle and stored-task prompt tests share one terminal helper. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/bundle/preflight.py | 32 ++++++++++------- src/agent_env/cli/task/_checks.py | 31 ++++------------ src/agent_env/cli/task/run.py | 24 ++++++------- src/agent_env/config/runtime.py | 8 +++-- src/agent_env/env/bootstrap.py | 15 ++++---- .../task_step/task_steps/deploy_env.py | 14 +++++--- tst/unit/_terminal.py | 25 +++++++++++++ tst/unit/bundle/requires_test.py | 18 ++-------- tst/unit/cli/task_run_requires_test.py | 35 +++++++------------ tst/unit/env/bootstrap_test.py | 10 ++++++ 10 files changed, 111 insertions(+), 101 deletions(-) create mode 100644 tst/unit/_terminal.py diff --git a/src/agent_env/bundle/preflight.py b/src/agent_env/bundle/preflight.py index c05f96c6f..10dc4d11a 100644 --- a/src/agent_env/bundle/preflight.py +++ b/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 @@ -504,9 +505,10 @@ def _problem(self, where: str, problem: str) -> None: self.problems.setdefault(problem, {})[where] = None -def stored_infra(steps: Iterable[TaskStep]) -> frozenset[str]: - """The infra env kinds a stored task's ``steps`` deploy on through the gateway, 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.""" +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: @@ -516,14 +518,16 @@ def stored_infra(steps: Iterable[TaskStep]) -> frozenset[str]: 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) + kinds |= _gateway_infra(step, deployment, user_overrides or {}) return frozenset(kinds) -def _gateway_infra(step: DeployEnvTaskStep, deployment: _Deployment) -> set[str]: - """The infra env kinds ``step``'s deploy of ``deployment`` through the gateway runs on.""" +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) == LOCAL_POSTGRES_STATE_TYPE: + if _state_type(step, user_overrides or {}) == LOCAL_POSTGRES_STATE_TYPE: kinds.add(SERVICE_DB) if deployment.websites: kinds.add(WEBSITE_BROWSER) @@ -555,12 +559,14 @@ def _stored_images(env: Env) -> tuple[list[_Image], bool]: 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 diff --git a/src/agent_env/cli/task/_checks.py b/src/agent_env/cli/task/_checks.py index 8249cab48..0782db276 100644 --- a/src/agent_env/cli/task/_checks.py +++ b/src/agent_env/cli/task/_checks.py @@ -13,7 +13,8 @@ from agent_env.bundle.preflight import stored_infra from agent_env.bundle.requires import UnmetRequirements from agent_env.entity.refs import EntityKind, RefRole, ref_sites -from agent_env.env.bootstrap import GATEWAY, SERVICE_DB, default_env_id, ensure_default_envs, stores_are_local +from agent_env.config.runtime import GATEWAY_OVERRIDE, SERVICE_DB_OVERRIDE +from agent_env.env.bootstrap import GATEWAY, SERVICE_DB, ensure_default_envs, stores_are_local from agent_env.env.env import Env from agent_env.plugins._requirements import unmet_problems from agent_env.store.base import NotFoundError @@ -21,7 +22,6 @@ from agent_env.task import Task from agent_env.task.task import UnknownStepTypeError from agent_env.task_step.task_step import TaskStep -from agent_env.task_step.task_steps.deploy_env import DeployEnvTaskStep # The entities a step reads whose type a plugin may register. _GETTERS = {EntityKind.ENV: Env.get, EntityKind.ARTIFACT: Artifact.get} @@ -51,30 +51,12 @@ def runnable_task(task_id: str, version: int | None, metadata: Mapping[str, Any] raise UnmetRequirements(problems, unmet) if problems: raise click.ClickException("\n".join(problems)) - given = {kind for kind, key in ((GATEWAY, "gateway_env_id"), (SERVICE_DB, "service_db_env_id")) if overrides.get(key)} - if stores_are_local() and (kinds := _missing(stored_infra(_with_state(steps, overrides)) - given)): - ensure_default_envs(kinds, say=say) + given = {kind for kind, key in ((GATEWAY, GATEWAY_OVERRIDE), (SERVICE_DB, SERVICE_DB_OVERRIDE)) if overrides.get(key)} + if stores_are_local() and (kinds := stored_infra(steps, overrides) - given): + ensure_default_envs(kinds, say=say, missing_only=True) return task -def _missing(kinds: set[str] | frozenset[str]) -> set[str]: - """Those of ``kinds`` whose default infra env the store doesn't hold.""" - missing = set() - for kind in kinds: - try: - Env.get(default_env_id(kind)) - except NotFoundError: - missing.add(kind) - return missing - - -def _with_state(steps: list[TaskStep], overrides: Mapping[str, Any]) -> list[TaskStep]: - """``steps``, each deploy_env step with the run's env state overrides, which it reads over its own.""" - state = {key: overrides[key] for key in ("env_state_type", "env_state_instance_id") if overrides.get(key)} - return [type(step).from_dict({**step.to_dict(), **state}) if state and isinstance(step, DeployEnvTaskStep) - else step for step in steps] - - def _unknown_types(steps: list[TaskStep], params: Mapping[str, Mapping[str, Any]]) -> list[str]: """Each env and artifact one of ``steps`` reads, as the store holds it, whose type nothing here registers. A step's ``params`` (the run's step_params) naming one of its fields that references a kind replace what it reads of that @@ -83,8 +65,7 @@ def _unknown_types(steps: list[TaskStep], params: Mapping[str, Mapping[str, Any] for step in steps: refs = type(step).entity_refs or () given = set(params.get(step.id) or {}) - replaced = {ref.kind for ref in refs - if {ref.path.split(".")[0].removesuffix("[]"), ref.version_field} & given} + replaced = {ref.kind for ref in refs if {ref.field, ref.version_field} & given} for site in ref_sites(refs, step.to_dict()): get = _GETTERS.get(site.ref.kind) if site.ref.role is RefRole.OUTPUT or site.ref.kind in replaced or get is None \ diff --git a/src/agent_env/cli/task/run.py b/src/agent_env/cli/task/run.py index 4c8ba6186..7473da0a2 100644 --- a/src/agent_env/cli/task/run.py +++ b/src/agent_env/cli/task/run.py @@ -15,6 +15,7 @@ from agent_env.cli._install_prompt import offer_to_install from agent_env.cli.banner import print_banner from agent_env.cli.task._checks import runnable_task +from agent_env.config.runtime import GATEWAY_OVERRIDE, SERVICE_DB_OVERRIDE from agent_env.cli import submitted_runs from agent_env.cli.identity import get_agent_env_client_id from agent_env.cli.teardown_output import echo_teardown @@ -412,9 +413,9 @@ def run(task_id: str, task_version: int | None, output_dir: str | None, k: int, if env_state_instance_id: initial_context.metadata.setdefault("user_overrides", {})["env_state_instance_id"] = env_state_instance_id if gateway_env_id: - initial_context.metadata.setdefault("user_overrides", {})["gateway_env_id"] = gateway_env_id + initial_context.metadata.setdefault("user_overrides", {})[GATEWAY_OVERRIDE] = gateway_env_id if service_db_env_id: - initial_context.metadata.setdefault("user_overrides", {})["service_db_env_id"] = service_db_env_id + initial_context.metadata.setdefault("user_overrides", {})[SERVICE_DB_OVERRIDE] = service_db_env_id initial_context.metadata["start_step"] = start_step initial_context.metadata[KEEP_KEY] = keep initial_context.metadata["task_id"] = task_id @@ -529,9 +530,16 @@ def run_batch(task_id: str, task_version: int | None, seeds: str, concurrency: i if not seed_rows: raise click.ClickException("No seeds found in CSV") + metadata = {"run_group_id": uuid4().hex, "task_id": task_id, KEEP_KEY: keep} + _stamp_agent_env_client_metadata(metadata, get_agent_env_client_id()) + given = {"litellm_api_key": litellm_api_key, "judge_litellm_api_key": judge_litellm_api_key, + "apply_trajectory_filter": apply_trajectory_filter, "user_model": user_model, "judge_model": judge_model, + "agent_sandbox": agent_sandbox, "env_sandbox": env_sandbox, "env_state_type": env_state_type, + "env_state_instance_id": env_state_instance_id} + if overrides := {key: value for key, value in given.items() if value is not None and value != ""}: + metadata["user_overrides"] = overrides click.echo(f"Fetching task: id={task_id} version={task_version or 'latest'}...") - task = _runnable(task_id, task_version, {"user_overrides": {"env_state_type": env_state_type, - "env_state_instance_id": env_state_instance_id}}) + task = _runnable(task_id, task_version, metadata) click.echo(f"Found task: id={task.id} version={task.version} steps={len(task.steps)}") suffixes = {step.universe_id_suffix or "" for step in task.steps if isinstance(step, CollectArtifactsTaskStep)} if is_local_id(task.id) and suffixes: @@ -547,14 +555,6 @@ def run_batch(task_id: str, task_version: int | None, seeds: str, concurrency: i print_banner() - metadata = {"run_group_id": uuid4().hex, "task_id": task_id, KEEP_KEY: keep} - _stamp_agent_env_client_metadata(metadata, get_agent_env_client_id()) - given = {"litellm_api_key": litellm_api_key, "judge_litellm_api_key": judge_litellm_api_key, - "apply_trajectory_filter": apply_trajectory_filter, "user_model": user_model, "judge_model": judge_model, - "agent_sandbox": agent_sandbox, "env_sandbox": env_sandbox, "env_state_type": env_state_type, - "env_state_instance_id": env_state_instance_id} - if overrides := {key: value for key, value in given.items() if value is not None and value != ""}: - metadata["user_overrides"] = overrides submissions = [] for index, seed in enumerate(seed_rows, 1): diff --git a/src/agent_env/config/runtime.py b/src/agent_env/config/runtime.py index 952d931df..056aead3f 100644 --- a/src/agent_env/config/runtime.py +++ b/src/agent_env/config/runtime.py @@ -1032,6 +1032,8 @@ def get_runner(): _RUN_ENV_IDS: ContextVar[Mapping[str, str]] = ContextVar("agent_env_run_env_ids", default={}) +# The user_overrides keys naming the gateway and service-db envs a run deploys on, over the configured defaults. +GATEWAY_OVERRIDE, SERVICE_DB_OVERRIDE = "gateway_env_id", "service_db_env_id" @contextlib.contextmanager @@ -1039,7 +1041,7 @@ def run_env_ids(user_overrides: Mapping[str, Any]) -> Iterator[None]: """Within a run, the ``gateway_env_id`` and ``service_db_env_id`` its ``user_overrides`` name win over the configured defaults, in whichever process runs it.""" token = _RUN_ENV_IDS.set( - {key: user_overrides[key] for key in ("gateway_env_id", "service_db_env_id") if user_overrides.get(key)}) + {key: user_overrides[key] for key in (GATEWAY_OVERRIDE, SERVICE_DB_OVERRIDE) if user_overrides.get(key)}) try: yield finally: @@ -1048,12 +1050,12 @@ def run_env_ids(user_overrides: Mapping[str, Any]) -> Iterator[None]: def gateway_env_id(config: Config) -> str: """The gateway env a deploy uses: the run's override, else ``config.default_gateway_env_id``.""" - return _RUN_ENV_IDS.get().get("gateway_env_id") or config.default_gateway_env_id + return _RUN_ENV_IDS.get().get(GATEWAY_OVERRIDE) or config.default_gateway_env_id def service_db_env_id(config: Config) -> str: """The service-db env a deploy uses: the run's override, else ``config.default_service_db_env_id``.""" - return _RUN_ENV_IDS.get().get("service_db_env_id") or config.default_service_db_env_id + return _RUN_ENV_IDS.get().get(SERVICE_DB_OVERRIDE) or config.default_service_db_env_id def get_config() -> Config: diff --git a/src/agent_env/env/bootstrap.py b/src/agent_env/env/bootstrap.py index 7cdfc1a59..a72cf1885 100644 --- a/src/agent_env/env/bootstrap.py +++ b/src/agent_env/env/bootstrap.py @@ -49,6 +49,7 @@ WEBSITE_BROWSER_CONTEXT = WEBSITE_BROWSER_DOCKERFILE.parent GATEWAY, SERVICE_DB, WEBSITE_BROWSER = "gateway", "service-db", "website-browser" # the kinds, named as their commands +MISSING = "missing" # the reason an infra env the store doesn't hold is set up def _quiet(line: str) -> None: @@ -268,7 +269,7 @@ def infra_to_set_up(kinds: Iterable[str]) -> list[InfraSetup]: env = Env.get(env_id) except NotFoundError: if local: - setups.append(_pulled_or_built(kind, env_id, "missing")) + setups.append(_pulled_or_built(kind, env_id, MISSING)) else: problems.append(f"the {kind} env {env_id!r} isn't in the store, and agent-env builds infra envs only into " f"local stores; put it with `{put_command(kind)}`") @@ -298,16 +299,18 @@ def needs_docker(builds: Iterable[InfraSetup], reason: str) -> str: return f"building the {', '.join(build.kind for build in builds)} env needs docker, and {reason}" -def ensure_default_envs(kinds: Iterable[str], *, say: Callable[[str], None] = _quiet) -> list[InfraSetup]: - """Set up the infra envs of ``kinds`` that ``infra_to_set_up`` names, each id locked so concurrent runs set it up - once, and return them. Raises InfraError before setting any up when one can't be: missing from stores that aren't - local, or to be built with docker unreachable.""" +def ensure_default_envs(kinds: Iterable[str], *, say: Callable[[str], None] = _quiet, + missing_only: bool = False) -> list[InfraSetup]: + """Set up the infra envs of ``kinds`` that ``infra_to_set_up`` names, or with ``missing_only`` those of them the + store doesn't hold, leaving one it holds as it is, each id locked so concurrent runs set it up once, and return + them. Raises InfraError before setting any up when one can't be: missing from stores that aren't local, or to be + built with docker unreachable.""" kinds = _in_order(kinds) if not kinds: return [] ids = [default_env_id(kind) for kind in kinds] with holding_locks(ids, on_wait=lambda: say("waiting for another agent-env run to finish setting up the infra envs")): - setups = infra_to_set_up(kinds) + setups = [setup for setup in infra_to_set_up(kinds) if not missing_only or setup.reason == MISSING] if (built := [setup for setup in setups if setup.builds]) and (reason := docker_unreachable()): raise InfraError([needs_docker(built, reason)]) for setup in setups: diff --git a/src/agent_env/task_step/task_steps/deploy_env.py b/src/agent_env/task_step/task_steps/deploy_env.py index 4efc589cc..375a4c515 100644 --- a/src/agent_env/task_step/task_steps/deploy_env.py +++ b/src/agent_env/task_step/task_steps/deploy_env.py @@ -4,7 +4,8 @@ import inspect import logging -from typing import ClassVar, Optional +from collections.abc import Mapping +from typing import Any, ClassVar, Optional from agent_env.config import get_config from agent_env.env.gateway import GatewayMode @@ -66,6 +67,12 @@ def __init__( # The env type's own deploy options, which this step passes to its deploy() without reading them. self.env_params = env_params + def run_state(self, user_overrides: Mapping[str, Any]) -> tuple[str | None, str | None]: + """The env state type and attached state instance a run with ``user_overrides`` deploys on: each the run's, + over this step's own.""" + return (user_overrides.get("env_state_type") or self.env_state_type, + user_overrides.get("env_state_instance_id") or self.env_state_instance_id) + def to_dict(self) -> dict: base = super().to_dict() base["env_id"] = self.env_id @@ -138,10 +145,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: env = Env.get(self.env_id, self.env_version) user_overrides = context.metadata.get("user_overrides", {}) resolved_sandbox_type = user_overrides.get("env_sandbox") or self.sandbox_type - resolved_env_state_type = user_overrides.get("env_state_type") or self.env_state_type - resolved_env_state_instance_id = ( - user_overrides.get("env_state_instance_id") or self.env_state_instance_id - ) + resolved_env_state_type, resolved_env_state_instance_id = self.run_state(user_overrides) _ttl_override = user_overrides.get("ttl_seconds") resolved_ttl = _ttl_override if _ttl_override is not None else self.ttl_seconds attribution = deploy_attribution(self, context) diff --git a/tst/unit/_terminal.py b/tst/unit/_terminal.py new file mode 100644 index 000000000..813e54293 --- /dev/null +++ b/tst/unit/_terminal.py @@ -0,0 +1,25 @@ +"""A person at a terminal for the install prompt, which the bundle and stored-task commands share.""" + +import os + +from agent_env.cli import _install_prompt + + +def at_a_terminal(monkeypatch) -> dict: + """Put stdin and stdout at a terminal, with an install that records what it's asked for and returns + ``calls["code"]``, and a re-exec that records its argv and the rerun marker, then stops the command; the calls.""" + calls = {"added": [], "exec": [], "code": 0} + + def add(specs, **kwargs): + calls["added"].append((specs, kwargs["yes"])) + return calls["code"] + + def execv(path, argv): + calls["exec"].append((path, argv, os.environ.get(_install_prompt._ASKED))) + raise SystemExit(0) + + monkeypatch.setattr(_install_prompt, "_interactive", lambda: True) + monkeypatch.setattr(_install_prompt._plugin_changes, "add", add) + monkeypatch.setattr(_install_prompt.os, "execv", execv) + monkeypatch.setenv(_install_prompt._ASKED, "") + return calls diff --git a/tst/unit/bundle/requires_test.py b/tst/unit/bundle/requires_test.py index b085d64b4..949670aec 100644 --- a/tst/unit/bundle/requires_test.py +++ b/tst/unit/bundle/requires_test.py @@ -34,6 +34,7 @@ from agent_env.store.routing import namespace_routing from agent_env.task import Task from agent_env.task_step.task_step import TaskStep +from tst.unit._terminal import at_a_terminal from tst.unit.bundle._support import fake_plugins, layout ACME = "agentenv-acme" @@ -496,21 +497,8 @@ def test_the_cli_refuses_a_run_whose_requirements_arent_met_without_a_terminal(s @pytest.fixture def terminal(stores, monkeypatch): - """A person at a terminal, an install that records what it's asked for, and a re-exec that stops the test.""" - calls = {"added": [], "exec": []} - - def add(specs, **kwargs): - calls["added"].append((specs, kwargs["yes"])) - return calls.get("code", 0) - - def execv(path, argv): - calls["exec"].append((path, argv, os.environ.get(_install_prompt._ASKED))) - raise SystemExit(0) - - monkeypatch.setattr(_install_prompt, "_interactive", lambda: True) - monkeypatch.setattr(_install_prompt._plugin_changes, "add", add) - monkeypatch.setattr(_install_prompt.os, "execv", execv) - monkeypatch.setenv(_install_prompt._ASKED, "") + """At a terminal, a bundle requiring what isn't installed, run by its path.""" + calls = at_a_terminal(monkeypatch) root = _bundle(stores / "b", _Acme.type, toml='requires = ["agentenv-missing>=1"]\n') monkeypatch.setattr(sys, "argv", ["agent-env", "run", str(root)]) return root, calls diff --git a/tst/unit/cli/task_run_requires_test.py b/tst/unit/cli/task_run_requires_test.py index 7f39b765e..6c4fb6221 100644 --- a/tst/unit/cli/task_run_requires_test.py +++ b/tst/unit/cli/task_run_requires_test.py @@ -9,18 +9,19 @@ from agent_env.artifact.artifacts.docker_image import DockerImageArtifact from agent_env.artifact.store import get_artifact_store -from agent_env.cli import _install_prompt, cli, submitted_runs +from agent_env.cli import cli, submitted_runs from agent_env.cli.task import _checks from agent_env.config import get_config from agent_env.env import MCPServerEnv from agent_env.entity.refs import EntityKind -from agent_env.env.bootstrap import GATEWAY, SERVICE_DB, default_env_id +from agent_env.env.bootstrap import GATEWAY, SERVICE_DB from agent_env.env.store import ENVS_COLLECTION from agent_env.plugins import _requirements from agent_env.store.document_store import UnknownDocumentTypeError from agent_env.store.routing import namespace_routing from agent_env.task import Task from agent_env.task.store import TASKS_COLLECTION +from tst.unit._terminal import at_a_terminal MISSING = "agentenv-e2e-missing>=1" REFUSED = f"task t v1: requires '{MISSING}', which isn't installed; install it with: agent-env plugin add '{MISSING}'" @@ -110,18 +111,7 @@ def test_a_task_whose_requirements_are_met_is_run(stores, monkeypatch): @pytest.fixture def terminal(stores, monkeypatch): - """A person at a terminal, an install that records what it's asked for, and a re-exec that stops the test.""" - calls = {"added": [], "exec": []} - - def execv(path, argv): - calls["exec"].append(argv) - raise SystemExit(0) - - monkeypatch.setattr(_install_prompt, "_interactive", lambda: True) - monkeypatch.setattr(_install_prompt._plugin_changes, "add", lambda specs, **kwargs: calls["added"].append(specs)) - monkeypatch.setattr(_install_prompt.os, "execv", execv) - monkeypatch.setenv(_install_prompt._ASKED, "") - return calls + return at_a_terminal(monkeypatch) @pytest.mark.parametrize("store", [ @@ -135,7 +125,7 @@ def test_at_a_terminal_it_offers_to_install_what_the_task_requires_and_runs_it_a assert result.exit_code == 0, result.output assert f"Install {MISSING} with agent-env plugin add? [y/N]: y" in result.output - assert terminal["added"] == [[MISSING]] and len(terminal["exec"]) == 1 + assert terminal["added"] == [([MISSING], True)] and len(terminal["exec"]) == 1 def test_a_task_needing_a_newer_agent_env_says_to_upgrade_it_and_never_asks(terminal): @@ -162,11 +152,12 @@ def test_the_infra_envs_its_deploys_run_on_are_set_up_in_local_stores_but_for_th _stored([{"id": "env", "type": "deploy_env", "env_id": "crm"}]) set_up = [] monkeypatch.setattr(_checks, "stores_are_local", lambda: local) - monkeypatch.setattr(_checks, "ensure_default_envs", lambda kinds, say: set_up.append(set(kinds))) + monkeypatch.setattr(_checks, "ensure_default_envs", + lambda kinds, say, missing_only: set_up.append((set(kinds), missing_only))) _checks.runnable_task("t", None, {"user_overrides": overrides}, say=print) - assert set_up == ([kinds] if kinds else []) + assert set_up == ([(kinds, True)] if kinds else []) def _mcp_env(id): @@ -237,17 +228,17 @@ def get(id, version): "step 'load0'", "step 'load1'", "step 'load2'"] -def test_an_infra_env_the_store_holds_is_left_as_it_is(stores, monkeypatch): +def test_the_infra_envs_are_set_up_only_where_the_store_doesnt_hold_them(stores, monkeypatch): _mcp_env("crm") - _mcp_env(default_env_id(GATEWAY)) # whatever it was built from _stored([{"id": "env", "type": "deploy_env", "env_id": "crm"}]) set_up = [] monkeypatch.setattr(_checks, "stores_are_local", lambda: True) - monkeypatch.setattr(_checks, "ensure_default_envs", lambda kinds, say: set_up.append(set(kinds))) + monkeypatch.setattr(_checks, "ensure_default_envs", + lambda kinds, say, missing_only: set_up.append((set(kinds), missing_only))) _checks.runnable_task("t", None, {}, say=print) - assert set_up == [{SERVICE_DB}] + assert set_up == [({GATEWAY, SERVICE_DB}, True)] def test_a_record_that_isnt_a_requirement_is_refused_and_never_offered(terminal): @@ -258,4 +249,4 @@ def test_a_record_that_isnt_a_requirement_is_refused_and_never_offered(terminal) assert result.exit_code == 0, result.output assert "task t v1: requires 'not a requirement!', which isn't a pip requirement that can be shown as written" in \ result.output - assert terminal["added"] == [[MISSING]] + assert terminal["added"] == [([MISSING], True)] diff --git a/tst/unit/env/bootstrap_test.py b/tst/unit/env/bootstrap_test.py index 3d7b4771b..507f3b771 100644 --- a/tst/unit/env/bootstrap_test.py +++ b/tst/unit/env/bootstrap_test.py @@ -106,6 +106,16 @@ def test_infra_built_from_other_inputs_is_rebuilt_once(puts, monkeypatch): assert len(puts) == 2 +def test_missing_only_sets_up_what_the_store_lacks_and_leaves_what_it_holds_as_it_is(puts, monkeypatch): + ensure_default_envs([GATEWAY]) + monkeypatch.setattr(bootstrap, "build_inputs_digest", lambda kind: "this release's") + + builds = ensure_default_envs([GATEWAY, SERVICE_DB], missing_only=True) + + assert [str(build) for build in builds] == ["service-db env 'default-db' (missing)"] + assert len(puts) == 2 + + def test_infra_built_before_its_inputs_were_recorded_is_rebuilt(puts): artifact = get_artifact_store().put_document(DockerImageArtifact( id="gateway-default", description="g", image_name="g:v1", tar_gz_s3_url="file:///g.tar.gz")) From 7e2c43bf1d037420402dfc9abfeb8f6bde9f6d33 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Wed, 28 Oct 2026 23:59:17 -0700 Subject: [PATCH 6/6] perf(env): missing-only infra setup reads an env the store holds no further ensure_default_envs(missing_only=True) used to weigh each held infra env (hashing its build inputs, reading the registry for its images) before dropping it, which a task run on a source checkout paid at every start. infra_to_set_up now skips a held one as soon as it finds it. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/env/bootstrap.py | 9 +++++---- tst/unit/env/bootstrap_test.py | 3 +++ 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/src/agent_env/env/bootstrap.py b/src/agent_env/env/bootstrap.py index a72cf1885..b08ceac47 100644 --- a/src/agent_env/env/bootstrap.py +++ b/src/agent_env/env/bootstrap.py @@ -254,10 +254,11 @@ def put_command(kind: str) -> str: return f"agent-env env {kind} put --id {default_env_id(kind)}" -def infra_to_set_up(kinds: Iterable[str]) -> list[InfraSetup]: +def infra_to_set_up(kinds: Iterable[str], *, missing_only: bool = False) -> list[InfraSetup]: """The infra envs of ``kinds`` a run sets up, in local stores only: stock ones missing or built from other inputs, pulled when their images are published for these sources and built otherwise, and stock ones built here switched - once they're published. Raises InfraError naming the put command for each one missing from stores not all local.""" + once they're published; with ``missing_only``, only those the store doesn't hold, one it holds read no further. + Raises InfraError naming the put command for each one missing from stores not all local.""" kinds = _in_order(kinds) if not kinds: return [] @@ -274,7 +275,7 @@ def infra_to_set_up(kinds: Iterable[str]) -> list[InfraSetup]: problems.append(f"the {kind} env {env_id!r} isn't in the store, and agent-env builds infra envs only into " f"local stores; put it with `{put_command(kind)}`") continue - if not (local and _built_from_stock(env, kind)): + if missing_only or not (local and _built_from_stock(env, kind)): continue recorded = (env.metadata or {}).get(BUILD_INPUTS_KEY) kept = {key: value for key, value in (env.metadata or {}).items() if not detected(key) and key != BUILD_INPUTS_KEY} @@ -310,7 +311,7 @@ def ensure_default_envs(kinds: Iterable[str], *, say: Callable[[str], None] = _q return [] ids = [default_env_id(kind) for kind in kinds] with holding_locks(ids, on_wait=lambda: say("waiting for another agent-env run to finish setting up the infra envs")): - setups = [setup for setup in infra_to_set_up(kinds) if not missing_only or setup.reason == MISSING] + setups = infra_to_set_up(kinds, missing_only=missing_only) if (built := [setup for setup in setups if setup.builds]) and (reason := docker_unreachable()): raise InfraError([needs_docker(built, reason)]) for setup in setups: diff --git a/tst/unit/env/bootstrap_test.py b/tst/unit/env/bootstrap_test.py index 507f3b771..c583f21c3 100644 --- a/tst/unit/env/bootstrap_test.py +++ b/tst/unit/env/bootstrap_test.py @@ -110,6 +110,9 @@ def test_missing_only_sets_up_what_the_store_lacks_and_leaves_what_it_holds_as_i ensure_default_envs([GATEWAY]) monkeypatch.setattr(bootstrap, "build_inputs_digest", lambda kind: "this release's") + # Before what an env it holds is built from is weighed: its files hashed, the registry read for its images. + monkeypatch.setattr(bootstrap, "_built_from_stock", lambda env, kind: pytest.fail("weighed the env it holds")) + builds = ensure_default_envs([GATEWAY, SERVICE_DB], missing_only=True) assert [str(build) for build in builds] == ["service-db env 'default-db' (missing)"]