diff --git a/AGENTS.md b/AGENTS.md index be1458e6..39f6fdd3 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -92,7 +92,7 @@ and deserialized through a registry. |---|---| | `env/` | `Env` base class: a subclass declares `type`, implements `from_dict` and `async deploy(...) -> DeployedEnv`, and may override `reset` and `load_file_artifact_universe`. Built-ins in `env/envs/`: `MCPServerEnv` (`mcp_server`), `MultiEnv` (`multi`, MCP servers plus websites), `ServiceDBEnv` (`service_db`, PostgreSQL), `WebsiteEnv` (`website`), `GatewayEnv` (`gateway_server`). `env/gateway/gateway.py` is the env-side gateway server, container-only code. | | `artifact/` | Immutable versioned artifacts; every `put()` creates a new version. The document goes to the document store, binary payloads to the object store. Built-ins: `CliArtifact` (`cli`), `FileArtifact` (`file`), `FileArtifactUniverse` (`file_artifact_universe`), `DockerImageArtifact` (`docker_image`), `EnvironmentArtifact` (`environment`; legacy `service`), `EnvironmentUniverseArtifact` (`environment_universe`; legacy `service_universe`), `SkillArtifact` (`skill`). A `DockerImageArtifact` holds a tar.gz of its image, or, written by `put_ref` (`--image-ref` on `env mcp-server put` and `a2a-agent put`), only a registry reference with its tag pinned to a digest, which sandboxes pull, or, written by `put_context`, only its build context, which a VM sandbox builds before running it (`VmSandbox.load_docker_images`) and Modal builds before creating its container (`ModalSandboxProvider.prepare_image`, outside a fallback chain's deadline); anything else that runs images by name (a container on the local provider, a lone server on any provider but Modal) refuses such an image. A build context (`utils/build_context.py`) is what `docker build` would send, `.dockerignore` applied, written so equal files give equal bytes and an equal `source_digest`. | -| `task/`, `task_step/` | A `Task` holds its `TaskStep`s inline; `Task.run()` executes them as a DAG. `depends_on` (None means all prior steps) gates a step, independent steps run concurrently, `fail_task_on_error` makes a failure fatal or tolerated, `retry_config` rolls a failed span back through the step journal and re-dispatches it. Built-in steps live in `task_step/task_steps/` (`deploy_env`, `deploy_agent`, `prompt_agent`, the verifiers under `verifiers/`, and more). | +| `task/`, `task_step/` | A `Task` holds its `TaskStep`s inline; `Task.run()` executes them as a DAG. `depends_on` (None means all prior steps) gates a step, independent steps run concurrently, `fail_task_on_error` makes a failure fatal or tolerated, `retry_config` rolls a failed span back through the step journal and re-dispatches it. Built-in steps live in `task_step/task_steps/` (`deploy_env`, `deploy_agent`, `prompt_agent`, the verifiers under `verifiers/`, and more). `task/episode/` runs a task with some of its agents played from outside, as a training loop does (`EpisodeSpec`). | | `store/` | Four store ABCs with local and cloud implementations: `DocumentStore` (SQLite, MongoDB), `ObjectStore` (filesystem, S3, Cloud Storage), `ImageStore` (local OCI registry, ECR), `SecretStore` (env vars or file, AWS Secrets Manager, Google Cloud Secret Manager). `VersionedEntityStore` implements the shared versioned get/put logic, `QueryBuilder` is the immutable chained query API, `store/base.py` holds the error types. A new backend must pass the conformance kits in `tst/store/`. | | `config/` | The `Config` singleton (`get_config`, `configure`, `reset_config`) in `config/runtime.py`, file discovery in `config/loader.py`, and `load_impl`, which resolves `module:Class` pointers. `agent_env.store` re-exports the config names for compatibility. | | `providers/` | `providers/sandbox_providers/` holds the sandbox providers `local`, `modal`, `modal_vm`, `e2b`, `sail_vm` (the `sail` extra); `[sandbox] default` and `agent_default` accept a comma-separated fallback chain, whose deploys skip a provider that can't run them. A sandbox provider declares what it runs (`ON_THIS_MACHINE`, `CONTAINER_ACCEPTS`, `PRIVATE_NETWORK`; `SUPPORTED_SANDBOX_MODES`, derived from whether it implements `create_sandbox` and `create_vm`; and `SANDBOX_ACCEPTS`, the images an agent's sandbox runs, which defaults to those a VM loads when it implements `create_vm` and to `CONTAINER_ACCEPTS` otherwise, so only a provider that implements `create_vm` yet gives an agent a container, like the local one, declares it), each sandbox carries the `platform` of the images it runs (linux/amd64 unless its provider sets another when it creates it; None, any, unchecked, for a local sandbox), and code outside `providers/sandbox_providers/` reads those, never which provider it has. `providers/env_providers/` holds the environment providers: `EnvironmentProvider` (an env's containers and state store) and `EnvironmentGatewayProvider`, which renders a docker-compose for the gateway and its MCP servers inside the sandbox; `providers/env_state/` holds env-state providers (`local_postgres` built in). | @@ -155,6 +155,12 @@ orchestrators call it the same way. What callers rely on: a context with `TaskStepContext.to_dict()`, not `dataclasses.asdict`. - A run is resumable: pass a previously persisted `context`, a `start_step` and the existing `instance_id`; completed steps are seeded as SUCCESS. +- An episode (`EpisodeSpec`, `task/episode/`) is one `Task.run()` of a copy of the task in which + the `deploy_agent` and `prompt_agent` steps of each of its `trainer_agents` are the plain-data + `episode_agent` and `episode_turn`, which wait for the caller instead of deploying or prompting. Its + record is tagged `context.metadata["episode"]` and names the original task, so it isn't resumable + with `start_step`. The episode tears down what its run deployed, and an abandoned one is recorded + `cancelled`. ## Tests diff --git a/README.md b/README.md index 67aa102e..b7595571 100644 --- a/README.md +++ b/README.md @@ -44,6 +44,24 @@ uv sync --extra dev source .venv/bin/activate ``` +## Run a task from your own loop + +An episode is one run of a task in which some of its agents are played from outside AgentEnv, by a training loop's model or any code of your own. AgentEnv runs the rest exactly as `agent-env run` does, deploying the envs, running the other agents and scoring, and waits at each of your agent's turns: + +```python +from agent_env.task import EpisodeSpec, Task + +spec = EpisodeSpec.from_task(Task.get("my-task"), turn_timeout=900) +episode = await spec.start() +while (turn := await episode.next_turn()) is not None: + async with turn.tools() as tools: # the agent's env tools, called as its role + ... # your model reads turn.system_prompt and turn.prompt, and acts with tools.list() and tools.call() + await turn.end(response="Done.") +scores = await episode.result() # every score the run recorded, by step id +``` + +The scores are the ones `agent-env run` reports, so making a reward of them is yours. `trainer_agents=` names the agents you play when the task prompts several, `turn_timeout` is each turn's deadline in seconds (the prompt's own `timeout_seconds` otherwise), `episode.abandon(reason)` stops a run early, and `from_task` names any step an outside player can't stand behind. [Running tasks](https://www.agentenvframework.com/docs/tasks/running) covers the rest. + ## Plugins A package of your own can add env types, providers, task steps and stores to agent-env, and `agent-env plugin` lists, checks, adds and removes installed plugins. The [plugin docs](https://www.agentenvframework.com/docs/plugins) show how. diff --git a/src/agent_env/a2a_agent/store.py b/src/agent_env/a2a_agent/store.py index 758a2290..11fb983f 100644 --- a/src/agent_env/a2a_agent/store.py +++ b/src/agent_env/a2a_agent/store.py @@ -3,7 +3,7 @@ import dataclasses import logging -import random +import secrets import string from datetime import datetime, timedelta, timezone from typing import TYPE_CHECKING, Any, Optional, Self @@ -160,7 +160,7 @@ def _doc_store(self): def create_instance(self, deployed: "DeployedA2AAgent", ttl_seconds: int) -> "DeployedA2AAgent": from agent_env.a2a_agent.a2a_agent import DeployedA2AAgent - suffix = "".join(random.choices(string.ascii_lowercase + string.digits, k=8)) + suffix = "".join(secrets.choice(string.ascii_lowercase + string.digits) for _ in range(8)) instance_id = f"{deployed.agent_id}-{suffix}" now = datetime.now(timezone.utc) created_at_utc = now.strftime("%Y-%m-%d %H:%M UTC") diff --git a/src/agent_env/bundle/run.py b/src/agent_env/bundle/run.py index d36e5b2e..73244065 100644 --- a/src/agent_env/bundle/run.py +++ b/src/agent_env/bundle/run.py @@ -27,6 +27,7 @@ from agent_env.store.routing import namespace_routing, run_scope from agent_env.task import Task, record_task_cancelled from agent_env.task.interrupts import Interrupts +from agent_env.task.task import recorded_scores from agent_env.task.teardown import TeardownReport, teardown_run from agent_env.task_step.context import TaskStepContext from agent_env.utils.docker_build import DEFAULT_BUILD_PLATFORM @@ -78,13 +79,7 @@ def failed_steps(self) -> list[dict]: def scores(self) -> dict[str, float]: """Every score the run recorded, by the id of the step that recorded it, or by its verification's key when that isn't a step's verifier id.""" - steps = {step.verifier_id: step.id for step in self.task.steps if getattr(step, "verifier_id", None)} - scores = {} - for key, verification in (self.context.metadata.get("verifications") or {}).items(): - score = verification.get("score") - if isinstance(score, (int, float)): - scores[steps.get(key, key)] = float(score) - return scores + return recorded_scores(self.task.steps, self.context) @property def outcome(self) -> Outcome: diff --git a/src/agent_env/env/store.py b/src/agent_env/env/store.py index b6a23859..24e83ffe 100644 --- a/src/agent_env/env/store.py +++ b/src/agent_env/env/store.py @@ -4,7 +4,7 @@ import dataclasses import logging -import random +import secrets import string from datetime import datetime, timedelta, timezone from typing import TYPE_CHECKING, Any, Optional, Self, TypeVar @@ -209,7 +209,7 @@ def create_instance(self, deployed_env: DeployedEnv, ttl_seconds: int) -> Deploy """Persist a deployed env as an instance record. Returns a new DeployedEnv with instance_id set.""" from agent_env.env.env import DeployedEnv as DeployedEnvCls - suffix = "".join(random.choices(string.ascii_lowercase + string.digits, k=8)) + suffix = "".join(secrets.choice(string.ascii_lowercase + string.digits) for _ in range(8)) instance_id = f"{deployed_env.env_id}-{suffix}" now = datetime.now(timezone.utc) created_at_utc = now.strftime("%Y-%m-%d %H:%M UTC") diff --git a/src/agent_env/providers/env_state/store.py b/src/agent_env/providers/env_state/store.py index 5db56815..9ff0cb48 100644 --- a/src/agent_env/providers/env_state/store.py +++ b/src/agent_env/providers/env_state/store.py @@ -3,7 +3,7 @@ from __future__ import annotations import logging -import random +import secrets import string from datetime import datetime, timedelta, timezone from typing import Optional @@ -38,7 +38,7 @@ def _doc_store(self): def create_instance(self, instance: EnvStateInstance, ttl_seconds: int) -> EnvStateInstance: now = datetime.now(timezone.utc) if not instance.instance_id: - suffix = "".join(random.choices(string.ascii_lowercase + string.digits, k=8)) + suffix = "".join(secrets.choice(string.ascii_lowercase + string.digits) for _ in range(8)) instance.instance_id = f"esi-{suffix}" instance.created_at_utc = now.strftime(_TS_FMT) instance.expires_at_utc = (now + timedelta(seconds=ttl_seconds)).strftime(_TS_FMT) diff --git a/src/agent_env/providers/sandbox_providers/local_sandbox.py b/src/agent_env/providers/sandbox_providers/local_sandbox.py index 3ece7364..6e61bdee 100644 --- a/src/agent_env/providers/sandbox_providers/local_sandbox.py +++ b/src/agent_env/providers/sandbox_providers/local_sandbox.py @@ -16,8 +16,8 @@ import logging import os import platform -import random import re +import secrets import shlex import shutil import signal @@ -139,7 +139,7 @@ def _free_host_port() -> int: try: fcntl.flock(fd, fcntl.LOCK_EX) last = os.read(fd, 16).strip() - start = int(last) + 1 if last.isdigit() else random.choice(_HOST_PORTS) + start = int(last) + 1 if last.isdigit() else secrets.choice(_HOST_PORTS) for i in range(len(_HOST_PORTS)): port = _HOST_PORTS[(start + i - _HOST_PORTS.start) % len(_HOST_PORTS)] if _bindable(port): diff --git a/src/agent_env/task/__init__.py b/src/agent_env/task/__init__.py index 47743161..f7d2f7c2 100644 --- a/src/agent_env/task/__init__.py +++ b/src/agent_env/task/__init__.py @@ -1,6 +1,7 @@ """Task module for AgentEnv.""" from .task import Task +from .episode import Episode, EpisodeError, EpisodeSpec, Turn from .store import ( TaskStore, TaskQuery, get_task_store, set_task_store, reset_task_store, TaskInstance, TaskInstanceStore, get_task_instance_store, set_task_instance_store, reset_task_instance_store, @@ -9,6 +10,10 @@ __all__ = [ "Task", + "Episode", + "EpisodeError", + "EpisodeSpec", + "Turn", "TaskStore", "TaskQuery", "get_task_store", diff --git a/src/agent_env/task/episode/__init__.py b/src/agent_env/task/episode/__init__.py new file mode 100644 index 00000000..bd00530f --- /dev/null +++ b/src/agent_env/task/episode/__init__.py @@ -0,0 +1,9 @@ +"""Run a task with some of its agents played from outside AgentEnv, by a training loop's model or any other caller. + +AgentEnv runs everything else exactly as ``Task.run`` does, so its scores are the ones ``agent-env run`` would report. +""" + +from .run import Episode, EpisodeError, Turn +from .spec import EpisodeSpec + +__all__ = ["Episode", "EpisodeError", "EpisodeSpec", "Turn"] diff --git a/src/agent_env/task/episode/rules.py b/src/agent_env/task/episode/rules.py new file mode 100644 index 00000000..8d766b42 --- /dev/null +++ b/src/agent_env/task/episode/rules.py @@ -0,0 +1,160 @@ +"""Which steps a task may have when some of its agents are played from outside AgentEnv, and why the others are refused. + +Every registered step type is either checked here (``CHECKS``) or listed as untouched by an outside agent +(``UNAFFECTED``). A test holds the registry to the two, so a new step type is classified before an episode runs it. +Plugin step types aren't inspected. +""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from agent_env.task_step.task_step import TaskStep + from agent_env.task_step.task_steps.collect_artifacts import CollectArtifactsTaskStep + from agent_env.task_step.task_steps.deploy_agent import DeployAgentTaskStep + from agent_env.task_step.task_steps.peer_agents import PeerAgentsTaskStep + from agent_env.task_step.task_steps.prompt_agent import PromptAgentTaskStep + from agent_env.task_step.task_steps.register_env_triggers import RegisterEnvTriggersStep + from agent_env.task_step.task_steps.verifiers.rubrics_verifier import RubricsVerifierTaskStep + + +@dataclass(frozen=True) +class Outside: + """The agents played from outside, and their prompts by step id, each with its prompt id.""" + + agents: frozenset[str] + prompts: dict[str, str] + + +def problems(steps: list[TaskStep], outside: list[str]) -> list[str]: + """Why each outside agent, and each step, can't run with those agents played from outside; empty when it can.""" + found = [] + for name in outside: + deploys = [step for step in steps if step.type == "deploy_agent" and step.agent_name == name] + if len(deploys) != 1: + found.append(f"agent {name!r} is played from outside, so it needs exactly one deploy_agent step, " + f"not {len(deploys)}") + scope = Outside(frozenset(outside), { + step.id: step.prompt_id for step in steps if step.type == "prompt_agent" and step.agent_name in outside + }) + for step in steps: + why = ["retry_config can replay an outside turn"] if step.retry_config is not None else [] + if check := CHECKS.get(step.type): + why += check(step, scope) + found.extend(f"step {step.id!r}: {reason}" for reason in why) + return found + + +def _prompt_agent(step: PromptAgentTaskStep, outside: Outside) -> list[str]: + if step.agent_name in outside.agents: + return [ + f"an outside agent's prompt can't have {what}" for what, present in ( + ("parts, only a text prompt", step.parts is not None), + ("a user simulator (max_conversation_turns > 1)", step.max_conversation_turns > 1), + ("output_format", step.output_format), + ("agentenv_tools", step.agentenv_tools), + ("snapshot_config", step.snapshot_config), + ) if present + ] + if step.max_conversation_turns > 1 and step.user_agent_name in outside.agents: + return [f"its user simulator {step.user_agent_name!r} is played from outside"] + return [] + + +def _deploy_agent(step: DeployAgentTaskStep, outside: Outside) -> list[str]: + if step.agent_name not in outside.agents: + return [] + return [ + f"an outside agent can't have {what}" for what, present in ( + ("skills, instructions its player wouldn't get", step.skills), + ("agent_snapshot_files_artifact_id", step.agent_snapshot_files_artifact_id), + ("agent_changelog_object_url", step.agent_changelog_object_url), + ("enable_agent_changelog", step.enable_agent_changelog), + ) if present + ] + + +def _deploys_an_agent(step: TaskStep, outside: Outside) -> list[str]: + if step.agent_name not in outside.agents: + return [] + return [f"{step.type} deploys {step.agent_name!r}, which is played from outside"] + + +def _needs_the_agent(step: TaskStep, outside: Outside) -> list[str]: + if step.agent_name not in outside.agents: + return [] + return [f"{step.type} needs {step.agent_name!r} as a real agent, and it is played from outside"] + + +def _validates_an_agent(step: TaskStep, outside: Outside) -> list[str]: + return ["validates the task's deployed A2A agent, and an episode plays one from outside"] + + +def _collect_artifacts(step: CollectArtifactsTaskStep, outside: Outside) -> list[str]: + why = [] + if step.agent_name in outside.agents and not (step.env_id or step.container_name or step.sandbox_name): + why.append(f"collects from {step.agent_name!r}'s sandbox, and it is played from outside") + if step.manifest_step_id in outside.prompts: + why.append(f"its manifest is outside prompt {step.manifest_step_id!r}") + return why + + +def _register_env_triggers(step: RegisterEnvTriggersStep, outside: Outside) -> list[str]: + if step.executor_agent_name not in outside.agents: + return [] + return [f"its executor {step.executor_agent_name!r} is played from outside"] + + +def _peer_agents(step: PeerAgentsTaskStep, outside: Outside) -> list[str]: + named = dict.fromkeys( + name for peering in step.peerings for name in (peering.source_agent_name, *peering.peer_agent_names) + ) + return [f"peers {name!r}, which is played from outside" for name in named if name in outside.agents] + + +def _rubrics_verifier(step: RubricsVerifierTaskStep, outside: Outside) -> list[str]: + why = [] + if step.use_trajectory and step.prompt_id in outside.prompts.values(): + why.append("grades an outside prompt's trajectory, which an outside player doesn't hand back") + if step.use_agent_judge and step.agent_name in outside.agents: + why.append(f"judges in {step.agent_name!r}'s sandbox, and it is played from outside") + return why + + +CHECKS: dict[str, Callable[..., list[str]]] = { + "prompt_agent": _prompt_agent, + "deploy_agent": _deploy_agent, + "install_agent": _deploys_an_agent, + "deploy_human_agent": _deploys_an_agent, + "add_skills": _needs_the_agent, + "register_agent_triggers": _needs_the_agent, + "snapshot_agent_state": _needs_the_agent, + "load_artifact": _needs_the_agent, + "verify_sandbox": _needs_the_agent, + "run_code": _needs_the_agent, + "collect_artifacts": _collect_artifacts, + "register_env_triggers": _register_env_triggers, + "peer_agents": _peer_agents, + "rubrics_verifier": _rubrics_verifier, + **dict.fromkeys(( + "verify_a2a_agent_card", "verify_a2a_agent_config_identity", "verify_a2a_agent_mcp", "verify_a2a_core_protocol", + "verify_a2a_install", "verify_a2a_modalities", "verify_a2a_peer_agents", "verify_a2a_role", + "verify_a2a_skill_config", "verify_a2a_snapshot", "verify_a2a_system_prompt", "verify_a2a_trajectory", + ), _validates_an_agent), +} + +UNAFFECTED = frozenset({ + # They name no agent. + "aggregate_verifiers", "apply_server_config", "build_mcp_cli", "combine_universe_verdicts", "deploy_env", + "deploy_sandbox", "env_outcome_verifier", "modify_env_tool_access", "reset_env", "review", + "run_container_unit_tests_verifier", "run_docker_container", "snapshot_env", "sync_env_clock", + "validation_gate_aggregator", "verify_env_card", "verify_env_core_protocol", "verify_mcp_tool_schema", + "verify_spec_conformance", "verify_universe_load_export_roundtrip", + # They read only what an outside turn leaves: its reply. + "agent_prompt_response_verifier", "verify_mcp_env_assessment", + # It skips an agent that isn't deployed. + "teardown_sandboxes", +}) diff --git a/src/agent_env/task/episode/run.py b/src/agent_env/task/episode/run.py new file mode 100644 index 00000000..8ee4ddfd --- /dev/null +++ b/src/agent_env/task/episode/run.py @@ -0,0 +1,234 @@ +from __future__ import annotations + +import asyncio +import logging +from collections import deque +from typing import TYPE_CHECKING, Any + +from agentenv_protocol import EnvironmentTool, ToolResult, ToolSession + +from agent_env.providers.sandbox_providers.sandbox_provider import sandbox_request_headers_for_url +from agent_env.task.store import record_task_cancelled +from agent_env.task.task import Task, recorded_scores, substitute_seed_in_steps +from agent_env.task.teardown import teardown_run +from agent_env.task_step.context import TaskStepContext + +if TYPE_CHECKING: + from agent_env.task.episode.steps import _EpisodeTurnStep + +logger = logging.getLogger(__name__) + +_live: dict[str, Episode] = {} + + +class EpisodeError(RuntimeError): + """A task that can't run as an episode, or an episode asked for what it can't give.""" + + +class Turn: + """An outside agent's turn. Its model reads ``system_prompt`` and ``prompt``, acts through ``tools()`` (or an MCP + client of its own on ``mcp_urls``, sending ``headers`` and ``role``), and hands its reply back with ``end()``.""" + + def __init__( + self, *, agent: str, step_id: str, prompt: str, system_prompt: str | None, role: str | None, + mcp_urls: dict[str, str], + ): + self.agent = agent + self.step_id = step_id + self.prompt = prompt + self.system_prompt = system_prompt + self.role = role + self.mcp_urls = mcp_urls + self.headers = {env_id: sandbox_request_headers_for_url(url) for env_id, url in mcp_urls.items()} + self._reply: asyncio.Future[str] = asyncio.get_running_loop().create_future() + + def tools(self, *, timeout: float = 600.0) -> TurnTools: + """A client for the tools of the envs this agent acts in, calling them as its role with the headers each URL + needs; use it with ``async with``. ``timeout`` bounds each call, in seconds.""" + return TurnTools(self, timeout) + + async def end(self, response: str | None = None) -> None: + """Hand the turn back, and the run carries on. No reply records an empty one, as an agent that answers with + nothing leaves. Once the turn is over, ended or not, this does nothing.""" + if self._open: + self._reply.set_result(response or "") + + @property + def _open(self) -> bool: + return not self._reply.done() + + async def _replied_within(self, seconds: float) -> str: + return await asyncio.wait_for(self._reply, seconds) + + def _fail(self, error: Exception) -> None: + if self._open: + self._reply.set_exception(error) + + +class TurnTools: + """The tools of a turn's envs, listed once when it opens: ``list()`` names every one, and ``call()`` routes a + call to the env that offers it.""" + + def __init__(self, turn: Turn, timeout: float): + self._step_id = turn.step_id + self._sessions = { + env_id: ToolSession(mcp_url=url, role=turn.role, headers=turn.headers[env_id], timeout=timeout) + for env_id, url in turn.mcp_urls.items() + } + self._tools: list[EnvironmentTool] = [] + self._offered_by: dict[str, str] = {} + + async def __aenter__(self) -> TurnTools: + try: + for env_id, session in self._sessions.items(): + await session.open() + for tool in await session.list_tools(): + if (other := self._offered_by.get(tool.name)) is not None: + raise EpisodeError(f"tool {tool.name!r} is offered by both env {other!r} and env {env_id!r}") + self._offered_by[tool.name] = env_id + self._tools.append(tool) + except BaseException: + await self._close() + raise + return self + + async def __aexit__(self, *exc_info: object) -> None: + await self._close() + + def list(self) -> list[EnvironmentTool]: + return list(self._tools) + + async def call(self, name: str, arguments: dict[str, Any] | None = None) -> ToolResult: + if name not in self._offered_by: + raise EpisodeError( + f"no env of turn {self._step_id!r} offers tool {name!r}; they offer {sorted(self._offered_by)}") + return await self._sessions[self._offered_by[name]].call_tool(name, arguments) + + async def _close(self) -> None: + await asyncio.gather(*(session.close() for session in self._sessions.values())) + + +class Episode: + """One run of an ``EpisodeSpec``'s task, which waits at each outside turn for the caller to play it. + + Cancelling a caller of ``next_turn()``, ``result()`` or ``abandon()`` cancels that call, never the run. Use an + episode on the event loop that started it.""" + + def __init__(self, id: str, task: Task, context: TaskStepContext): + self.id = id + self._task = task + self._context = context + self._run: asyncio.Task[None] + self._bindings: dict[str, tuple[str | None, dict[str, str]]] = {} + self._waiting: deque[Turn] = deque() + self._turns: set[Turn] = set() + self._changed: asyncio.Future[None] = asyncio.get_running_loop().create_future() + self._tearing_down = False + self._abandoned: str | None = None + + async def next_turn(self) -> Turn | None: + """The next outside turn the run reaches, or None once the run has finished. Raises the run's error if it + failed first. End the turns you hold before waiting on one that depends on them.""" + while True: + changed = self._changed + while self._waiting and self._abandoned is None: + turn = self._waiting.popleft() + if turn._open: + return turn + if self._run.done(): + self._raise_unless_completed() + return None + await self._wait(changed) + + async def result(self) -> dict[str, float]: + """Every score the run recorded, by the id of the step that recorded it, once the run and its teardown are + over: what ``agent-env run`` reports. Making a reward of them is the caller's.""" + while not self._run.done(): + changed = self._changed + if open_turn := next((turn for turn in self._turns if turn._open), None): + raise EpisodeError( + f"turn {open_turn.step_id!r} of agent {open_turn.agent!r} isn't ended, so the run can't finish" + ) + await self._wait(changed) + self._raise_unless_completed() + seeded = substitute_seed_in_steps(self._task.steps, self._context.metadata.get("seed") or {}) + return recorded_scores(seeded, self._context) + + async def abandon(self, reason: str) -> None: + """Stop the run and wait for its teardown; it is recorded cancelled with ``reason``. A run a step has + already failed is only finishing what's in flight, so this ends its open turns and the run's error stands. + Once the run has ended, or while it tears down, this only waits.""" + if not (self._run.done() or self._tearing_down or self._abandoned is not None): + if any(failure.get("is_fatal") for failure in self._context.metadata.get("failed_steps") or []): + for turn in self._turns: + turn._fail(EpisodeError(f"episode abandoned: {reason}")) + else: + self._abandoned = reason + self._run.cancel(msg=reason) + await asyncio.wait((self._run,)) + + async def _launch(self) -> None: + """Start the run and return once its record exists. Cancelled, it takes the run down and waits for teardown.""" + _live[self.id] = self + running = asyncio.Event() + self._run = asyncio.create_task(self._drive(running)) + try: + await running.wait() + except asyncio.CancelledError: + self._abandoned = "start() was cancelled" + self._run.cancel() + await asyncio.wait((self._run,)) + raise + + async def _drive(self, running: asyncio.Event) -> None: + running.set() + try: + await self._task.run(context=self._context, instance_id=self.id) + finally: + self._tearing_down = True + try: + report = await teardown_run(self._context) + if report.still_up: + logger.warning("episode %s: %s may still be up", self.id, [s.sandbox_id for s in report.still_up]) + if self._abandoned is not None: + await asyncio.to_thread(record_task_cancelled, self.id, self._abandoned, Task._utc_now_str()) + finally: + del _live[self.id] + + async def _wait(self, changed: asyncio.Future[None]) -> None: + """Until the run ends or what was ``changed`` when the caller last looked changes again.""" + await asyncio.wait((self._run, changed), return_when=asyncio.FIRST_COMPLETED) + + def _raise_unless_completed(self) -> None: + if self._abandoned is not None: + raise EpisodeError(f"episode abandoned: {self._abandoned}") + if (error := self._run.exception()) is not None: + raise error + + def _bind(self, agent: str, role: str | None, mcp_urls: dict[str, str]) -> None: + self._bindings[agent] = (role, mcp_urls) + + async def _play(self, step: _EpisodeTurnStep) -> str: + if step.agent_name not in self._bindings: + raise RuntimeError( + f"agent {step.agent_name!r} isn't deployed yet; turn {step.id!r} must depend on its deploy_agent step" + ) + role, mcp_urls = self._bindings[step.agent_name] + turn = Turn(agent=step.agent_name, step_id=step.id, prompt=step.prompt, system_prompt=step.system_prompt, + role=role, mcp_urls=dict(mcp_urls)) + self._waiting.append(turn) + self._turns.add(turn) + self._changed.set_result(None) + self._changed = asyncio.get_running_loop().create_future() + try: + return await turn._replied_within(step.timeout_seconds) + except TimeoutError: + raise TimeoutError( + f"turn {step.id!r} of agent {step.agent_name!r} wasn't ended within {step.timeout_seconds}s" + ) from None + finally: + self._turns.discard(turn) + + +def _episode_of(context: TaskStepContext) -> Episode: + return _live[context.metadata["episode"]["id"]] diff --git a/src/agent_env/task/episode/spec.py b/src/agent_env/task/episode/spec.py new file mode 100644 index 00000000..0ddddda0 --- /dev/null +++ b/src/agent_env/task/episode/spec.py @@ -0,0 +1,91 @@ +from __future__ import annotations + +import uuid +from collections.abc import Mapping, Sequence +from typing import Any + +from agent_env.task.episode.rules import problems +from agent_env.task.episode.run import Episode, EpisodeError +from agent_env.task.episode.steps import _EpisodeAgentStep, _EpisodeTurnStep +from agent_env.task.task import Task, substitute_seed_in_steps +from agent_env.task_step.context import TaskStepContext +from agent_env.task_step.task_step import TaskStep +from agent_env.task_step.task_steps.prompt_agent import PromptAgentTaskStep + + +class EpisodeSpec: + """A task read as an episode: the agents a loop plays, and the copy of the task that runs. Built by + ``from_task``.""" + + def __init__(self, *, copy: Task, trainer_agents: tuple[str, ...], seed: dict[str, str] | None): + self._copy = copy + self.trainer_agents = trainer_agents + self._seed = seed + + @classmethod + def from_task( + cls, task: Task, *, trainer_agents: Sequence[str] | None = None, seed: Mapping[str, str] | None = None, + turn_timeout: float | None = None, + ) -> EpisodeSpec: + """``trainer_agents`` are the agents the loop plays, by default the task's only prompted agent. + ``turn_timeout`` is every turn's deadline in seconds, in place of its prompt's own ``timeout_seconds``, + which budgets an agent's run rather than the loop playing it.""" + seed = dict(seed) if seed else None + seeded = substitute_seed_in_steps(task.steps, seed or {}) + outside = _trainer_agents(task, seeded, trainer_agents) + if found := problems(seeded, outside): + raise EpisodeError(f"task {task.id!r} can't run as an episode: " + "; ".join(found)) + deploys = { + as_run.agent_name: step for step, as_run in zip(task.steps, seeded) + if step.type == "deploy_agent" and as_run.agent_name in outside + } + steps = [_stand_in(step, as_run, deploys, turn_timeout) for step, as_run in zip(task.steps, seeded)] + return cls(copy=Task(id=task.id, version=task.version, steps=steps), trainer_agents=tuple(outside), seed=seed) + + async def start(self) -> Episode: + """Start one run of the task, a new episode each call. It returns once the run's record exists, before + anything is deployed; ``next_turn()`` is where the deploys and setup are waited for.""" + episode_id = uuid.uuid4().hex + metadata: dict[str, Any] = {"episode": {"id": episode_id, "trainer_agents": list(self.trainer_agents)}} + if self._seed: + metadata["seed"] = dict(self._seed) + episode = Episode(episode_id, self._copy, TaskStepContext(metadata=metadata)) + await episode._launch() + return episode + + +def _trainer_agents(task: Task, steps: list[TaskStep], named: Sequence[str] | None) -> list[str]: + prompted = list(dict.fromkeys(step.agent_name for step in steps if step.type == "prompt_agent")) + if named is None: + if len(prompted) != 1: + raise EpisodeError( + f"task {task.id!r} prompts {prompted}; name the ones to play with trainer_agents=" if prompted + else f"task {task.id!r} prompts no agent" + ) + return prompted + outside = list(dict.fromkeys(named)) + if not outside: + raise EpisodeError(f"trainer_agents= names no agent; task {task.id!r} prompts {prompted}") + if unprompted := [name for name in outside if name not in prompted]: + raise EpisodeError(f"task {task.id!r} doesn't prompt {unprompted}; it prompts {prompted}") + return outside + + +def _stand_in(step: TaskStep, as_run: TaskStep, deploys: dict[str, TaskStep], turn_timeout: float | None) -> TaskStep: + """The step the episode's copy of the task runs in place of ``step``, judged by ``as_run``, the step as the + seed makes it; ``deploys`` are the outside agents', by name. A stand-in keeps the step's own fields, which the + run seeds.""" + if as_run.type == "deploy_agent" and as_run.agent_name in deploys: + return _EpisodeAgentStep( + id=step.id, version=step.version, depends_on=step.depends_on, fail_task_on_error=step.fail_task_on_error, + agent_name=step.agent_name, env_ids=step.env_ids, env_step_id=step.env_step_id, role=step.role, + ) + if as_run.type == "prompt_agent" and as_run.agent_name in deploys: + timeout = step.timeout_seconds if turn_timeout is None else turn_timeout + return _EpisodeTurnStep( + id=step.id, version=step.version, depends_on=step.depends_on, fail_task_on_error=step.fail_task_on_error, + agent_name=step.agent_name, prompt_id=step.prompt_id, prompt=step.prompt, + system_prompt=step.system_prompt or deploys[as_run.agent_name].system_prompt, + timeout_seconds=PromptAgentTaskStep.DEFAULT_PROMPT_TIMEOUT_SECONDS if timeout is None else timeout, + ) + return step diff --git a/src/agent_env/task/episode/steps.py b/src/agent_env/task/episode/steps.py new file mode 100644 index 00000000..66ffeb93 --- /dev/null +++ b/src/agent_env/task/episode/steps.py @@ -0,0 +1,80 @@ +from __future__ import annotations + +from typing import Any, ClassVar, Optional + +from agent_env.entity_refs import EntityRef +from agent_env.task.episode.run import _episode_of +from agent_env.task_step.context import PromptResponse, TaskStepContext +from agent_env.task_step.task_step import TaskStep +from agent_env.task_step.task_steps.deploy_agent import bound_deployment, live_mcp_url + + +class _EpisodeAgentStep(TaskStep): + """An outside agent's ``deploy_agent``: deploys nothing, and binds the agent to its envs where it stood.""" + + type: ClassVar[str] = "episode_agent" + entity_refs = (EntityRef.env("env_ids[]"),) + + def __init__( + self, id: str, version: Optional[int], *, agent_name: str, env_ids: list[str], env_step_id: str | None, + role: str | None, depends_on: Optional[list] = None, fail_task_on_error: bool = True, + ): + super().__init__(id, version, depends_on=depends_on, fail_task_on_error=fail_task_on_error) + self.agent_name = agent_name + self.env_ids = list(env_ids) + self.env_step_id = env_step_id + self.role = role + + def to_dict(self) -> dict[str, Any]: + return {**super().to_dict(), "agent_name": self.agent_name, "env_ids": self.env_ids, + "env_step_id": self.env_step_id, "role": self.role} + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> _EpisodeAgentStep: + return cls(**cls._base_from_dict(data), agent_name=data["agent_name"], env_ids=data["env_ids"], + env_step_id=data["env_step_id"], role=data["role"]) + + async def execute(self, context: TaskStepContext) -> TaskStepContext: + mcp_urls = {} + for env_id in self.env_ids: + deployed = bound_deployment(context, env_id, env_step_id=self.env_step_id, step_id=self.id) + mcp_urls[env_id] = deployed.mcp_url if deployed is not None else live_mcp_url(env_id) + _episode_of(context)._bind(self.agent_name, self.role, mcp_urls) + return context + + +class _EpisodeTurnStep(TaskStep): + """An outside agent's ``prompt_agent``: hands the turn to the caller and records the reply it gets back.""" + + type: ClassVar[str] = "episode_turn" + entity_refs = () + + def __init__( + self, id: str, version: Optional[int], *, agent_name: str, prompt_id: str, prompt: str, + system_prompt: str | None, timeout_seconds: float, depends_on: Optional[list] = None, + fail_task_on_error: bool = True, + ): + super().__init__(id, version, depends_on=depends_on, fail_task_on_error=fail_task_on_error) + self.agent_name = agent_name + self.prompt_id = prompt_id + self.prompt = prompt + self.system_prompt = system_prompt + self.timeout_seconds = timeout_seconds + + def to_dict(self) -> dict[str, Any]: + return {**super().to_dict(), "agent_name": self.agent_name, "prompt_id": self.prompt_id, + "prompt": self.prompt, "system_prompt": self.system_prompt, "timeout_seconds": self.timeout_seconds} + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> _EpisodeTurnStep: + return cls(**cls._base_from_dict(data), agent_name=data["agent_name"], prompt_id=data["prompt_id"], + prompt=data["prompt"], system_prompt=data["system_prompt"], + timeout_seconds=data["timeout_seconds"]) + + async def execute(self, context: TaskStepContext) -> TaskStepContext: + reply = await _episode_of(context)._play(self) + context.prompt_responses.append(PromptResponse( + prompt_id=self.prompt_id, response=reply, prompt_text=self.prompt, agent_name=self.agent_name, + step_id=self.id, + )) + return context diff --git a/src/agent_env/task/store.py b/src/agent_env/task/store.py index 2aa650cb..7ea6c3d9 100644 --- a/src/agent_env/task/store.py +++ b/src/agent_env/task/store.py @@ -6,7 +6,7 @@ import copy import dataclasses import logging -import random +import secrets import string import time from dataclasses import dataclass, field @@ -299,7 +299,7 @@ def create_instance( completed_step_docs = [ dataclasses.asdict(e) for e in completed_steps or [] ] - suffix = "".join(random.choices(string.ascii_lowercase + string.digits, k=8)) + suffix = "".join(secrets.choice(string.ascii_lowercase + string.digits) for _ in range(8)) instance_id = f"{task_id}-{suffix}" created_at_utc = datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M UTC") doc = _new_instance_doc( diff --git a/src/agent_env/task/task.py b/src/agent_env/task/task.py index 30c54cd2..46960671 100644 --- a/src/agent_env/task/task.py +++ b/src/agent_env/task/task.py @@ -96,7 +96,6 @@ def substitute_seed_in_steps(steps: list["TaskStep"], seed: dict) -> list["TaskS objects, so it executes exactly what the original run executed.""" if not seed: return list(steps) - from agent_env.task_step.registry import get_task_step_registry def _sub(value): if isinstance(value, str): @@ -111,14 +110,25 @@ def _sub(value): from agent_env.task_step.task_step import attach_retry_config - registry = get_task_step_registry() out: list["TaskStep"] = [] for step in steps: data = _sub(step.to_dict()) - out.append(attach_retry_config(registry[step.type].from_dict(data), data)) + out.append(attach_retry_config(type(step).from_dict(data), data)) return out +def recorded_scores(steps: list["TaskStep"], context: TaskStepContext) -> dict[str, float]: + """Every score a run of ``steps`` recorded, by the id of the step that recorded it, or by its verification's key + when that isn't a step's verifier id.""" + step_ids = {step.verifier_id: step.id for step in steps if getattr(step, "verifier_id", None)} + scores = {} + for key, verification in (context.metadata.get("verifications") or {}).items(): + score = verification.get("score") + if isinstance(score, (int, float)): + scores[step_ids.get(key, key)] = float(score) + return scores + + def _dependency_ids(steps: list["TaskStep"]) -> dict[str, set[str]]: """Direct dependency ids per step; depends_on None = all prior (scheduler semantics).""" deps: dict[str, set[str]] = {} diff --git a/src/agent_env/task_step/task_steps/deploy_agent.py b/src/agent_env/task_step/task_steps/deploy_agent.py index 4cc81176..b4e1c6e8 100644 --- a/src/agent_env/task_step/task_steps/deploy_agent.py +++ b/src/agent_env/task_step/task_steps/deploy_agent.py @@ -85,7 +85,7 @@ def _choose_mcp_url(agent, env) -> str: return reachable_url(env.mcp_url, from_sandbox_type=env.sandbox_type, to_sandbox_type=agent.sandbox_type) -def _live_mcp_url(env_id: str) -> str: +def live_mcp_url(env_id: str) -> str: """The MCP URL of an env that is referenced by id but never deployed here. Such an env is a pointer at a live endpoint: it exposes ``mcp_url`` and has no @@ -100,6 +100,35 @@ def _live_mcp_url(env_id: str) -> str: return url.strip() +def _deploy_step_of(deployed_env: DeployedEnv) -> Optional[str]: + return (deployed_env.metadata or {}).get("deploy_step_id") + + +def bound_deployment( + context: TaskStepContext, env_id: str, *, env_step_id: Optional[str], step_id: str, +) -> Optional[DeployedEnv]: + """The deployment of `env_id` that the agent deployed by step `step_id` is bound to, by `env_step_id`, or None if + it isn't deployed here.""" + matches = [d for d in context.deployed_envs if d.env_id == env_id] + if env_step_id is not None: + # A filter, not a hint: an unmatched id is an error, never a substitute. + # Envs with no deployments fall through to the caller's live-endpoint path. + named = [d for d in matches if _deploy_step_of(d) == env_step_id] + if matches and not named: + raise RuntimeError( + f"deploy_agent '{step_id}': env_step_id '{env_step_id}' names no deployment of " + f"env '{env_id}' (deployed by: {[_deploy_step_of(d) for d in matches]})" + ) + matches = named + if len(matches) > 1: + raise RuntimeError( + f"deploy_agent '{step_id}': {len(matches)} deployments of env '{env_id}' " + f"(deployed by: {[_deploy_step_of(d) for d in matches]}). Set env_step_id " + "to the deploy_env step this agent should bind to." + ) + return matches[0] if matches else None + + def _skill_fields(skill: dict) -> dict: """A ``skills`` entry as ``A2AAgent.register_skill`` arguments. ``content`` is an older name for ``skill_md``, and ``skill_s3_url`` and ``s3_uri`` for ``object_url``; ``skill_md`` is sent verbatim.""" @@ -276,7 +305,10 @@ async def _execute(self, context: TaskStepContext, deployed_agents: list) -> Tas ) # Resolved before anything is provisioned so an unbound env fails before a sandbox exists. - env_deployments = {env_id: self._deployed_env(context, env_id) for env_id in self.env_ids} + env_deployments = { + env_id: bound_deployment(context, env_id, env_step_id=self.env_step_id, step_id=self.id) + for env_id in self.env_ids + } existing = next((a for a in context.deployed_agents if a.agent_name == self.agent_name), None) if existing is not None: @@ -360,7 +392,7 @@ async def _execute(self, context: TaskStepContext, deployed_agents: list) -> Tas if not card_name: logger.warning(f"Env {env_id} has no env card name; the agent will pick its own MCP alias") else: - url = host_url_for(_live_mcp_url(env_id), deployed.sandbox_type) + url = host_url_for(live_mcp_url(env_id), deployed.sandbox_type) card_name = None if mcp_ext: body = _mcp_add_body(mcp_ext, url, sandbox_request_headers_for_url(url), card_name, self.role) @@ -456,31 +488,6 @@ async def _execute(self, context: TaskStepContext, deployed_agents: list) -> Tas )) return context - @staticmethod - def _deploy_step_of(deployed_env: DeployedEnv) -> Optional[str]: - return (deployed_env.metadata or {}).get("deploy_step_id") - - def _deployed_env(self, context: TaskStepContext, env_id: str) -> Optional[DeployedEnv]: - """The deployment of `env_id` this step is bound to, or None if it isn't deployed here.""" - matches = [d for d in context.deployed_envs if d.env_id == env_id] - if self.env_step_id is not None: - # A filter, not a hint: an unmatched id is an error, never a substitute. - # Envs with no deployments fall through to the caller's live-endpoint path. - named = [d for d in matches if self._deploy_step_of(d) == self.env_step_id] - if matches and not named: - raise RuntimeError( - f"deploy_agent '{self.id}': env_step_id '{self.env_step_id}' names no deployment of " - f"env '{env_id}' (deployed by: {[self._deploy_step_of(d) for d in matches]})" - ) - matches = named - if len(matches) > 1: - raise RuntimeError( - f"deploy_agent '{self.id}': {len(matches)} deployments of env '{env_id}' " - f"(deployed by: {[self._deploy_step_of(d) for d in matches]}). Set env_step_id " - "to the deploy_env step this agent should bind to." - ) - return matches[0] if matches else None - def _resolved_ttl_seconds(self, context: TaskStepContext) -> int: override = context.metadata.get("user_overrides", {}).get("ttl_seconds") if override is not None: diff --git a/tst/integration/env/envs/server_env_local_test.py b/tst/integration/env/envs/server_env_local_test.py index 69fa4c6c..2481076c 100644 --- a/tst/integration/env/envs/server_env_local_test.py +++ b/tst/integration/env/envs/server_env_local_test.py @@ -3,12 +3,14 @@ The env deploys, a later process restores its record and loads data into the server, and the teardown step, which rebuilds each sandbox from disk, removes the container. A load is timed by the size of the payload it staged, through the deploy's own handle and a restored one; and the server's state reads as JSON whether ``data/get`` answers with -data or with a file bundle. A loop of one's own lists, declares and calls the server's tools through ``ToolSession``. -Requires a docker daemon; spins up a throwaway ``registry:2`` and skips if it can't start. +data or with a file bundle. A loop of one's own lists, declares and calls the server's tools through ``ToolSession``, +and plays a task's agent as an episode, around the echo agent. Requires a docker daemon; spins up a throwaway +``registry:2`` and skips if it can't start. """ import asyncio import json +import re import shutil import socket import subprocess @@ -30,8 +32,11 @@ from agent_env.env.envs.mcp_server import MCPServerEnv from agent_env.env.gateway import constants from agent_env.store.image_store import LocalRegistryImageStore +from agent_env.store.object_store.local.grant_server import grant_server +from agent_env.task import EpisodeSpec, Task from agent_env.task_step.context import TaskStepContext from agent_env.task_step.task_steps.teardown_sandboxes import TORN_DOWN_KEY, TeardownSandboxesTaskStep +from tst.util.a2a_test_agent import put_test_agent pytestmark = pytest.mark.integration @@ -75,8 +80,10 @@ def local_stack(monkeypatch, tmp_path): pytest.skip("local registry did not become ready") owned: list[str] = [] + sandboxes = tmp_path / "sandboxes" monkeypatch.chdir(tmp_path) - monkeypatch.setenv("AGENT_ENV_LOCAL_SANDBOX_DIR", str(tmp_path / "sandboxes")) + monkeypatch.setenv("XDG_STATE_HOME", str(tmp_path / "state")) + monkeypatch.setenv("AGENT_ENV_LOCAL_SANDBOX_DIR", str(sandboxes)) for var in ("DOCUMENT", "OBJECT", "IMAGE"): monkeypatch.setenv(f"AGENT_ENV_{var}_STORE", "local") configure() @@ -85,8 +92,14 @@ def local_stack(monkeypatch, tmp_path): try: yield owned finally: + # A run's sandboxes aren't known up front: each has a work dir in this test's sandbox root, named after it. + for work_dir in sandboxes.glob("agent-env-*"): + if m := re.match(r"agent-env-(local-[0-9a-f]+)-", work_dir.name): + owned.append(f"agent-{m.group(1)}") for container_name in owned: _docker("rm", "-f", container_name) + # Agents upload through this process's grant server, whose certificate is from this test's state root. + grant_server(None, None).close() reset_artifact_store() reset_config() _docker("rm", "-f", reg_name) @@ -185,6 +198,86 @@ async def add(i: int) -> tuple[int, dict, str]: assert json.loads((await after_restart.call_tool("list_items")).text) == {"items": []} +VERIFY_CRATE = """ +from agentenv_protocol import client + + +async def verify(mcp_url): + items = (await client.get_data(mcp_url.removesuffix("/mcp"))).parts[0].data["items"] + return [{"id": "crate", "result": "crate" in items}] +""" + + +@pytest.mark.asyncio +async def test_a_tasks_agent_is_played_from_a_loop_of_ones_own_as_an_episode(local_stack, tmp_path): + """A trainer's view of a seeded task: the env deployed and loaded and the echo agent prompted around the outside + agent, which plays its turn through the env's tools as its role and is never deployed; its reply recorded and + scored under the verifier's step; and every episode, played, not played, abandoned or forgotten past its turn's + deadline, leaving nothing behind.""" + uid = uuid.uuid4().hex[:8] + env = MCPServerEnv.put(id=f"server-items-{uid}", docker_image_artifact=_put_items_image(f"server-items-{uid}"), + environment_name="items", env_provider_type="server") + verify = FileArtifact.put_bytes(id=f"verify-crate-{uid}", description="episode e2e", filename="verify.py", + content=VERIFY_CRATE.encode(), content_type="text/x-python") + task = Task.from_dict({"id": f"crate-{uid}", "steps": [ + {"id": "env", "type": "deploy_env", "env_id": ""}, + {"id": "load", "type": "load_artifact", "env_id": "", "artifact_id": _items_artifact(uid).id}, + {"id": "policy", "type": "deploy_agent", "agent_name": "policy", "env_ids": [""], + "role": "@example.com"}, + {"id": "echo", "type": "deploy_agent", "agent_name": "echo", "a2a_agent_id": put_test_agent(f"echo-{uid}").id}, + {"id": "ask-echo", "type": "prompt_agent", "agent_name": "echo", "prompt_id": "ask-echo", + "prompt": "Count the crates."}, + {"id": "play", "type": "prompt_agent", "agent_name": "policy", "prompt_id": "play", "prompt": "Add a ."}, + {"id": "check", "type": "env_outcome_verifier", "env_id": "", "file_artifact_id": verify.id, + "verifier_id": "crate-"}, + ]}) + seed = {"env": env.id, "who": "alice", "item": "crate", "week": "3"} + spec = EpisodeSpec.from_task(task, trainer_agents=["policy"], seed=seed) + + async def play(episode, add): + turn = await episode.next_turn() + assert turn.prompt == "Add a crate." and turn.role == "alice@example.com" + assert [agent["agent_name"] for agent in Task.get_instance(episode.id).context["deployed_agents"]] == ["echo"] + if add: + async with turn.tools() as tools: + assert {"items_add_item", "items_whoami"} <= {tool.name for tool in tools.list()} + await tools.call("items_add_item", {"item": "crate"}) + assert json.loads((await tools.call("items_whoami")).text)["role"] == "alice@example.com" + await turn.end("Added a crate." if add else None) + return await episode.result() + + sandboxes = tmp_path / "sandboxes" + played, unplayed = await spec.start(), await spec.start() + assert await asyncio.gather(play(played, True), play(unplayed, False)) == [{"check": 1.0}, {"check": 0.0}] + assert list(sandboxes.glob("agent-env-*")) == [] + record = Task.get_instance(played.id) + assert record.status == "completed" and record.context["metadata"]["episode"]["trainer_agents"] == ["policy"] + assert [agent["agent_name"] for agent in record.context["deployed_agents"]] == ["echo"] + replies = {response["step_id"]: response["response"] for response in record.context["prompt_responses"]} + assert replies["play"] == "Added a crate." and replies["ask-echo"].startswith("Echo:") + + abandoned = await spec.start() + await abandoned.next_turn() + assert len(list(sandboxes.glob("agent-env-*"))) == 2 + await abandoned.abandon("test") + assert (Task.get_instance(abandoned.id).status, Task.get_instance(abandoned.id).error) == ("cancelled", "test") + assert list(sandboxes.glob("agent-env-*")) == [] + + forgotten = await EpisodeSpec.from_task(task, trainer_agents=["policy"], seed=seed, turn_timeout=3).start() + await forgotten.next_turn() + with pytest.raises(TimeoutError, match="turn 'play' of agent 'policy' wasn't ended within 3s"): + await forgotten.next_turn() + assert Task.get_instance(forgotten.id).status == "failed" + assert list(sandboxes.glob("agent-env-*")) == [] + for episode in (played, unplayed, abandoned, forgotten): + context = Task.get_instance(episode.id).context + sandbox_ids = [entry["sandbox_id"] for key in ("deployed_envs", "deployed_agents") + for entry in context[key] if entry.get("sandbox_id")] + assert len(sandbox_ids) == 2 + for sandbox_id in sandbox_ids: + assert _docker("ps", "-aq", "--filter", f"name=^/agent-{sandbox_id}$").stdout.strip() == "" + + def _put_items_image(artifact_id: str) -> DockerImageArtifact: """The in-memory items server (card 'items'), built with the vendored agentenv_protocol and pushed to the registry.""" data = _REPO / "tst" / "data" / "agentenv_mcp" diff --git a/tst/integration/env/envs/server_env_modal_test.py b/tst/integration/env/envs/server_env_modal_test.py index 9895b746..ed7ef874 100644 --- a/tst/integration/env/envs/server_env_modal_test.py +++ b/tst/integration/env/envs/server_env_modal_test.py @@ -5,12 +5,14 @@ bundle. Needs the resolved config to reach Modal and an image store Modal can pull from; skipped otherwise. """ +import asyncio import json import shutil import tempfile import uuid from pathlib import Path +import modal import pytest from agentenv_protocol import client as protocol_v1 @@ -19,6 +21,7 @@ from agent_env.env.env import Env from agent_env.env.envs.mcp_server import MCPServerEnv from agent_env.env.gateway import constants +from agent_env.task import EpisodeSpec, Task from agent_env.task_step.context import TaskStepContext from agent_env.task_step.task_steps.teardown_sandboxes import TeardownSandboxesTaskStep from tst.util.capabilities import skip_without_remote_sandbox @@ -56,6 +59,90 @@ async def test_a_modal_load_is_timed_by_its_payload_and_the_servers_state_reads_ ) +VERIFY_CRATE = """ +from agentenv_protocol import client + + +async def verify(mcp_url): + items = (await client.get_data(mcp_url.removesuffix("/mcp"))).parts[0].data["items"] + return [{"id": "crate", "result": "crate" in items}] +""" + + +@pytest.mark.asyncio +async def test_an_episodes_env_on_modal_is_played_from_this_machine_and_no_ending_leaves_it_running(): + """A loop on this machine plays an agent whose env the task deploys on Modal: the env's tools reached over its + tunnel as the agent's role, the score read from the env's state there, and every episode, played, abandoned or + forgotten past its turn's deadline, leaving its sandbox stopped.""" + uid = uuid.uuid4().hex[:8] + env = MCPServerEnv.put(id=f"server-items-episode-{uid}", docker_image_artifact=_items_context(uid), + environment_name="items", env_provider_type="server") + verify = FileArtifact.put_bytes(id=f"verify-crate-modal-{uid}", description="episode on Modal", + filename="verify.py", content=VERIFY_CRATE.encode(), content_type="text/x-python") + task = Task.from_dict({"id": f"crate-modal-{uid}", "steps": [ + {"id": "env", "type": "deploy_env", "env_id": env.id, "sandbox_type": "modal", "ttl_seconds": 900}, + {"id": "load", "type": "load_artifact", "env_id": env.id, "artifact_id": _items_artifact(uid).id}, + {"id": "policy", "type": "deploy_agent", "agent_name": "policy", "env_ids": [env.id], + "role": "alice@example.com"}, + {"id": "play", "type": "prompt_agent", "agent_name": "policy", "prompt_id": "play", "prompt": "Add a crate."}, + {"id": "check", "type": "env_outcome_verifier", "env_id": env.id, "file_artifact_id": verify.id, + "verifier_id": "crate"}, + ]}) + spec = EpisodeSpec.from_task(task) + started = [] + + async def start(spec): + started.append(episode := await spec.start()) + return episode + + try: + played = await start(spec) + turn = await played.next_turn() + assert turn.mcp_urls[env.id].startswith("https://") + async with turn.tools() as tools: + await tools.call("items_add_item", {"item": "crate"}) + assert json.loads((await tools.call("items_whoami")).text)["role"] == "alice@example.com" + await turn.end("Added a crate.") + assert await played.result() == {"check": 1.0} + + abandoned = await start(spec) + await abandoned.next_turn() + await abandoned.abandon("test") + assert Task.get_instance(abandoned.id).status == "cancelled" + + forgotten = await start(EpisodeSpec.from_task(task, turn_timeout=3)) + await forgotten.next_turn() + with pytest.raises(TimeoutError, match="wasn't ended within 3s"): + await forgotten.next_turn() + + for episode in (played, abandoned, forgotten): + [sandbox_id] = [entry["sandbox_id"] for entry in Task.get_instance(episode.id).context["deployed_envs"]] + await _stopped(sandbox_id) + finally: + for episode in started: + await episode.abandon("the test ended") + +async def _stopped(sandbox_id: str) -> None: + sandbox = await modal.Sandbox.from_id.aio(sandbox_id) + async with asyncio.timeout(60): + while await sandbox.poll.aio() is None: + await asyncio.sleep(2) + + +def _items_context(uid: str) -> DockerImageArtifact: + """The items server as a build context, which Modal builds where it runs, so no registry Modal can pull from is + needed.""" + data = _REPO / "tst" / "data" / "agentenv_mcp" + with tempfile.TemporaryDirectory() as build_dir: + context = Path(build_dir) + shutil.copytree(_REPO / "packages" / "agentenv-protocol" / "src" / "agentenv_protocol", context / "agentenv_protocol", + ignore=shutil.ignore_patterns("__pycache__")) + for name in ("server.py", "Dockerfile", "seed.json"): + shutil.copy(data / name, context / name) + return DockerImageArtifact.put_context(f"server-items-episode-image-{uid}", description="episode on Modal", + context_path=str(context), dockerfile_path="Dockerfile") + + def _items_image() -> DockerImageArtifact: """The in-memory items server (card 'items'), built for linux/amd64 with the vendored agentenv_protocol.""" data = _REPO / "tst" / "data" / "agentenv_mcp" diff --git a/tst/unit/env/store_test.py b/tst/unit/env/store_test.py index 6fadbd08..1a401aa6 100644 --- a/tst/unit/env/store_test.py +++ b/tst/unit/env/store_test.py @@ -6,6 +6,7 @@ """ from __future__ import annotations +import random from dataclasses import dataclass import pytest @@ -219,3 +220,15 @@ def test_merge_metadata_is_a_noop_for_empty_values(tmp_path): assert recording.updates == [] assert _read(bundle)["metadata"] is None + + +def test_instance_ids_do_not_follow_a_reseeded_random(): + """A trainer reseeds ``random`` between deploys; the instance ids must still differ, or records collide.""" + store = EnvInstanceStore() + set_document_store(FakeDocumentStore()) + ids = [] + for _ in range(2): + random.seed(0) + ids.append(store.create_instance(DeployedEnv(env_id="portsim", env_version=1), ttl_seconds=60).instance_id) + + assert ids[0] != ids[1] diff --git a/tst/unit/providers/sandbox_providers/local_sandbox_host_ports_test.py b/tst/unit/providers/sandbox_providers/local_sandbox_host_ports_test.py index be4aab6f..251ad6d9 100644 --- a/tst/unit/providers/sandbox_providers/local_sandbox_host_ports_test.py +++ b/tst/unit/providers/sandbox_providers/local_sandbox_host_ports_test.py @@ -13,6 +13,7 @@ import errno import os import platform +import random import socket import threading @@ -132,13 +133,24 @@ def test_host_ports_wrap_round_their_range(busy): @pytest.mark.parametrize("cursor", [None, "", "not a port"]) def test_a_missing_or_unreadable_cursor_starts_at_a_random_port(busy, monkeypatch, cursor): - monkeypatch.setattr(local_sandbox_module.random, "choice", lambda ports: 25000) + monkeypatch.setattr(local_sandbox_module.secrets, "choice", lambda ports: 25000) if cursor is not None: _last_handed_out(cursor) assert local_sandbox_module._free_host_port() == 25000 +def test_the_first_port_of_a_state_root_does_not_follow_a_reseeded_random(busy, tmp_path, monkeypatch): + """A trainer reseeds ``random``; fresh state roots must still start apart in the range.""" + starts = [] + for root in ("a", "b", "c"): + monkeypatch.setenv("XDG_STATE_HOME", str(tmp_path / root)) + random.seed(0) + starts.append(local_sandbox_module._free_host_port()) + + assert len(set(starts)) > 1 + + def test_concurrent_draws_never_share_a_host_port(busy): """Each draw opens the cursor afresh, so the lock serialises threads as it does processes.""" drawn: list[int] = [] diff --git a/tst/unit/task/episode/__init__.py b/tst/unit/task/episode/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tst/unit/task/episode/run_test.py b/tst/unit/task/episode/run_test.py new file mode 100644 index 00000000..4cd5ccdb --- /dev/null +++ b/tst/unit/task/episode/run_test.py @@ -0,0 +1,588 @@ +"""The run of an episode, on fake steps: its turns, its scores, abandoning it, teardown and its record.""" + +import asyncio + +import pytest +import pytest_asyncio +from agentenv_protocol import EnvironmentTool, ToolResult + +import agent_env.task.episode.run as episode_run +from agent_env.config.runtime import Config +from agent_env.env.env import DeployedEnv +from agent_env.task import EpisodeError, EpisodeSpec, Task +from agent_env.task.teardown import TeardownReport +from agent_env.task_step.task_step import TaskStep + +class _Env(TaskStep): + """Deploys ``env_id`` with no sandbox, or fails as an env that doesn't come up.""" + + type = "episode_test_env" + entity_refs = () + + def __init__(self, id, version, env_id="items", fail=False, depends_on=None, fail_task_on_error=True): + super().__init__(id, version, depends_on=depends_on, fail_task_on_error=fail_task_on_error) + self.env_id = env_id + self.fail = fail + + def to_dict(self): + return {**super().to_dict(), "env_id": self.env_id, "fail": self.fail} + + @classmethod + def from_dict(cls, data): + return cls(**cls._base_from_dict(data), env_id=data.get("env_id", "items"), fail=data.get("fail", False)) + + async def execute(self, context): + if self.fail: + raise RuntimeError("the env didn't come up") + context.deployed_envs.append(DeployedEnv( + env_id=self.env_id, env_version=1, mcp_url=f"http://127.0.0.1:20031/{self.env_id}/mcp", + metadata={"deploy_step_id": self.id}, + )) + return context + + +GATES: dict[str, asyncio.Event] = {} + + +def _gate(step_id): + return GATES.setdefault(step_id, asyncio.Event()) + + +class _Gate(TaskStep): + """Waits until the test opens the gate of its id.""" + + type = "episode_test_gate" + entity_refs = () + + @classmethod + def from_dict(cls, data): + return cls(**cls._base_from_dict(data)) + + async def execute(self, context): + await _gate(self.id).wait() + return context + + +class _Replies(TaskStep): + """Scores 1 when the run recorded replies and every one has text, else 0.""" + + type = "episode_test_replies" + entity_refs = () + + def __init__(self, id, version, verifier_id="replies", depends_on=None, fail_task_on_error=True): + super().__init__(id, version, depends_on=depends_on, fail_task_on_error=fail_task_on_error) + self.verifier_id = verifier_id + + def to_dict(self): + return {**super().to_dict(), "verifier_id": self.verifier_id} + + @classmethod + def from_dict(cls, data): + return cls(**cls._base_from_dict(data), verifier_id=data["verifier_id"]) + + async def execute(self, context): + replies = [response.response for response in context.prompt_responses] + context.metadata.setdefault("verifications", {})[self.verifier_id] = { + "score": float(bool(replies) and all(replies))} + return context + + +class _Teardowns: + """Stands in for ``teardown_run``: records each run it tears down, and waits on ``hold`` when it's set.""" + + def __init__(self): + self.started = [] + self.finished = [] + self.hold = None + + async def __call__(self, context, **_): + self.started.append(context.instance_id) + if self.hold is not None: + await self.hold.wait() + self.finished.append(context.instance_id) + return TeardownReport() + + +class _ToolSessions: + """Stands in for ``ToolSession``: the session on each URL serves the tools ``offered`` names for it, and every + session opened, call made and session closed is recorded.""" + + def __init__(self, offered): + self.offered = offered + self.opened = [] + self.calls = [] + self.closed = [] + + def __call__(self, *, mcp_url, role, headers, timeout): + sessions = self + + class _Session: + async def open(self): + sessions.opened.append((mcp_url, role, headers, timeout)) + return self + + async def list_tools(self): + return [EnvironmentTool(name=name) for name in sessions.offered[mcp_url]] + + async def call_tool(self, name, arguments=None): + sessions.calls.append((mcp_url, name, arguments)) + return ToolResult(content=[{"type": "text", "text": f"{name} on {mcp_url}"}]) + + async def close(self): + sessions.closed.append(mcp_url) + + return _Session() + + +@pytest_asyncio.fixture +async def teardowns(monkeypatch, local_stores): + """Fake steps registered and ``teardown_run`` stood in for; an episode the test leaves running is abandoned + before they are undone.""" + steps = {**Config().task_step_registry(), **{step.type: step for step in (_Env, _Gate, _Replies)}} + monkeypatch.setattr(Config, "task_step_registry", lambda self: steps) + GATES.clear() + teardowns = _Teardowns() + monkeypatch.setattr(episode_run, "teardown_run", teardowns) + yield teardowns + for gate in GATES.values(): + gate.set() + if teardowns.hold is not None: + teardowns.hold.set() + for episode in list(episode_run._live.values()): + await episode.abandon("the test ended") + + +ITEMS = {"id": "items", "type": "episode_test_env"} +ALICE = {"id": "alice", "type": "deploy_agent", "agent_name": "alice", "env_ids": ["items"], + "role": "alice@example.com", "system_prompt": "You keep the items."} +TURN = {"id": "turn", "type": "prompt_agent", "agent_name": "alice", "prompt_id": "turn", "prompt": "Add a crate.", + "timeout_seconds": 5} +REPLIES = {"id": "replies", "type": "episode_test_replies", "verifier_id": "replies"} + + +def _spec(*steps, **kwargs): + task = Task.from_dict({"id": "dock", "version": 4, "steps": [dict(step) for step in steps]}) + return EpisodeSpec.from_task(task, **kwargs) + + +async def _ended(episode): + await _until(lambda: Task.get_instance(episode.id).status != "running") + + +async def _until(holds): + async with asyncio.timeout(5): + while not holds(): + await asyncio.sleep(0.01) + + +@pytest.mark.asyncio +async def test_the_run_waits_at_an_outside_turn_and_carries_on_once_it_ends(teardowns): + episode = await _spec(ITEMS, ALICE, TURN, REPLIES).start() + assert Task.get_instance(episode.id).status == "running" + + turn = await episode.next_turn() + assert (turn.agent, turn.step_id, turn.prompt, turn.system_prompt, turn.role) == ( + "alice", "turn", "Add a crate.", "You keep the items.", "alice@example.com") + assert (turn.mcp_urls, turn.headers) == ({"items": "http://127.0.0.1:20031/items/mcp"}, {"items": {}}) + + await turn.end("Added a crate.") + assert await episode.result() == {"replies": 1.0} + assert await episode.next_turn() is None + assert teardowns.finished == [episode.id] + record = Task.get_instance(episode.id) + assert record.status == "completed" + assert record.context["metadata"]["episode"] == {"id": episode.id, "trainer_agents": ["alice"]} + assert [(r["step_id"], r["prompt_text"], r["response"]) for r in record.context["prompt_responses"]] == [ + ("turn", "Add a crate.", "Added a crate.")] + assert record.context["deployed_agents"] == [] + + +@pytest.mark.asyncio +async def test_a_turn_carries_the_headers_each_url_needs(teardowns, monkeypatch): + monkeypatch.setattr(episode_run, "sandbox_request_headers_for_url", lambda url: {"X-Sandbox": url}) + episode = await _spec(ITEMS, ALICE, TURN, REPLIES).start() + + turn = await episode.next_turn() + + assert turn.headers == {"items": {"X-Sandbox": "http://127.0.0.1:20031/items/mcp"}} + + +ITEMS_MCP, PANTRY_MCP = "http://127.0.0.1:20031/items/mcp", "http://127.0.0.1:20031/pantry/mcp" +PANTRY = {"id": "pantry", "type": "episode_test_env", "env_id": "pantry"} +STOCKED = {**ALICE, "env_ids": ["items", "pantry"]} + + +def _tool_sessions(monkeypatch, offered): + sessions = _ToolSessions(offered) + monkeypatch.setattr(episode_run, "ToolSession", sessions) + monkeypatch.setattr(episode_run, "sandbox_request_headers_for_url", lambda url: {"X-Sandbox": url}) + return sessions + + +@pytest.mark.asyncio +async def test_a_turns_tools_call_each_env_as_its_role_with_its_headers(teardowns, monkeypatch): + sessions = _tool_sessions(monkeypatch, {ITEMS_MCP: ["items_add_item"], PANTRY_MCP: ["pantry_count"]}) + episode = await _spec(ITEMS, PANTRY, STOCKED, TURN, REPLIES).start() + turn = await episode.next_turn() + + async with turn.tools(timeout=30) as tools: + assert [tool.name for tool in tools.list()] == ["items_add_item", "pantry_count"] + assert (await tools.call("pantry_count")).text == f"pantry_count on {PANTRY_MCP}" + + assert sessions.opened == [(url, "alice@example.com", {"X-Sandbox": url}, 30) for url in (ITEMS_MCP, PANTRY_MCP)] + assert sessions.calls == [(PANTRY_MCP, "pantry_count", None)] + assert sorted(sessions.closed) == [ITEMS_MCP, PANTRY_MCP] + + +@pytest.mark.asyncio +async def test_a_tool_two_envs_offer_is_refused_as_the_tools_open_and_nothing_is_left_open(teardowns, monkeypatch): + sessions = _tool_sessions(monkeypatch, {ITEMS_MCP: ["count"], PANTRY_MCP: ["count"]}) + episode = await _spec(ITEMS, PANTRY, STOCKED, TURN, REPLIES).start() + turn = await episode.next_turn() + + with pytest.raises(EpisodeError, match="tool 'count' is offered by both env 'items' and env 'pantry'"): + async with turn.tools(): + pass + assert sorted(sessions.closed) == [ITEMS_MCP, PANTRY_MCP] + + +@pytest.mark.asyncio +async def test_a_tool_no_env_offers_is_refused(teardowns, monkeypatch): + _tool_sessions(monkeypatch, {ITEMS_MCP: ["items_add_item"]}) + episode = await _spec(ITEMS, ALICE, TURN, REPLIES).start() + turn = await episode.next_turn() + + async with turn.tools() as tools: + with pytest.raises(EpisodeError) as raised: + await tools.call("nope") + assert str(raised.value) == "no env of turn 'turn' offers tool 'nope'; they offer ['items_add_item']" + + +@pytest.mark.asyncio +async def test_two_outside_agents_turns_are_each_handed_out_once(teardowns): + bob = {**ALICE, "id": "bob", "agent_name": "bob", "role": None, "depends_on": ["items"]} + bob_turn = {**TURN, "id": "bob-turn", "agent_name": "bob", "prompt_id": "bob-turn", "depends_on": ["bob"]} + episode = await _spec( + ITEMS, {**ALICE, "depends_on": ["items"]}, bob, {**TURN, "depends_on": ["alice"]}, bob_turn, + {**REPLIES, "depends_on": ["turn", "bob-turn"]}, trainer_agents=["alice", "bob"], + ).start() + + first = await episode.next_turn() + second = await episode.next_turn() + assert {first.agent, second.agent} == {"alice", "bob"} + await first.end("one") + await second.end("two") + + assert await episode.result() == {"replies": 1.0} + + +@pytest.mark.asyncio +async def test_episodes_of_one_spec_run_side_by_side_each_with_its_own_record(teardowns): + spec = _spec(ITEMS, ALICE, TURN, REPLIES) + + async def play(n): + episode = await spec.start() + await (await episode.next_turn()).end(f"crate {n}") + return episode.id, await episode.result() + + played = dict(await asyncio.gather(*(play(n) for n in range(8)))) + + assert list(played.values()) == [{"replies": 1.0}] * 8 + assert sorted(teardowns.finished) == sorted(played) + assert sorted(Task.get_instance(id).context["prompt_responses"][0]["response"] for id in played) == [ + f"crate {n}" for n in range(8)] + + +@pytest.mark.asyncio +async def test_a_seed_reaches_the_turns_prompt_env_binding_and_role(teardowns): + episode = await _spec( + {**ITEMS, "env_id": ""}, + {**ALICE, "env_ids": [""], "role": "", "system_prompt": "You keep the ."}, + {**TURN, "prompt": "Add a crate to the ."}, REPLIES, + seed={"env": "pantry", "who": "bob@example.com"}, + ).start() + + turn = await episode.next_turn() + + assert (turn.prompt, turn.system_prompt) == ("Add a crate to the pantry.", "You keep the pantry.") + assert (turn.mcp_urls, turn.role) == ({"pantry": "http://127.0.0.1:20031/pantry/mcp"}, "bob@example.com") + + +@pytest.mark.asyncio +async def test_a_setup_failure_reaches_the_caller_waiting_for_a_turn(teardowns): + episode = await _spec({**ITEMS, "fail": True}, ALICE, TURN, REPLIES).start() + + with pytest.raises(RuntimeError, match="the env didn't come up"): + await episode.next_turn() + with pytest.raises(RuntimeError, match="the env didn't come up"): + await episode.result() + assert Task.get_instance(episode.id).status == "failed" + assert teardowns.finished == [episode.id] + + +@pytest.mark.asyncio +async def test_a_turn_not_taken_by_its_deadline_fails_the_run_and_is_never_handed_out(teardowns): + episode = await _spec(ITEMS, ALICE, {**TURN, "timeout_seconds": 0.05}, REPLIES).start() + await _ended(episode) + + with pytest.raises(TimeoutError, match="turn 'turn' of agent 'alice' wasn't ended within 0.05s"): + await episode.next_turn() + assert Task.get_instance(episode.id).status == "failed" + + +@pytest.mark.asyncio +async def test_a_turn_ended_after_its_deadline_changes_nothing(teardowns): + episode = await _spec(ITEMS, ALICE, {**TURN, "timeout_seconds": 0.5}, REPLIES).start() + turn = await episode.next_turn() + await _ended(episode) + + await turn.end("too late") + + with pytest.raises(TimeoutError, match="wasn't ended within 0.5s"): + await episode.result() + record = Task.get_instance(episode.id) + assert record.status == "failed" and "wasn't ended within 0.5s" in record.error + assert record.context["prompt_responses"] == [] + + +@pytest.mark.asyncio +async def test_a_tolerant_turn_past_its_deadline_is_scored_without_a_reply(teardowns): + tolerant = {**TURN, "timeout_seconds": 0.05, "fail_task_on_error": False} + episode = await _spec(ITEMS, ALICE, tolerant, REPLIES).start() + await _ended(episode) + + assert await episode.result() == {"replies": 0.0} + assert await episode.next_turn() is None + assert [f["step_id"] for f in Task.get_instance(episode.id).context["metadata"]["failed_steps"]] == ["turn"] + + +@pytest.mark.asyncio +async def test_the_result_is_refused_while_a_turn_waits_or_is_open(teardowns): + episode = await _spec(ITEMS, ALICE, TURN, REPLIES).start() + + with pytest.raises(EpisodeError, match="turn 'turn' of agent 'alice' isn't ended, so the run can't finish"): + await episode.result() + turn = await episode.next_turn() + with pytest.raises(EpisodeError, match="isn't ended"): + await episode.result() + await turn.end("done") + assert await episode.result() == {"replies": 1.0} + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("reply", "score"), [("Added a crate.", 1.0), (None, 0.0)]) +async def test_the_reply_is_graded_as_a_prompt_agents_would_be(teardowns, reply, score): + graded = {"id": "graded", "type": "agent_prompt_response_verifier", "prompt_id": "turn", "verifier_id": "graded", + "criteria": [{"type": "response_contains", "needles": ["crate"]}]} + episode = await _spec(ITEMS, ALICE, TURN, graded).start() + + await (await episode.next_turn()).end(reply) + + assert await episode.result() == {"graded": score} + + +@pytest.mark.asyncio +async def test_an_episode_abandoned_as_it_starts_is_recorded_cancelled_with_the_reason(teardowns): + episode = await _spec(ITEMS, ALICE, TURN, REPLIES).start() + + await episode.abandon("the trainer moved on") + + record = Task.get_instance(episode.id) + assert (record.status, record.error) == ("cancelled", "the trainer moved on") + assert teardowns.finished == [episode.id] + for ask in (episode.next_turn, episode.result): + with pytest.raises(EpisodeError, match="episode abandoned: the trainer moved on"): + await ask() + + +@pytest.mark.asyncio +async def test_an_episode_abandoned_mid_turn_drops_the_turn(teardowns): + episode = await _spec(ITEMS, ALICE, TURN, REPLIES).start() + turn = await episode.next_turn() + + await episode.abandon("a tool call failed") + await turn.end("too late") + + record = Task.get_instance(episode.id) + assert (record.status, record.error, record.context["prompt_responses"]) == ( + "cancelled", "a tool call failed", []) + assert teardowns.finished == [episode.id] + + +@pytest.mark.asyncio +async def test_a_turn_waiting_when_the_episode_is_abandoned_is_never_handed_out(teardowns): + episode = await _spec(ITEMS, ALICE, TURN, REPLIES).start() + await _until(lambda: episode._waiting) + + abandoning = asyncio.create_task(episode.abandon("stop")) + await asyncio.sleep(0) + with pytest.raises(EpisodeError, match="episode abandoned: stop"): + await episode.next_turn() + await abandoning + + +@pytest.mark.asyncio +async def test_abandoning_during_teardown_waits_for_it_and_leaves_the_runs_outcome(teardowns): + teardowns.hold = asyncio.Event() + episode = await _spec(ITEMS, ALICE, TURN, REPLIES).start() + await (await episode.next_turn()).end("done") + await _until(lambda: teardowns.started) + + abandoning = asyncio.create_task(episode.abandon("too late")) + await asyncio.sleep(0.05) + assert not abandoning.done() + teardowns.hold.set() + await abandoning + + assert teardowns.finished == [episode.id] + assert await episode.result() == {"replies": 1.0} + assert Task.get_instance(episode.id).status == "completed" + + +@pytest.mark.asyncio +async def test_cancelling_a_caller_cancels_only_its_wait(teardowns): + gate = {"id": "before", "type": "episode_test_gate"} + episode = await _spec(ITEMS, ALICE, gate, TURN, {**gate, "id": "after"}, REPLIES).start() + + waiting = asyncio.create_task(episode.next_turn()) + await asyncio.sleep(0.01) + waiting.cancel() + with pytest.raises(asyncio.CancelledError): + await waiting + _gate("before").set() + await (await episode.next_turn()).end("done") + + scoring = asyncio.create_task(episode.result()) + await asyncio.sleep(0.01) + scoring.cancel() + with pytest.raises(asyncio.CancelledError): + await scoring + _gate("after").set() + assert await episode.result() == {"replies": 1.0} + + +@pytest.mark.asyncio +async def test_an_agent_named_by_the_seed_is_played_from_outside_in_the_run(teardowns): + episode = await _spec( + ITEMS, {**ALICE, "agent_name": ""}, {**TURN, "agent_name": ""}, REPLIES, seed={"who": "bob"}, + ).start() + + turn = await episode.next_turn() + + assert turn.agent == "bob" + await turn.end("done") + assert await episode.result() == {"replies": 1.0} + + +@pytest.mark.asyncio +async def test_a_turn_ended_twice_keeps_its_first_reply(teardowns): + episode = await _spec(ITEMS, ALICE, TURN, REPLIES).start() + turn = await episode.next_turn() + + await turn.end("first") + await turn.end("second") + await episode.result() + + assert [r["response"] for r in Task.get_instance(episode.id).context["prompt_responses"]] == ["first"] + + +@pytest.mark.asyncio +async def test_the_scores_are_every_verifiers_by_step_with_an_aggregate_beside_its_parts(teardowns): + def graded(step_id, needle): + return {"id": step_id, "type": "agent_prompt_response_verifier", "prompt_id": "turn", + "verifier_id": f"has-{needle}", "criteria": [{"type": "response_contains", "needles": [needle]}]} + overall = {"id": "overall", "type": "aggregate_verifiers", "verifier_ids": ["has-crate", "has-pallet"], + "verifier_id": "overall"} + episode = await _spec(ITEMS, ALICE, TURN, graded("crates", "crate"), graded("pallets", "pallet"), overall).start() + + await (await episode.next_turn()).end("Added a crate.") + + assert await episode.result() == {"crates": 1.0, "pallets": 0.0, "overall": 0.5} + + +@pytest.mark.asyncio +async def test_a_score_is_keyed_by_its_step_when_the_seed_names_its_verifier(teardowns): + episode = await _spec(ITEMS, ALICE, TURN, {**REPLIES, "verifier_id": "replies-"}, seed={"week": "3"}).start() + + await (await episode.next_turn()).end("done") + + assert await episode.result() == {"replies": 1.0} + + +@pytest.mark.asyncio +async def test_a_tolerated_verifier_that_recorded_no_score_is_missing_from_the_scores(teardowns): + unscored = {"id": "graded", "type": "agent_prompt_response_verifier", "prompt_id": "missing", + "verifier_id": "graded", "criteria": [{"type": "response_contains", "needles": ["crate"]}], + "fail_task_on_error": False} + episode = await _spec(ITEMS, ALICE, TURN, unscored, REPLIES).start() + await (await episode.next_turn()).end("done") + + assert await episode.result() == {"replies": 1.0} + + +BREAKS_WHILE_THE_TURN_IS_OPEN = ( + ITEMS, {**ALICE, "depends_on": ["items"]}, {**TURN, "depends_on": ["alice"]}, + {"id": "later", "type": "episode_test_gate", "depends_on": ["items"]}, + {"id": "broken", "type": "episode_test_env", "env_id": "crane", "fail": True, "depends_on": ["later"]}, + {**REPLIES, "depends_on": ["turn", "broken"]}, +) + + +@pytest.mark.asyncio +async def test_a_failure_elsewhere_in_the_task_surfaces_once_the_open_turn_ends(teardowns): + episode = await _spec(*BREAKS_WHILE_THE_TURN_IS_OPEN).start() + turn = await episode.next_turn() + _gate("later").set() + await _until(lambda: episode._context.metadata.get("failed_steps")) + + await turn.end("done") + + with pytest.raises(RuntimeError, match="the env didn't come up"): + await episode.result() + assert Task.get_instance(episode.id).status == "failed" + + +@pytest.mark.asyncio +async def test_abandoning_a_run_a_step_has_failed_keeps_its_error(teardowns): + episode = await _spec(*BREAKS_WHILE_THE_TURN_IS_OPEN).start() + await episode.next_turn() + _gate("later").set() + await _until(lambda: episode._context.metadata.get("failed_steps")) + + await episode.abandon("the trainer moved on") + + with pytest.raises(RuntimeError, match="the env didn't come up"): + await episode.next_turn() + record = Task.get_instance(episode.id) + assert (record.status, record.error) == ("failed", "the env didn't come up") + assert teardowns.finished == [episode.id] + + +@pytest.mark.asyncio +async def test_cancelling_the_caller_of_abandon_still_abandons_the_run(teardowns): + episode = await _spec(ITEMS, ALICE, TURN, REPLIES).start() + await episode.next_turn() + + abandoning = asyncio.create_task(episode.abandon("stop")) + await asyncio.sleep(0) + abandoning.cancel() + with pytest.raises(asyncio.CancelledError): + await abandoning + await _until(lambda: Task.get_instance(episode.id).status == "cancelled") + + assert (Task.get_instance(episode.id).status,Task.get_instance(episode.id).error) == ("cancelled", "stop") + assert teardowns.finished == [episode.id] + + +@pytest.mark.asyncio +async def test_a_cancelled_start_takes_its_run_down_with_it(teardowns): + spec = _spec(ITEMS, ALICE, TURN, REPLIES) + + starting = asyncio.create_task(spec.start()) + await asyncio.sleep(0) + starting.cancel() + with pytest.raises(asyncio.CancelledError): + await starting + + assert not episode_run._live + assert teardowns.finished == teardowns.started and len(teardowns.finished) == 1 diff --git a/tst/unit/task/episode/spec_test.py b/tst/unit/task/episode/spec_test.py new file mode 100644 index 00000000..14e4d4f2 --- /dev/null +++ b/tst/unit/task/episode/spec_test.py @@ -0,0 +1,268 @@ +"""What ``EpisodeSpec.from_task`` reads off a task (the agents a loop plays), what it refuses, and the copy of the task +it runs.""" + +import pytest + +from agent_env.task import EpisodeError, EpisodeSpec, Task +from agent_env.task.episode.rules import CHECKS, UNAFFECTED +from agent_env.task.episode.steps import _EpisodeAgentStep, _EpisodeTurnStep +from agent_env.task.task import substitute_seed_in_steps +from agent_env.task_step.registry import _builtin_registry +from agent_env.task_step.task_steps.deploy_agent import DeployAgentTaskStep +from agent_env.task_step.task_steps.prompt_agent import PromptAgentTaskStep + +ENV = {"id": "env", "type": "deploy_env", "env_id": "portsim"} +LOAD = {"id": "load", "type": "load_artifact", "env_id": "portsim", "artifact_id": "week-"} +AGENT = {"id": "agent", "type": "deploy_agent", "env_ids": ["portsim"], "role": "planner@", + "system_prompt": "You plan ."} +PLAY = {"id": "play", "type": "prompt_agent", "prompt_id": "play", "prompt": "Plan week .", + "depends_on": ["load"], "fail_task_on_error": False} +VERIFY = {"id": "verify", "type": "env_outcome_verifier", "env_id": "portsim", "file_artifact_id": "verify-py", + "verifier_id": "berth"} +BOB = {"id": "bob", "type": "deploy_agent", "agent_name": "bob", "a2a_agent_id": "echo"} +ASK_BOB = {"id": "ask-bob", "type": "prompt_agent", "agent_name": "bob", "prompt": "Check the plan."} +SEED = {"week": "3", "port": "Valencia"} + + +def _task(*steps): + return Task.from_dict({"id": "dock", "version": 4, "steps": [dict(step) for step in steps]}) + + +def _portsim(*extra, **replace): + steps = {step["id"]: step for step in (ENV, LOAD, AGENT, PLAY, VERIFY)} + for step_id, step in replace.items(): + steps[step_id] = step + return _task(*(step for step in steps.values() if step is not None), *extra) + + +def _refusal(task, trainer_agents=None): + with pytest.raises(EpisodeError) as raised: + EpisodeSpec.from_task(task, trainer_agents=trainer_agents) + return str(raised.value) + + +def test_a_task_with_one_prompted_agent_is_played_as_that_agent(): + assert EpisodeSpec.from_task(_portsim(), seed=SEED).trainer_agents == ("default-agent",) + + +def test_an_agent_named_by_the_seed_is_the_one_played_from_outside(): + spec = EpisodeSpec.from_task( + _portsim(agent={**AGENT, "agent_name": ""}, play={**PLAY, "agent_name": ""}), + seed={**SEED, "who": "planner"}, + ) + + assert spec.trainer_agents == ("planner",) + assert [step.type for step in spec._copy.steps] == [ + "deploy_env", "load_artifact", "episode_agent", "episode_turn", "env_outcome_verifier"] + + +def test_a_prompts_own_system_prompt_wins_over_the_agents(): + spec = EpisodeSpec.from_task(_portsim(play={**PLAY, "system_prompt": "Plan only the berths."})) + + assert spec._copy.steps[3].system_prompt == "Plan only the berths." + + +def test_each_prompt_of_a_trainer_agent_becomes_a_turn(): + again = {**PLAY, "id": "replan", "prompt_id": "replan", "prompt": "Replan.", "depends_on": ["play"]} + spec = EpisodeSpec.from_task(_portsim(again)) + + assert [step.id for step in spec._copy.steps if step.type == "episode_turn"] == ["play", "replan"] + + +def test_a_task_prompting_several_agents_needs_the_ones_to_play_named(): + task = _portsim(BOB, ASK_BOB) + + assert _refusal(task) == ( + "task 'dock' prompts ['default-agent', 'bob']; name the ones to play with trainer_agents=") + assert EpisodeSpec.from_task(task, trainer_agents=["bob"]).trainer_agents == ("bob",) + assert EpisodeSpec.from_task(task, trainer_agents=["bob", "default-agent", "bob"]).trainer_agents == ( + "bob", "default-agent") + + +@pytest.mark.parametrize(("trainer_agents", "message"), [ + (["carol"], "task 'dock' doesn't prompt ['carol']; it prompts ['default-agent']"), + ([], "trainer_agents= names no agent; task 'dock' prompts ['default-agent']"), +]) +def test_only_a_prompted_agent_can_be_played(trainer_agents, message): + assert _refusal(_portsim(), trainer_agents=trainer_agents) == message + + +def test_a_task_prompting_no_agent_has_nothing_to_play(): + assert _refusal(_task(ENV, VERIFY)) == "task 'dock' prompts no agent" + + +def test_the_copy_stands_in_for_the_outside_agents_steps_and_keeps_everything_else(): + task = _portsim(BOB, ASK_BOB) + copy = EpisodeSpec.from_task(task, trainer_agents=["default-agent"])._copy + + assert (copy.id, copy.version) == ("dock", 4) + assert [step.type for step in copy.steps] == [ + "deploy_env", "load_artifact", "episode_agent", "episode_turn", "env_outcome_verifier", "deploy_agent", + "prompt_agent", + ] + assert [copy.steps[i] is task.steps[i] for i in range(len(task.steps))] == [ + True, True, False, False, True, True, True] + for original, stand_in in zip(task.steps, copy.steps): + assert (stand_in.id, stand_in.version, stand_in.depends_on, stand_in.fail_task_on_error) == ( + original.id, original.version, original.depends_on, original.fail_task_on_error) + + +def test_the_stand_ins_carry_what_the_run_needs_unseeded_and_round_trip(): + agent, turn = EpisodeSpec.from_task(_portsim(), seed=SEED)._copy.steps[2:4] + + assert isinstance(agent, _EpisodeAgentStep) and isinstance(turn, _EpisodeTurnStep) + assert (agent.agent_name, agent.env_ids, agent.env_step_id, agent.role) == ( + "default-agent", ["portsim"], None, "planner@") + assert (turn.agent_name, turn.prompt_id, turn.prompt, turn.system_prompt, turn.timeout_seconds) == ( + "default-agent", "play", "Plan week .", "You plan .", 600) + for step in (agent, turn): + assert type(step).from_dict(step.to_dict()).to_dict() == step.to_dict() + + +@pytest.mark.parametrize( + ("timeout", "deadline"), [(30, 30), (None, PromptAgentTaskStep.DEFAULT_PROMPT_TIMEOUT_SECONDS)]) +def test_a_turns_deadline_is_its_prompts_timeout(timeout, deadline): + turn = EpisodeSpec.from_task(_portsim(play={**PLAY, "timeout_seconds": timeout}))._copy.steps[3] + + assert turn.timeout_seconds == deadline + + +def test_a_turn_timeout_replaces_every_outside_prompts_own(): + again = {**PLAY, "id": "replan", "prompt_id": "replan", "timeout_seconds": 60} + copy = EpisodeSpec.from_task(_portsim(again, play={**PLAY, "timeout_seconds": 7200}), turn_timeout=90)._copy + + assert [step.timeout_seconds for step in copy.steps if step.type == "episode_turn"] == [90, 90] + + +def test_an_inside_agent_keeps_its_own_steps(): + copy = EpisodeSpec.from_task(_portsim(BOB, ASK_BOB), trainer_agents=["bob"])._copy + + assert type(copy.steps[2]) is DeployAgentTaskStep and type(copy.steps[3]) is PromptAgentTaskStep + assert [step.type for step in copy.steps[5:]] == ["episode_agent", "episode_turn"] + + +AT_THE_AGENT = {"agent_name": "default-agent"} +REFUSED = { + "no deploy_agent": ({}, {"agent": None}, "agent 'default-agent' is played from outside, so it needs exactly one " + "deploy_agent step, not 0"), + "two deploy_agents": (({**AGENT, "id": "agent-2"},), {}, "needs exactly one deploy_agent step, not 2"), + "deploy_human_agent": (({"id": "human", "type": "deploy_human_agent", **AT_THE_AGENT},), {"agent": None}, + "step 'human': deploy_human_agent deploys 'default-agent', which is played from outside"), + "install_agent": (({"id": "install", "type": "install_agent", "sandbox_name": "box", "a2a_agent_id": "x", + **AT_THE_AGENT},), {}, + "step 'install': install_agent deploys 'default-agent'"), + "parts": ({}, {"play": {**PLAY, "prompt": None, "parts": [{"kind": "text", "text": "Plan."}]}}, + "step 'play': an outside agent's prompt can't have parts, only a text prompt"), + "user simulator": ({}, {"play": {**PLAY, "max_conversation_turns": 2}}, + "step 'play': an outside agent's prompt can't have a user simulator " + "(max_conversation_turns > 1)"), + "output_format": ({}, {"play": {**PLAY, "output_format": {"type": "json_schema", "schema": {}}}}, + "step 'play': an outside agent's prompt can't have output_format"), + "agentenv_tools": ({}, {"play": {**PLAY, "agentenv_tools": ["complete_step"]}}, + "step 'play': an outside agent's prompt can't have agentenv_tools"), + "snapshot_config": ({}, {"play": {**PLAY, "snapshot_config": {"env_id": "portsim"}}}, + "step 'play': an outside agent's prompt can't have snapshot_config"), + "skills": ({}, {"agent": {**AGENT, "skills": [{"name": "s", "skill_md": "Plan berths first."}]}}, + "step 'agent': an outside agent can't have skills, instructions its player wouldn't get"), + "snapshot files": ({}, {"agent": {**AGENT, "agent_snapshot_files_artifact_id": "files"}}, + "step 'agent': an outside agent can't have agent_snapshot_files_artifact_id"), + "changelog url": ({}, {"agent": {**AGENT, "agent_changelog_object_url": "file:///changelog"}}, + "step 'agent': an outside agent can't have agent_changelog_object_url"), + "changelog": ({}, {"agent": {**AGENT, "enable_agent_changelog": True}}, + "step 'agent': an outside agent can't have enable_agent_changelog"), + "add_skills": (({"id": "s", "type": "add_skills", "skills": [{"name": "s", "description": "d", "body": "x"}]},), {}, + "step 's': add_skills needs 'default-agent' as a real agent, and it is played from outside"), + "register_agent_triggers": (({"id": "s", "type": "register_agent_triggers", "triggers": [], **AT_THE_AGENT},), {}, + "step 's': register_agent_triggers needs 'default-agent' as a real agent"), + "snapshot_agent_state": (({"id": "s", "type": "snapshot_agent_state", "artifact_id": "a", "prompt_id": "play"},), + {}, "step 's': snapshot_agent_state needs 'default-agent' as a real agent"), + "load_artifact": (({"id": "s", "type": "load_artifact", "artifact_id": "a", **AT_THE_AGENT},), {}, + "step 's': load_artifact needs 'default-agent' as a real agent"), + "verify_sandbox": (({"id": "s", "type": "verify_sandbox", "criteria": []},), {}, + "step 's': verify_sandbox needs 'default-agent' as a real agent"), + "run_code": (({"id": "s", "type": "run_code", "script_artifact_id": "s", "entrypoint": "main.py"},), {}, + "step 's': run_code needs 'default-agent' as a real agent"), + "collect from the agent": (({"id": "s", "type": "collect_artifacts", "artifact_paths": ["/out"]},), {}, + "step 's': collects from 'default-agent''s sandbox, and it is played from outside"), + "collect a manifest": (({"id": "s", "type": "collect_artifacts", "env_id": "portsim", "manifest_step_id": "play"},), + {}, "step 's': its manifest is outside prompt 'play'"), + "trigger executor": (({"id": "s", "type": "register_env_triggers", "env_id": "portsim", "triggers": [], + "executor_agent_name": "default-agent"},), {}, + "step 's': its executor 'default-agent' is played from outside"), + "peered": ((BOB, {"id": "s", "type": "peer_agents", "peerings": [ + {"source_agent_name": "bob", "peer_agent_names": ["default-agent"]}]}), {}, + "step 's': peers 'default-agent', which is played from outside"), + "rubric on the trajectory": (({"id": "s", "type": "rubrics_verifier", "prompt_id": "play", "criteria": [], + "use_agent_judge": False},), {}, + "step 's': grades an outside prompt's trajectory, which an outside player doesn't " + "hand back"), + "rubric judged in the agent": (({"id": "s", "type": "rubrics_verifier", "prompt_id": "other", "criteria": [], + "use_trajectory": False, **AT_THE_AGENT},), {}, + "step 's': judges in 'default-agent''s sandbox, and it is played from outside"), + "retry_config": ({}, {"verify": {**VERIFY, "retry_config": {"retry_from_step_id": "verify"}}}, + "step 'verify': retry_config can replay an outside turn"), + "an A2A validator": (({"id": "s", "type": "verify_a2a_agent_card", "a2a_agent_id": "x"},), {}, + "step 's': validates the task's deployed A2A agent, and an episode plays one from outside"), +} + + +@pytest.mark.parametrize(("extra", "replace", "message"), REFUSED.values(), ids=REFUSED.keys()) +def test_a_step_an_outside_player_cant_stand_behind_is_refused_naming_it(extra, replace, message): + refusal = _refusal(_portsim(*extra, **replace)) + + assert refusal.startswith("task 'dock' can't run as an episode: ") + assert message in refusal + + +def test_a_user_simulator_played_from_outside_is_refused(): + ask = {**ASK_BOB, "max_conversation_turns": 3, "user_agent_name": "default-agent"} + + assert "step 'ask-bob': its user simulator 'default-agent' is played from outside" in _refusal( + _portsim(BOB, ask), trainer_agents=["default-agent"]) + + +def test_every_refusal_is_reported_at_once(): + refusal = _refusal(_portsim(play={**PLAY, "agentenv_tools": ["complete_step"]}, + agent={**AGENT, "enable_agent_changelog": True})) + + assert refusal == ("task 'dock' can't run as an episode: step 'agent': an outside agent can't have " + "enable_agent_changelog; step 'play': an outside agent's prompt can't have agentenv_tools") + + +ALLOWED = { + "a reply verifier on an outside prompt": ( + {"verify": {"id": "graded", "type": "agent_prompt_response_verifier", "prompt_id": "play", + "criteria": [{"type": "response_contains", "needles": ["berths"]}]}},), + "a rubric of the reply alone": ( + {"verify": {"id": "rubric", "type": "rubrics_verifier", "prompt_id": "play", "criteria": [], + "use_trajectory": False, "use_agent_judge": False}},), + "collecting from an env": ({"collect": {"id": "collect", "type": "collect_artifacts", "env_id": "portsim", + "artifact_paths": ["/out"]}},), + "tearing down an outside agent": ({"down": {"id": "down", "type": "teardown_sandboxes", + "agent_names": ["default-agent"]}},), +} + + +@pytest.mark.parametrize("replace", [case[0] for case in ALLOWED.values()], ids=ALLOWED.keys()) +def test_steps_that_need_only_what_an_outside_player_leaves_are_allowed(replace): + EpisodeSpec.from_task(_portsim(**replace)) + +def test_every_built_in_step_type_is_checked_or_unaffected(): + registered, checked = set(_builtin_registry()), set(CHECKS) + + assert sorted(registered - checked - UNAFFECTED) == [] + assert sorted((checked | UNAFFECTED) - registered) == [] + assert sorted(checked & UNAFFECTED) == [] + + +def test_the_seed_pass_rebuilds_each_step_as_its_own_class(): + class _Unregistered(PromptAgentTaskStep): + pass + + steps = [*_portsim().steps, _Unregistered.from_dict({**PLAY, "id": "extra", "prompt": "Berth at ."})] + + seeded = substitute_seed_in_steps(steps, SEED) + + assert [type(step) for step in seeded] == [type(step) for step in steps] + assert (seeded[2].role, seeded[3].prompt, seeded[-1].prompt) == ( + "planner@Valencia", "Plan week 3.", "Berth at Valencia.") diff --git a/tst/unit/task_step/test_deploy_agent_live_endpoint.py b/tst/unit/task_step/test_deploy_agent_live_endpoint.py index d05b4645..00e56763 100644 --- a/tst/unit/task_step/test_deploy_agent_live_endpoint.py +++ b/tst/unit/task_step/test_deploy_agent_live_endpoint.py @@ -15,7 +15,7 @@ def _env_get_returns(env): @pytest.mark.parametrize("raw", ["https://mcp.example.com/mcp", " https://mcp.example.com/mcp\n"]) def test_env_with_a_live_mcp_url_resolves(raw): with _env_get_returns(SimpleNamespace(mcp_url=raw)): - assert deploy_agent._live_mcp_url("live-env") == "https://mcp.example.com/mcp" + assert deploy_agent.live_mcp_url("live-env") == "https://mcp.example.com/mcp" @pytest.mark.parametrize( @@ -33,4 +33,4 @@ def test_env_with_a_live_mcp_url_resolves(raw): ) def test_env_without_a_live_mcp_url_is_an_error(env): with _env_get_returns(env), pytest.raises(RuntimeError, match="not in context.deployed_envs"): - deploy_agent._live_mcp_url("deployable-env") + deploy_agent.live_mcp_url("deployable-env") diff --git a/tst/unit/task_step/test_parallel_deploy_env_binding.py b/tst/unit/task_step/test_parallel_deploy_env_binding.py index 5041db9a..5b95f743 100644 --- a/tst/unit/task_step/test_parallel_deploy_env_binding.py +++ b/tst/unit/task_step/test_parallel_deploy_env_binding.py @@ -12,7 +12,7 @@ from agent_env.env.env import DeployedGatewayEnv from agent_env.task_step.context import TaskStepContext -from agent_env.task_step.task_steps.deploy_agent import DeployAgentTaskStep +from agent_env.task_step.task_steps.deploy_agent import DeployAgentTaskStep, bound_deployment from agent_env.task_step.task_steps.deploy_env import DeployEnvTaskStep ENV_ID = "multi-env" @@ -39,13 +39,9 @@ async def _deploy(**_kwargs): return env -def _agent_bound_to(step_id: str | None) -> DeployAgentTaskStep: - return DeployAgentTaskStep( - id=f"ag-{step_id or 'unbound'}", - version=None, - env_ids=[ENV_ID], - env_step_id=step_id, - ) +def _bound(context: TaskStepContext, env_id: str, env_step_id: str | None): + """The deployment an agent bound to ``env_step_id`` gets, as deploy_agent resolves it.""" + return bound_deployment(context, env_id, env_step_id=env_step_id, step_id=f"ag-{env_step_id or 'unbound'}") async def _deploy_two_in_parallel() -> TaskStepContext: @@ -80,8 +76,8 @@ async def test_each_agent_binds_to_the_env_it_depends_on(): context = await _deploy_two_in_parallel() by_step = {(d.metadata or {}).get("deploy_step_id"): d for d in context.deployed_envs} - bound_a = _agent_bound_to("env-a")._deployed_env(context, ENV_ID) - bound_b = _agent_bound_to("env-b")._deployed_env(context, ENV_ID) + bound_a = _bound(context, ENV_ID, "env-a") + bound_b = _bound(context, ENV_ID, "env-b") assert bound_a.mcp_url == by_step["env-a"].mcp_url assert bound_b.mcp_url == by_step["env-b"].mcp_url @@ -96,7 +92,7 @@ async def test_agent_without_env_step_id_refuses_to_guess(): context = await _deploy_two_in_parallel() with pytest.raises(RuntimeError, match="2 deployments of env"): - _agent_bound_to(None)._deployed_env(context, ENV_ID) + _bound(context, ENV_ID, None) @pytest.mark.asyncio @@ -108,7 +104,7 @@ async def test_single_deployment_still_resolves_without_env_step_ids(): Env.get.return_value = _fake_env_deploying_to("https://only.modal.host/mcp", "sb-ONLY") await step.execute(context) - bound = _agent_bound_to(None)._deployed_env(context, ENV_ID) + bound = _bound(context, ENV_ID, None) assert bound.sandbox_id == "sb-ONLY" @@ -134,7 +130,7 @@ async def test_env_step_id_naming_no_deployment_is_an_error(): context = await _deploy_two_in_parallel() with pytest.raises(RuntimeError, match="names no deployment"): - _agent_bound_to("env-typo")._deployed_env(context, ENV_ID) + _bound(context, ENV_ID, "env-typo") @pytest.mark.asyncio @@ -150,7 +146,7 @@ async def test_stale_env_step_id_errors_even_with_a_single_deployment(): await step.execute(context) with pytest.raises(RuntimeError, match="names no deployment"): - _agent_bound_to("env-typo")._deployed_env(context, ENV_ID) + _bound(context, ENV_ID, "env-typo") @pytest.mark.asyncio @@ -161,7 +157,7 @@ async def test_env_never_deployed_here_still_resolves_to_none(): """ context = await _deploy_two_in_parallel() - assert _agent_bound_to("env-a")._deployed_env(context, "oc-halt-ask-human") is None + assert _bound(context, "oc-halt-ask-human", "env-a") is None def test_deploy_agent_env_step_id_round_trips():