diff --git a/AGENTS.md b/AGENTS.md index 60ede76c..ea6e2f8f 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -154,7 +154,10 @@ Out-of-tree code plugs in either through `config.toml`, as a `module:Class` poin outside a bundle run gets it from `A2AAgent.get`, which reads the store first and an alias after (the planner reads the store itself, as `STORE_GETTERS` does). agent-env's own bundles' aliases are theirs, and an alias two other bundles give is refused, naming both. `agent-env a2a-agent list` lists the installed agents with their aliases, then - the store's; `a2a-agent get` prints the agent an alias names. + the store's; `a2a-agent get` prints the agent an alias names. `agent-env run --agent ` deploys that agent for + the solver for one run (each step deploying or installing an agent a `prompt_agent` step prompts, not its user + simulator: `plan.solver_steps`), pulled in and checked as a task's reference to it would be, through per-step + `step_params` overrides, so the tasks stay as written. - 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/agents/README.md b/src/agent_env/agents/README.md index b4f66a09..5bb8c733 100644 --- a/src/agent_env/agents/README.md +++ b/src/agent_env/agents/README.md @@ -6,7 +6,8 @@ turn's trajectory. It's the default agent, so a `deploy_agent` or an agent judge from another bundle as `claude-code`, the alias `aliases.toml` gives it, or by its id, `@local/agentenv-framework/agents/claude-code`; an agent of your own with the id `claude-code`, in your bundle or your store, is the one `claude-code` names instead. `agent-env a2a-agent list` lists the agents a task can name, with -their aliases. +their aliases, and `agent-env run --agent pi` runs a bundle's solver as another of them for one run, its +tasks left as written. It needs Docker, and a model endpoint serving Anthropic's Messages API, `/v1/messages`, which a LiteLLM gateway does: the `[model]` base URL and key reach it as `ANTHROPIC_BASE_URL` and a bearer token, and `[model.headers]` are sent diff --git a/src/agent_env/bundle/materialize.py b/src/agent_env/bundle/materialize.py index 4abb3dd5..446e6e50 100644 --- a/src/agent_env/bundle/materialize.py +++ b/src/agent_env/bundle/materialize.py @@ -127,7 +127,7 @@ def through_ledger(write: Write, write_fn: Callable[[], int]) -> None: for write in tasks: with _noted(plan, write, "preflighting"): built[write.id] = _build(write) - found, skipped = _preflight(write, built[write.id], unwritten) + found, skipped = _preflight(plan, write, built[write.id], unwritten) problems.extend(f"{_path(plan, write)}: {problem}" for problem in found) unchecked.extend((write, step) for step in skipped) if problems: @@ -325,14 +325,19 @@ def _build(write: Write) -> Task: return Task(id=write.id, version=None, steps=[build_step(step) for step in write.source.config]) -def _preflight(write: Write, task: Task, unwritten: set[tuple[str, str]]) -> tuple[list[str], list[TaskStep]]: +def _preflight(plan: Plan, write: Write, task: Task, + unwritten: set[tuple[str, str]]) -> tuple[list[str], list[TaskStep]]: """``task``'s preflight problems, and the steps with a preflight left unchecked because they read a (store, id) in ``unwritten``, which a dry run would write first; an artifact and an agent can share an id. A step reading one of the task's own outputs is skipped too, and not listed: the output only exists once the - task runs.""" + task runs. A step the solver deploys or installs at is checked for the agent the run gives it (``Plan.agent``).""" outputs = {output.id for output in write.source.outputs} + solver = plan.agent.steps.get(write.id, ()) if plan.agent is not None else () problems, unchecked = [], [] for config, step in zip(write.source.config, task.steps): + if step.id in solver: + config = {**config, "a2a_agent_id": plan.agent.agent_id, "a2a_agent_version": plan.agent.version} + step = build_step(config) problems.extend(step.config_problems()) reads = _reads(config) if any(id in outputs for _, id in reads): diff --git a/src/agent_env/bundle/plan.py b/src/agent_env/bundle/plan.py index 63c087a9..5067eb6e 100644 --- a/src/agent_env/bundle/plan.py +++ b/src/agent_env/bundle/plan.py @@ -54,7 +54,10 @@ from agent_env.task import Task from agent_env.task_step.registry import get_task_step_registry from agent_env.task_step.task_step import TaskStep +from agent_env.task_step.task_steps.deploy_agent import DeployAgentTaskStep +from agent_env.task_step.task_steps.install_agent import InstallAgentTaskStep from agent_env.task_step.task_steps.load_artifact import LoadArtifactTaskStep +from agent_env.task_step.task_steps.prompt_agent import PromptAgentTaskStep from ._fs import fold, relative, with_article from .authoring import AuthoringContext, entry_files, staged_build_context @@ -127,6 +130,17 @@ def agent_id(self) -> str: return self.alias if self.store_version is not None else self.target_id +@dataclass(frozen=True) +class AgentOverride: + """The agent a run deploys for the solver in place of the one each task names: at each step that deploys or + installs an agent a step prompts, not as its user simulator (``solver_steps``), by task id.""" + + name: str # as the run was given it + agent_id: str + steps: Mapping[str, tuple[str, ...]] + version: int | None = None # the store's version the plan read; None for an agent the run writes + + @dataclass(frozen=True) class Write: """An entity or built image a run writes.""" @@ -150,6 +164,7 @@ class Plan: # run registers in place of its build context alone. published: Mapping[str, Publication] = field(default_factory=dict) aliases: Mapping[str, AliasUse] = field(default_factory=dict) # each agent alias the run names, by alias + agent: AgentOverride | None = None # the agent --agent names for the solver def bundle_of(self, entry: BundleEntry) -> Bundle: """The bundle ``entry`` comes from: the installed one for an entry pulled in from one, else the run's.""" @@ -162,12 +177,13 @@ def path(self, entry: BundleEntry) -> str: def plan_bundle(resolved: ResolvedBundle, *, tasks: Sequence[str] = (), evals: Sequence[str] = (), - installed: Installed | None = None) -> Plan: + installed: Installed | None = None, agent: str | None = None) -> Plan: """Plan a run of ``tasks`` and ``evals``, each a name or id in this bundle. With neither, the run is every eval or, in a bundle without evals, every task. ``installed`` lists the installed bundles whose - entities a reference pulls in (``installed_bundles``); without it, an id under one is a store id. Raises + entities a reference pulls in (``installed_bundles``); without it, an id under one is a store id. ``agent`` names + the agent the solver runs as, as a task's reference would, leaving the tasks as written (``Plan.agent``). Raises BundleError listing every problem found.""" - return _Planner(resolved, installed).plan(tasks, evals) + return _Planner(resolved, installed).plan(tasks, evals, agent) def check_bundle(resolved: ResolvedBundle, *, installed: Installed | None = None) -> None: @@ -206,6 +222,8 @@ def __init__(self, resolved: ResolvedBundle, installed: Installed | None = None) self.pullable: dict[str, _Pullable | str] = {} # each installed bundle read, by id root, or why it can't be self.unpulled: dict[tuple[EntityKind, str, int | None], str] = {} # why an installed bundle's id isn't pulled self.alias_table: InstalledAliases | None = None # the installed aliases, once a plan reads them + self.alias_uses: dict[str, AliasUse] = {} # each alias the plan names, by alias + self.alias_stored: dict[str, int | None] = {} # the store's latest version of each alias's id, as read self.pulling: set[_Key] = set() self.own = _by_id(resolved.bundle) # what a pulled entity's id names in this bundle images: dict[BundleEntry, list[BundleImage]] = {} @@ -265,7 +283,7 @@ def _resolve_aliases(self) -> dict[str, AliasUse]: pointed at the alias's agent is left with the problem the store ref check reports.""" if self.installed is None: return {} - table, uses, stored, default = installed_aliases(self.installed()), {}, {}, _UNREAD + table, uses, stored, default = installed_aliases(self.installed()), self.alias_uses, self.alias_stored, _UNREAD self.alias_table = table for index, entry in enumerate(self.entries): if entry.entry.kind is not BundleKind.TASK: @@ -292,6 +310,54 @@ def _resolve_aliases(self) -> dict[str, AliasUse]: self._replace(index, replace(entry, config=_renamed(entry.config, aliased), references=tuple(references))) return uses + def _agent_override(self, name: str) -> AgentOverride | None: + """The agent ``name`` names, as a task's reference to it would: this bundle's agent, an installed bundle's + pulled in, an alias's, or the store's, read once. Each task with solver steps gets a reference to it, so it's + written and checked with what the task needs, and its config is left as written. Records a problem and gives + None when ``name`` names no agent, or no task has a step for the solver.""" + where = f"--agent {name!r}" + agents = [entry.entry for entry in self.resolved.entries if entry.entry.kind is BundleKind.AGENT] + own = next((agent for agent in agents if _nfc(name) in (agent.name, agent.id)), None) + near = next((agent for agent in agents if fold(name) in (fold(agent.name), fold(agent.id))), None) + if own is None and near is not None: + self.problems.append(f"{where}: isn't in this bundle, but {near.name!r} is; names are case-sensitive") + return None + if own is not None: + ref = Reference(EntityKind.AGENT, own.id, None, own, where, None) + else: + ref = Reference(EntityKind.AGENT, name, None, None, where, None) + use = None if self.alias_table is None else self._alias_use( + self.alias_table, name, None, self.alias_uses, self.alias_stored) + if isinstance(use, str): + self.problems.append(f"{where}: {use}") + return None + ref = self._pulled(replace(ref, id=use.target_id) if use is not None and use.store_version is None else ref) + if (problem := self.unpulled.get((ref.kind, ref.id, None)) or self._unstored(ref)) is not None: + self.problems.append(f"{where}: {problem}") + return None + steps = {} + for index, entry in enumerate(self.entries): + if entry.entry.kind is BundleKind.TASK and (solver := solver_steps(entry.config)): + steps[entry.entry.id] = solver + self._replace(index, replace(entry, references=(*entry.references, ref))) + if not steps: + self.problems.append(f"{where}: no step in this bundle's tasks prompts an agent, so none deploys a solver") + return None + return AgentOverride(name, ref.local.id if ref.local is not None else ref.id, steps) + + def _unstored(self, ref: Reference) -> str | None: + """Why the store has no agent ``ref`` names, a store ref, or None when it has one or ``ref`` is local.""" + if ref.local is not None: + return None + try: + STORE_GETTERS[EntityKind.AGENT](ref.id, None) + except NotFoundError: + hint = self.alias_table.hint(ref.id) if self.alias_table is not None else "" + return f"there is no agent {ref.id!r} in this bundle, an installed bundle or the store{hint}" + except (ValueError, KeyError, TypeError) as e: + return f"agent {ref.id!r} can't be read ({type(e).__name__}: {e})" + return None + def _alias_use(self, table: InstalledAliases, agent_id: str, version: int | None, uses: dict[str, AliasUse], stored: dict[str, int | None]) -> AliasUse | str | None: """What ``agent_id`` names when it's an installed alias, the store's agent of that id read first (its latest @@ -403,13 +469,16 @@ def _pull(self, pullable: _Pullable, target: BundleEntry | BundleImage) -> None: self.pulling.discard(key) self._add(key, source, _needs(source)) - def plan(self, tasks: Sequence[str], evals: Sequence[str]) -> Plan: + def plan(self, tasks: Sequence[str], evals: Sequence[str], agent: str | None = None) -> Plan: aliases = self._resolve_aliases() + override = self._agent_override(agent) if agent is not None else None self._check_writers() self._check_loops() selected_tasks, selected_evals = self._selection(tasks, evals) closure = self._closure(selected_tasks + selected_evals) store_refs, store_latest = self._check_store_refs(closure) + if override is not None: + override = replace(override, version=store_latest.get((EntityKind.AGENT, override.agent_id))) self._build_tasks(closure) self._check_files(closure) self._check_cli_loads(closure) @@ -421,7 +490,7 @@ def plan(self, tasks: Sequence[str], evals: Sequence[str]) -> Plan: raise BundleError(self.problems) writes = self._ordered(closure) return Plan(self.resolved, writes, tuple(selected_tasks), tuple(selected_evals), store_refs, store_latest, - dict(self.origins), self._published(writes), aliases) + dict(self.origins), self._published(writes), aliases, override) def check(self) -> None: self._check_writers() @@ -986,6 +1055,17 @@ def folder_files(bundle: Bundle, entry: BundleEntry) -> Mapping[str, Path]: _UNREAD: Any = object() # the configured default agent, before a task needs it +def solver_steps(steps: Any) -> tuple[str, ...]: + """The ids of the steps in ``steps``, a task's config, that deploy or install an agent a ``prompt_agent`` step + prompts, and not as its user simulator: the solver's.""" + prompts = [step for step in steps if isinstance(step, dict) and step.get("type") == PromptAgentTaskStep.type] + prompted = {step.get("agent_name") or TaskStep.DEFAULT_AGENT_NAME for step in prompts} + prompted -= {step.get("user_agent_name") or PromptAgentTaskStep.DEFAULT_USER_AGENT_NAME for step in prompts} + return tuple(step["id"] for step in steps if isinstance(step, dict) and isinstance(step.get("id"), str) + and step.get("type") in (DeployAgentTaskStep.type, InstallAgentTaskStep.type) + and (step.get("agent_name") or TaskStep.DEFAULT_AGENT_NAME) in prompted) + + def _leaves_an_agent_unset(task: ResolvedEntry) -> bool: """Whether a step of ``task`` leaves an agent field unset, so deploys the configured default agent.""" registry = get_task_step_registry() diff --git a/src/agent_env/bundle/preflight.py b/src/agent_env/bundle/preflight.py index 47d97637..394e0535 100644 --- a/src/agent_env/bundle/preflight.py +++ b/src/agent_env/bundle/preflight.py @@ -77,7 +77,7 @@ from agent_env.utils.docker.build import DEFAULT_BUILD_PLATFORM, docker_plugin_missing, docker_unreachable from .parse import BundleError, BundleKind -from .plan import Plan +from .plan import AgentOverride, Plan from .resolve import BuiltImage, Reference, ReferencedImage, ResolvedEntry, build_step _SHOWN = 3 @@ -147,18 +147,20 @@ def __init__(self, plan: Plan, sandbox: str | None): def task(self, task: ResolvedEntry) -> None: steps = [build_step(config) for config in task.config] sandboxes = {step.sandbox_name: step for step in steps if isinstance(step, DeploySandboxTaskStep)} + solver = self.plan.agent.steps.get(task.entry.id, ()) if self.plan.agent else () for step in steps: where = f"{self.plan.path(task.entry)}: step {step.id!r}" + override = self.plan.agent if step.id in solver else None try: if isinstance(step, DeployEnvTaskStep): self._env(where, step) elif isinstance(step, DeployAgentTaskStep): - self._agent(where, step, sandboxes) + self._agent(where, step, sandboxes, override) elif isinstance(step, DeploySandboxTaskStep): self._sandbox(where, step) elif isinstance(step, RubricsVerifierTaskStep): self._judge(where, step) - elif isinstance(step, InstallAgentTaskStep) and step.a2a_agent_id is None: + elif isinstance(step, InstallAgentTaskStep) and step.a2a_agent_id is None and override is None: self.default_agent_users.append((where, "names no agent")) except ValueError as e: # a provider that names no backend, or can't be built from its config self._problem(where, str(e)) @@ -301,7 +303,9 @@ def _planned_version(self, kind: EntityKind, entity_id: str, version: int | None """The version a write names a store entity at: its pin, else the one the plan read, which the writer pins.""" return version if version is not None else self.plan.store_latest.get((kind, entity_id)) - def _agent(self, where: str, step: DeployAgentTaskStep, sandboxes: dict[str, DeploySandboxTaskStep]) -> None: + def _agent(self, where: str, step: DeployAgentTaskStep, sandboxes: dict[str, DeploySandboxTaskStep], + override: AgentOverride | None = None) -> None: + """Check a deploy of the step's agent, or of ``override``, the agent ``--agent`` names for the solver.""" accepts: Callable[[SandboxProvider], Accepts] = A2AAgent.accepts if step.sandbox_name: linked = sandboxes.get(step.sandbox_name) @@ -315,8 +319,9 @@ def _agent(self, where: str, step: DeployAgentTaskStep, sandboxes: dict[str, Dep else: provider = _provider(self.sandbox or step.sandbox_type, get_agent_sandbox_provider) given = frozenset({*step.env_vars, *([ENV_BASE_URL] if step.litellm_base_url else [])}) - self._agent_deploy(where, provider, step.a2a_agent_id, step.a2a_agent_version, "names no agent", accepts, - given=given) + agent_id, version = ((override.agent_id, override.version) if override is not None + else (step.a2a_agent_id, step.a2a_agent_version)) + self._agent_deploy(where, provider, agent_id, version, "names no agent", accepts, given=given) def _judge(self, where: str, step: RubricsVerifierTaskStep) -> None: if not step.use_agent_judge or step.agent_name is not None: diff --git a/src/agent_env/bundle/run.py b/src/agent_env/bundle/run.py index 6b1f9d43..b05bd02c 100644 --- a/src/agent_env/bundle/run.py +++ b/src/agent_env/bundle/run.py @@ -18,7 +18,7 @@ import shlex import time import uuid -from collections.abc import Callable, Sequence +from collections.abc import Callable, Mapping, Sequence from dataclasses import dataclass, replace from enum import Enum from pathlib import Path @@ -196,6 +196,7 @@ def run_bundle( model: str | None = None, user_model: str | None = None, judge_model: str | None = None, + agent: str | None = None, sandbox: str | None = None, on_progress: Callable[[str], None] | None = None, id_root: str | None = None, @@ -203,7 +204,8 @@ def run_bundle( ) -> BundleRun: """Write what the tasks and evals selected in ``root`` need, as ``plan_bundle`` selects them, then run each of their tasks once. ``model``, ``user_model`` and ``judge_model`` replace the models the solver, the - user simulator and the judges run on, leaving the tasks as written; ``sandbox`` is the provider that envs, + user simulator and the judges run on, and ``agent`` the agent the solver runs as (``Plan.agent``), leaving the + tasks as written; ``sandbox`` is the provider that envs, agents, sandboxes and the judge deploy on. ``on_progress`` is given a line of text at each write and step. ``id_root`` roots the ids of a bundle an installed package provides (``InstalledBundle.id_root``); a folder's ids otherwise come from its path. @@ -228,9 +230,8 @@ def run_bundle( build_sandbox_provider(sandbox) say = _progress(on_progress) with namespace_routing(): - plan, preflight = _planned(root, tasks, evals, id_root, sandbox) - for use in plan.aliases.values(): - say(_aliased(use)) + plan, preflight = _planned(root, tasks, evals, id_root, sandbox, agent) + _say_agents(plan, say) materialization = materialize( plan, on_wait=lambda: say("waiting for another agent-env run to finish writing this bundle's ids"), @@ -244,7 +245,8 @@ def run_bundle( ensure_default_envs(preflight.infra_kinds, say=say) with Interrupts() as interrupts: overrides = _overrides(sandbox, user_model, judge_model) - runs = interrupts.run(_run_all(plan, to_run, model, overrides, keep, say, interrupts)) + runs = interrupts.run(_run_all(plan, to_run, model, overrides, _solver_params(plan, materialization), keep, + say, interrupts)) _mark_cancelled(runs, interrupts.reason) result = _bundle_run(plan, materialization, runs, _skipped(plan, entries, every=not tasks and not evals)) if not interrupts.count: @@ -263,6 +265,7 @@ def dry_run_bundle( tasks: Sequence[str] = (), evals: Sequence[str] = (), sandbox: str | None = None, + agent: str | None = None, on_progress: Callable[[str], None] | None = None, id_root: str | None = None, ) -> DryRun: @@ -278,9 +281,8 @@ def dry_run_bundle( build_sandbox_provider(sandbox) say = _progress(on_progress) with namespace_routing(): - plan, preflight = _planned(root, tasks, evals, id_root, sandbox) - for use in plan.aliases.values(): - say(_aliased(use)) + plan, preflight = _planned(root, tasks, evals, id_root, sandbox, agent) + _say_agents(plan, say) materialization = materialize(plan, dry_run=True, on_write=lambda done: say(_written(plan, done))) runs = _to_run(plan) return DryRun(materialization, runs, _skipped(plan, runs, every=not tasks and not evals), preflight.infra, @@ -309,9 +311,9 @@ def _bundle_run(plan: Plan, materialization: Materialization, runs: tuple[TaskRu def _planned(root: Path | str, tasks: Sequence[str], evals: Sequence[str], id_root: str | None, - sandbox: str | None) -> tuple[Plan, Preflight]: + sandbox: str | None, agent: str | None = None) -> tuple[Plan, Preflight]: plan = plan_bundle(resolve_bundle(parse_bundle(Path(root), id_root=id_root)), tasks=tasks, evals=evals, - installed=installed_bundles) + installed=installed_bundles, agent=agent) refuse_store_tasks(plan) to_run = set(_to_run(plan)) return plan, preflight_run(plan, [write.source for write in plan.writes @@ -345,6 +347,28 @@ def refuse_store_tasks(plan: Plan) -> None: raise BundleError(problems) +def _say_agents(plan: Plan, say: Callable[[str], None]) -> None: + """Say what each agent alias names, and what the solver deploys in place of the tasks' agents.""" + for use in plan.aliases.values(): + say(_aliased(use)) + if plan.agent is not None: + runs = _to_run(plan) + solved = sum(entry.id in plan.agent.steps for entry in runs) + say(f"--agent {plan.agent.name}: the solver deploys {plan.agent.agent_id}, in {solved} of {len(runs)} " + f"task{'s' if len(runs) != 1 else ''}") + + +def _solver_params(plan: Plan, materialization: Materialization) -> dict[str, dict[str, dict]]: + """Each task's ``step_params`` that deploy ``Plan.agent`` at its solver steps: at the version this run wrote, for + an agent it writes, or the store's version the plan read.""" + if plan.agent is None: + return {} + written = any(write.kind is BundleKind.AGENT and write.id == plan.agent.agent_id for write in plan.writes) + version = materialization.version_of(BundleKind.AGENT.store, plan.agent.agent_id) if written else plan.agent.version + choice = {"a2a_agent_id": plan.agent.agent_id, "a2a_agent_version": version} + return {task_id: {step_id: dict(choice) for step_id in steps} for task_id, steps in plan.agent.steps.items()} + + def _overrides(sandbox: str | None, user_model: str | None, judge_model: str | None) -> dict[str, str]: """The run's ``user_overrides``: where everything deploys, and the models the user simulator and judges run on.""" overrides = {"env_sandbox": sandbox, "agent_sandbox": sandbox, "sandbox": sandbox} if sandbox else {} @@ -353,8 +377,8 @@ def _overrides(sandbox: str | None, user_model: str | None, judge_model: str | N async def _run_all( - plan: Plan, to_run: list[tuple[BundleEntry, Task]], model: str | None, overrides: dict[str, str], keep: bool, - say: Callable[[str], None], interrupts: Interrupts, + plan: Plan, to_run: list[tuple[BundleEntry, Task]], model: str | None, overrides: dict[str, str], + step_params: Mapping[str, dict], keep: bool, say: Callable[[str], None], interrupts: Interrupts, ) -> tuple[TaskRun, ...]: semaphore = asyncio.Semaphore(MAX_CONCURRENT_RUNS) run_group_id = uuid.uuid4().hex @@ -374,8 +398,9 @@ async def settle(context: TaskStepContext, cancelled: bool) -> TeardownReport | async def run_one(entry: BundleEntry, task: Task) -> TaskRun: tag = f"[{plan.path(entry)}]" metadata = {"run_group_id": run_group_id, "task_id": task.id} - if overrides: - metadata["user_overrides"] = dict(overrides) + user_overrides = {**overrides, **({"step_params": step_params[entry.id]} if entry.id in step_params else {})} + if user_overrides: + metadata["user_overrides"] = user_overrides context = TaskStepContext(metadata=metadata) try: await semaphore.acquire() diff --git a/src/agent_env/cli/run.py b/src/agent_env/cli/run.py index 6c22099c..d0138616 100644 --- a/src/agent_env/cli/run.py +++ b/src/agent_env/cli/run.py @@ -35,6 +35,9 @@ @click.option("--model", default=None, help="The model the solver agent runs on, whatever the tasks name.") @click.option("--user-model", default=None, help="The model the user simulator runs on, whatever the tasks name.") @click.option("--judge-model", default=None, help="The model the judges grade with, whatever the tasks name.") +@click.option("--agent", default=None, + help="The agent the solver runs as, whatever the tasks name: an alias such as pi, an agent of this " + "bundle or an installed one, or the store's.") @click.option("--sandbox", default=None, help="The sandbox provider envs, agents, sandboxes and the judge deploy on; comma-separated for a " "fallback chain.") @@ -44,7 +47,8 @@ help="Check the bundle as the run does and show what it would write and run; write and run nothing.") @click.pass_context def run(ctx: click.Context, bundle: str | None, tasks: tuple[str, ...], evals: tuple[str, ...], model: str | None, - user_model: str | None, judge_model: str | None, sandbox: str | None, keep: bool, dry_run: bool): + user_model: str | None, judge_model: str | None, agent: str | None, sandbox: str | None, keep: bool, + dry_run: bool): """Run a bundle's tasks and evals, or list the installed bundles. BUNDLE is a folder, or the name of a bundle an installed package provides (PACKAGE/NAME when several @@ -53,11 +57,14 @@ def run(ctx: click.Context, bundle: str | None, tasks: tuple[str, ...], evals: t With no --task or --eval, it runs every eval or, in a bundle without evals, every task. Each task runs once, at most four at a time, and the run exits 1 if any of them failed. With --verbose, a run that raised prints its traceback. --model, --user-model and --judge-model choose the models for this run only, - leaving the tasks as written, and the summary names the model each role ran on. + leaving the tasks as written, and the summary names the model each role ran on. --agent likewise chooses the + agent the solver runs as: each step that deploys or installs the agent a task prompts deploys it instead, + while its user simulator and judges keep theirs. --dry-run makes the checks the run makes before its first task, reading the stores as the run does, and prints what it would write and run, writing and running nothing. It exits 1 on a problem the run would - stop at, and 0 otherwise; the models and --keep change nothing it shows. It takes no lock, so another run + stop at, and 0 otherwise; the models and --keep change nothing it shows, and --agent what it writes and checks + for the solver. It takes no lock, so another run can write first and change what it shows. Each run's sandboxes are torn down as it ends; its instance and outputs stay. --keep holds them up @@ -65,9 +72,9 @@ def run(ctx: click.Context, bundle: str | None, tasks: tuple[str, ...], evals: t exits 130 (143 for SIGTERM); a second one stops the teardown and prints what is still up. """ if bundle is None: - if tasks or evals or model or user_model or judge_model or sandbox or keep or dry_run: - raise click.UsageError("--task, --eval, --model, --user-model, --judge-model, --sandbox, --keep and " - "--dry-run need a BUNDLE to run") + if tasks or evals or model or user_model or judge_model or agent or sandbox or keep or dry_run: + raise click.UsageError("--task, --eval, --model, --user-model, --judge-model, --agent, --sandbox, --keep " + "and --dry-run need a BUNDLE to run") _list_installed() return root, id_root = _locate(bundle) @@ -75,12 +82,12 @@ def run(ctx: click.Context, bundle: str | None, tasks: tuple[str, ...], evals: t try: if dry_run: click.echo(_DRY_RUN) - _report_dry_run(dry_run_bundle(root, tasks=tasks, evals=evals, sandbox=sandbox, on_progress=click.echo, - id_root=id_root)) + _report_dry_run(dry_run_bundle(root, tasks=tasks, evals=evals, sandbox=sandbox, agent=agent, + on_progress=click.echo, id_root=id_root)) return result = run_bundle(root, tasks=tasks, evals=evals, model=model, user_model=user_model, - judge_model=judge_model, sandbox=sandbox, on_progress=click.echo, id_root=id_root, - keep=keep) + judge_model=judge_model, agent=agent, sandbox=sandbox, on_progress=click.echo, + id_root=id_root, keep=keep) except BundleError as e: _note_an_installed_namesake(e, bundle, root, id_root) raise diff --git a/src/agent_env/task_step/task_steps/install_agent.py b/src/agent_env/task_step/task_steps/install_agent.py index 40df1eb2..1da700a2 100644 --- a/src/agent_env/task_step/task_steps/install_agent.py +++ b/src/agent_env/task_step/task_steps/install_agent.py @@ -177,6 +177,18 @@ def preflight(self) -> list[str]: problem = _install_problem(agent) return [f"install_agent '{self.id}': {problem}"] if problem else [] + def _agent_ref(self, context: TaskStepContext) -> tuple[str, Optional[int]]: + """The agent id and version this run installs: this step's ``step_params`` override names both, else the + run-wide ``a2a_agent_id`` override or the step's agent, at the step's version.""" + step_overrides = self.step_param_overrides(context) + if "a2a_agent_id" in step_overrides: + return step_overrides["a2a_agent_id"], step_overrides.get("a2a_agent_version") + # Honor `--a2a-agent-id` override (same pattern as DeployAgentTaskStep) — lets the + # caller swap which agent gets installed without rewriting the task definition. + user_overrides = context.metadata.get("user_overrides", {}) + a2a_agent_id = user_overrides.get("a2a_agent_id") or self.a2a_agent_id or get_config().get_default_a2a_agent_id() + return a2a_agent_id, self.a2a_agent_version + async def execute(self, context: TaskStepContext) -> TaskStepContext: from agent_env.providers.sandbox_providers.sandbox_provider import ( SANDBOX_MODE_VM, @@ -218,11 +230,8 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: ) sandbox = await provider.get_sandbox(ds.sandbox_id) - # Honor `--a2a-agent-id` override (same pattern as DeployAgentTaskStep) — lets the - # caller swap which agent gets installed without rewriting the task definition. user_overrides = context.metadata.get("user_overrides", {}) - a2a_agent_id = user_overrides.get("a2a_agent_id") or self.a2a_agent_id or get_config().get_default_a2a_agent_id() - agent = A2AAgent.get(a2a_agent_id, self.a2a_agent_version) + agent = A2AAgent.get(*self._agent_ref(context)) if problem := _install_problem(agent): raise RuntimeError(problem) diff --git a/src/agent_env/task_step/task_steps/prompt_agent.py b/src/agent_env/task_step/task_steps/prompt_agent.py index 9d732679..069a06d5 100644 --- a/src/agent_env/task_step/task_steps/prompt_agent.py +++ b/src/agent_env/task_step/task_steps/prompt_agent.py @@ -117,6 +117,7 @@ def _file_uris(parts: list[dict]) -> list[str]: class PromptAgentTaskStep(TaskStep): type: ClassVar[str] = "prompt_agent" + DEFAULT_USER_AGENT_NAME: ClassVar[str] = "human_agent" entity_refs = (EntityRef.env("snapshot_config.env_id"),) # user_a2a_url stays the task's: setting it is also what tells a human counterpart from a simulator. field_scopes = {**TaskStep.field_scopes, **dict.fromkeys( @@ -147,7 +148,7 @@ def __init__( poll_interval_seconds: int = DEFAULT_POLL_INTERVAL_SECONDS, context_id: Optional[str] = None, max_conversation_turns: int = 1, - user_agent_name: str = "human_agent", + user_agent_name: str = DEFAULT_USER_AGENT_NAME, user_a2a_url: Optional[str] = None, user_agent_timeout_seconds: int = 600, user_output_format: Optional[dict[str, Any]] = None, @@ -301,7 +302,7 @@ def from_dict(cls, data: dict) -> PromptAgentTaskStep: poll_interval_seconds=data.get("poll_interval_seconds", cls.DEFAULT_POLL_INTERVAL_SECONDS), context_id=data.get("context_id"), max_conversation_turns=data.get("max_conversation_turns", 1), - user_agent_name=data.get("user_agent_name", "human_agent"), + user_agent_name=data.get("user_agent_name", cls.DEFAULT_USER_AGENT_NAME), user_a2a_url=data.get("user_a2a_url"), user_agent_timeout_seconds=data.get("user_agent_timeout_seconds", 600), user_output_format=data.get("user_output_format"), diff --git a/tst/integration/cli/run_bundle_built_agents_test.py b/tst/integration/cli/run_bundle_built_agents_test.py index 7fb2050f..7c12fd0a 100644 --- a/tst/integration/cli/run_bundle_built_agents_test.py +++ b/tst/integration/cli/run_bundle_built_agents_test.py @@ -4,6 +4,7 @@ import json import logging +import re import shutil import pytest @@ -13,6 +14,8 @@ from agent_env.cli import cli from agent_env.config import configure, reset_config from agent_env.store.object_store.local.grant_server import grant_server +from agent_env.store.routing import namespace_routing +from agent_env.task.store import get_task_instance_store from tst.util.a2a_test_agent import AGENT_DIR, PROTOCOL_DIR pytestmark = [pytest.mark.integration, pytest.mark.int_test_slow] @@ -146,8 +149,15 @@ def state(monkeypatch, tmp_path): reset_config() -def _run(root): - return CliRunner().invoke(cli, ["run", str(root)]) +def _run(root, *args): + return CliRunner().invoke(cli, ["run", str(root), *args]) + + +def _reply(output): + """What the agent of the run ``output`` reports replied, read from its instance.""" + instance_id = re.search(r"instance (@local/\S+)", output).group(1) + with namespace_routing(): + return get_task_instance_store().get(instance_id).context["prompt_responses"][0]["response"] def test_each_shape_of_agent_dockerfile_builds_runs_and_is_rebuilt_only_when_its_folder_changes(state): @@ -202,3 +212,24 @@ def test_a_dockerfile_that_fails_to_build_is_one_problem_with_dockers_output_and assert "about to fail" in result.output, result.output assert "Traceback" not in result.output, result.output assert "agents/broken: v1" not in result.output, result.output + + +def test_agent_runs_the_solver_as_another_agent_for_one_run_and_leaves_the_task_as_written(state): + root = state / "swap-bundle" + for name in ("plain", "subdir"): + _agent_folder(root, name, AGENTS[name]) + (root / "tasks").mkdir() + (root / "tasks/ask.json").write_text(_task("plain", "sees []")) + + swapped = _run(root, "--agent", "subdir") + + assert swapped.exit_code == 0, swapped.output + assert re.search(r"--agent subdir: the solver deploys @local/\S+/subdir, in 1 of 1 task", swapped.output), swapped.output + assert "tasks/ask.json v1: passed" in swapped.output, swapped.output + assert _reply(swapped.output) == REPLIES["subdir"] + + written = _run(root) + + assert written.exit_code == 0, written.output + assert "tasks/ask.json: v1, unchanged" in written.output, written.output + assert _reply(written.output) == REPLIES["plain"] diff --git a/tst/unit/agents/claude_code/test_claude_bridge_cancel.py b/tst/unit/agents/claude_code/test_claude_bridge_cancel.py index f8c80862..a4ff2353 100644 --- a/tst/unit/agents/claude_code/test_claude_bridge_cancel.py +++ b/tst/unit/agents/claude_code/test_claude_bridge_cancel.py @@ -23,7 +23,7 @@ def _state(pid: int) -> str: try: with open(f"/proc/{pid}/stat") as stat: return stat.read().rsplit(")", 1)[1].split()[0] - except FileNotFoundError: + except (FileNotFoundError, ProcessLookupError): # a process that exits mid-read gives ESRCH return "" return subprocess.run(["ps", "-o", "stat=", "-p", str(pid)], capture_output=True, text=True).stdout.strip()[:1] diff --git a/tst/unit/bundle/agent_override_test.py b/tst/unit/bundle/agent_override_test.py new file mode 100644 index 00000000..33c5401e --- /dev/null +++ b/tst/unit/bundle/agent_override_test.py @@ -0,0 +1,223 @@ +"""``--agent``: the agent a run deploys for the solver, at each step deploying or installing the agent a step prompts, +named as a task's reference would name it, with the tasks left as written.""" + +import json + +import pytest +from click.testing import CliRunner + +from agent_env.a2a_agent import A2AAgent +from agent_env.artifact.artifacts.docker_image import DockerImageArtifact +from agent_env.artifact.store import get_artifact_store +from agent_env.bundle import BundleError, dry_run_bundle, parse_bundle, run_bundle +from agent_env.bundle import preflight as preflight_module +from agent_env.bundle.installed import InstalledBundle +from agent_env.bundle.plan import AgentOverride, plan_bundle, solver_steps +from agent_env.bundle.resolve import resolve_bundle +from agent_env.cli import cli +from agent_env.plugins._bundles import ALIASES_FILE +from agent_env.store.routing import namespace_routing +from agent_env.task.task import Task +from agent_env.task_step.context import TaskStepContext +from agent_env.task_step.task_steps.install_agent import InstallAgentTaskStep +from tst.unit.bundle._support import layout + +ROOT = "@local/~/triage" +SOLVER = "@local/demo-dist/kit/solver" +DEPLOY = {"id": "agent", "type": "deploy_agent", "env_ids": [], "a2a_agent_id": "theirs"} +PROMPT = {"id": "ask", "type": "prompt_agent", "prompt_id": "ask", "prompt": "hi", "depends_on": ["agent"]} +INSTALL = [{"id": "box", "type": "deploy_sandbox", "sandbox_name": "box", "sandbox_mode": "vm", "sandbox_type": "local"}, + {"id": "agent", "type": "install_agent", "sandbox_name": "box", "agent_name": "coder", "a2a_agent_id": "theirs", + "depends_on": ["box"]}, + {**PROMPT, "agent_name": "coder"}] + + +@pytest.fixture +def bundle_dir(tmp_path, monkeypatch, local_stores, cli_routing): + monkeypatch.setenv("HOME", str(tmp_path)) + _store_agent("theirs") + return tmp_path / "triage" + + +@pytest.fixture +def kit(tmp_path): + """An installed bundle whose agent ``solver`` has the alias ``solver-x``.""" + root = layout(tmp_path / "site" / "kit", {"agents/solver/Dockerfile": "FROM scratch\n", + ALIASES_FILE: '[agents]\nsolver = ["solver-x"]\n'}) + return InstalledBundle("kit", "demo-dist", "1.0", "demo", root, None, None) + + +def _store_agent(agent_id): + with namespace_routing(): + image = get_artifact_store().put_document(DockerImageArtifact( + id=f"{agent_id}-image", description="d", image_name="ghcr.io/acme/agent:1", tar_gz_s3_url="s3://b/x.tar.gz")) + return A2AAgent.put(id=agent_id, docker_image_artifact=image) + + +def _bundle(bundle_dir, steps, files=None): + layout(bundle_dir, {"tasks/t.json": json.dumps(steps), **(files or {})}) + return bundle_dir + + +def _plan(bundle_dir, agent, *installed): + return plan_bundle(resolve_bundle(parse_bundle(bundle_dir)), installed=lambda: installed, agent=agent) + + +def _problems(bundle_dir, agent, *installed): + with pytest.raises(BundleError) as caught: + _plan(bundle_dir, agent, *installed) + return caught.value.problems + + +def _step(plan, step_id): + task = next(write for write in plan.writes if write.id == f"{ROOT}/t") + return next(step for step in task.source.config if step["id"] == step_id) + + +def test_an_alias_names_the_agent_the_solver_deploys_and_the_task_stays_as_written(bundle_dir, kit): + plan = _plan(_bundle(bundle_dir, [DEPLOY, PROMPT]), "solver-x", kit) + + assert plan.agent == AgentOverride("solver-x", SOLVER, {f"{ROOT}/t": ("agent",)}) + assert SOLVER in {write.id for write in plan.writes} + assert _step(plan, "agent")["a2a_agent_id"] == "theirs" and plan.aliases["solver-x"].target_id == SOLVER + + +def test_this_bundles_own_agent_is_named_by_its_folder(bundle_dir): + plan = _plan(_bundle(bundle_dir, [DEPLOY, PROMPT], {"agents/helper/Dockerfile": "FROM scratch\n"}), "helper") + + assert plan.agent.agent_id == f"{ROOT}/helper" and f"{ROOT}/helper" in {write.id for write in plan.writes} + + +def test_the_stores_agent_is_named_by_its_id_at_the_version_the_plan_read(bundle_dir): + _store_agent("mine") + _store_agent("mine") + + plan = _plan(_bundle(bundle_dir, [DEPLOY, PROMPT]), "mine") + + assert (plan.agent.agent_id, plan.agent.version) == ("mine", 2) + assert "mine" not in {write.id for write in plan.writes} + + +def test_a_name_that_names_no_agent_is_refused_pointing_at_a_close_alias(bundle_dir, kit): + assert _problems(_bundle(bundle_dir, [DEPLOY, PROMPT]), "solver_x", kit) == ( + "--agent 'solver_x': there is no agent 'solver_x' in this bundle, an installed bundle or the store; did you " + "mean the alias 'solver-x'? agent-env a2a-agent list lists the agents a task can name",) + + +def test_this_bundles_agent_is_named_with_its_folders_case_as_a_task_names_it(bundle_dir): + assert _problems(_bundle(bundle_dir, [DEPLOY, PROMPT], {"agents/Helper/Dockerfile": "FROM scratch\n"}), "helper") == ( + "--agent 'helper': isn't in this bundle, but 'Helper' is; names are case-sensitive",) + + +def test_a_bundle_whose_tasks_prompt_no_agent_is_refused(bundle_dir): + assert _problems(_bundle(bundle_dir, [DEPLOY]), "theirs") == ( + "--agent 'theirs': no step in this bundle's tasks prompts an agent, so none deploys a solver",) + + +def test_only_the_steps_deploying_or_installing_the_agent_a_step_prompts_are_the_solvers(): + steps = [ + {"id": "solver", "type": "deploy_agent"}, + {"id": "user", "type": "deploy_agent", "agent_name": "carol"}, + {"id": "judge", "type": "deploy_agent", "agent_name": "grader"}, + {"id": "installed", "type": "install_agent", "agent_name": "coder", "sandbox_name": "box"}, + {"id": "ask", "type": "prompt_agent", "user_agent_name": "carol"}, + {"id": "ask-coder", "type": "prompt_agent", "agent_name": "coder"}, + {"id": "grade", "type": "rubrics_verifier", "agent_name": "grader"}, + ] + + assert solver_steps(steps) == ("solver", "installed") + + +def test_preflight_checks_the_agent_the_solver_deploys(bundle_dir, monkeypatch): + checked = [] + original = preflight_module._Walk._agent_deploy + + def record(self, where, provider, agent_id, version, *args, **kwargs): + checked.append((agent_id, version)) + return original(self, where, provider, agent_id, version, *args, **kwargs) + + monkeypatch.setattr(preflight_module._Walk, "_agent_deploy", record) + _bundle(bundle_dir, [{**DEPLOY, "a2a_agent_version": 1, "sandbox_type": "local"}, PROMPT], + {"agents/helper/Dockerfile": "FROM scratch\n"}) + + dry_run_bundle(bundle_dir, agent="helper") + _store_agent("mine") + dry_run_bundle(bundle_dir, agent="mine") + + assert checked == [(f"{ROOT}/helper", None), ("mine", 1)] + + +def test_an_install_is_checked_for_the_agent_the_solver_installs_not_the_one_the_task_names(bundle_dir): + _store_agent("mine") + _bundle(bundle_dir, INSTALL, {"agents/helper/Dockerfile": "FROM scratch\n"}) + + with pytest.raises(BundleError, match="install_agent 'agent': Agent theirs v1 has no build_context_object_url"): + dry_run_bundle(bundle_dir) + dry_run_bundle(bundle_dir, agent="helper") + with pytest.raises(BundleError, match="install_agent 'agent': Agent mine v1 has no build_context_object_url"): + dry_run_bundle(bundle_dir, agent="mine") + + +def test_a_dry_run_says_what_the_solver_deploys(bundle_dir, monkeypatch): + lines = [] + _store_agent("mine") + + dry_run_bundle(_bundle(bundle_dir, [{**DEPLOY, "sandbox_type": "local"}, PROMPT]), agent="mine", + on_progress=lines.append) + + assert "--agent mine: the solver deploys mine, in 1 of 1 task" in lines + + +def test_a_run_deploys_the_agent_at_the_solvers_steps_alone_at_the_version_it_wrote(bundle_dir, builds, monkeypatch): + seen = {} + + async def run(self, *, context, **kwargs): + seen[self.id] = context.metadata.get("user_overrides") + return context + + monkeypatch.setattr(Task, "run", run) + user = {"id": "user", "type": "deploy_agent", "env_ids": [], "agent_name": "carol", "a2a_agent_id": "theirs"} + _bundle(bundle_dir, [DEPLOY, user, {**PROMPT, "user_agent_name": "carol", "depends_on": ["agent", "user"]}], + {"agents/helper/Dockerfile": "FROM scratch\n"}) + + run_bundle(bundle_dir, agent="helper", sandbox="local") + + assert seen[f"{ROOT}/t"]["step_params"] == {"agent": {"a2a_agent_id": f"{ROOT}/helper", "a2a_agent_version": 1}} + + +def test_a_run_deploys_the_stores_agent_at_the_version_the_plan_read_not_one_put_since(bundle_dir, monkeypatch): + seen = {} + + async def run(self, *, context, **kwargs): + _store_agent("mine") + seen[self.id] = context.metadata.get("user_overrides") + return context + + monkeypatch.setattr(Task, "run", run) + _store_agent("mine") + _bundle(bundle_dir, [DEPLOY, PROMPT]) + + run_bundle(bundle_dir, agent="mine", sandbox="local") + + assert seen[f"{ROOT}/t"]["step_params"] == {"agent": {"a2a_agent_id": "mine", "a2a_agent_version": 1}} + + +def test_an_install_takes_its_steps_override_over_the_run_wide_one_and_the_stored_agent(): + step = InstallAgentTaskStep(id="agent", version=None, sandbox_name="box", a2a_agent_id="stored", a2a_agent_version=3) + solver = {"step_params": {"agent": {"a2a_agent_id": "solver", "a2a_agent_version": None}}} + + assert step._agent_ref(TaskStepContext(metadata={"user_overrides": {"a2a_agent_id": "run-wide", **solver}})) == ( + "solver", None) + assert step._agent_ref(TaskStepContext(metadata={"user_overrides": {"a2a_agent_id": "run-wide"}})) == ("run-wide", 3) + assert step._agent_ref(TaskStepContext()) == ("stored", 3) + + +def test_the_cli_takes_the_agent_for_a_run_and_a_dry_run(bundle_dir, monkeypatch): + monkeypatch.setenv("LITELLM_BASE_URL", "http://litellm") + _store_agent("mine") + _bundle(bundle_dir, [{**DEPLOY, "sandbox_type": "local"}, PROMPT]) + + result = CliRunner().invoke(cli, ["run", str(bundle_dir), "--agent", "mine", "--dry-run"]) + + assert result.exit_code == 0, result.output + assert "--agent mine: the solver deploys mine, in 1 of 1 task" in result.output diff --git a/tst/unit/bundle/installed_test.py b/tst/unit/bundle/installed_test.py index 053bbe65..cc2ceca0 100644 --- a/tst/unit/bundle/installed_test.py +++ b/tst/unit/bundle/installed_test.py @@ -277,11 +277,12 @@ def test_the_cli_says_so_when_no_bundles_are_installed_and_needs_a_bundle_for_it _registered(monkeypatch) assert CliRunner().invoke(cli, ["run"]).output == "No bundles are installed. agent-env run PATH runs a folder.\n" - for option in (["--task", "t"], ["--keep"], ["--dry-run"], ["--user-model", "m"], ["--judge-model", "m"]): + for option in (["--task", "t"], ["--keep"], ["--dry-run"], ["--user-model", "m"], ["--judge-model", "m"], + ["--agent", "pi"]): usage = CliRunner().invoke(cli, ["run", *option]) assert usage.exit_code == 2 - assert ("--task, --eval, --model, --user-model, --judge-model, --sandbox, --keep and --dry-run need a BUNDLE " - "to run") in " ".join(usage.output.split()) + assert ("--task, --eval, --model, --user-model, --judge-model, --agent, --sandbox, --keep and --dry-run need a " + "BUNDLE to run") in " ".join(usage.output.split()) def test_the_cli_runs_an_installed_bundle_by_name_with_ids_rooted_at_its_package(site, monkeypatch, quiet_logs, diff --git a/tst/unit/providers/sandbox_providers/local_sandbox_test.py b/tst/unit/providers/sandbox_providers/local_sandbox_test.py index 0d55287b..6cebea5d 100644 --- a/tst/unit/providers/sandbox_providers/local_sandbox_test.py +++ b/tst/unit/providers/sandbox_providers/local_sandbox_test.py @@ -943,7 +943,7 @@ def _running(pid: int) -> bool: """Whether ``pid`` is still running: a killed process another parent hasn't reaped yet is a zombie.""" try: return Path(f"/proc/{pid}/stat").read_text().rsplit(")", 1)[1].split()[0] not in "ZX" - except FileNotFoundError: + except (FileNotFoundError, ProcessLookupError): # a process that exits mid-read gives ESRCH return False