From ed9cbfe0a5e591fdc2eb8b53256b889a1dab34a5 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Mon, 12 Oct 2026 16:23:00 -0700 Subject: [PATCH 01/17] feat(task): EpisodeSpec.from_task reads a task as an episode Which agents are played from outside, their turns (prompt and system prompt, seeded) and roles, the reward (the one verifier no aggregate_verifiers step lists), and the copy of the task that runs, in which the outside agents' deploy_agent and prompt_agent steps become episode_agent and episode_turn stand-ins. Refuses, naming each step, what an outside player can't stand behind. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/task/__init__.py | 5 + src/agent_env/task/episode.py | 252 +++++++++++++++++++++++++++++++++ tst/unit/task/episode_test.py | 244 +++++++++++++++++++++++++++++++ 3 files changed, 501 insertions(+) create mode 100644 src/agent_env/task/episode.py create mode 100644 tst/unit/task/episode_test.py diff --git a/src/agent_env/task/__init__.py b/src/agent_env/task/__init__.py index 47743161..bc39df0d 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 AgentSpec, EpisodeError, EpisodeSpec, TurnSpec 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", + "AgentSpec", + "EpisodeError", + "EpisodeSpec", + "TurnSpec", "TaskStore", "TaskQuery", "get_task_store", diff --git a/src/agent_env/task/episode.py b/src/agent_env/task/episode.py new file mode 100644 index 00000000..b433fd42 --- /dev/null +++ b/src/agent_env/task/episode.py @@ -0,0 +1,252 @@ +"""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 the reward is the score ``agent-env run`` would give. +""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from typing import Any, ClassVar, Optional + +from agent_env.entity_refs import EntityRef +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 + +_AGENT_DEPLOYERS = ("install_agent", "deploy_human_agent") +_AGENT_TARGETING = ( + "add_skills", "register_agent_triggers", "snapshot_agent_state", "load_artifact", "verify_sandbox", "run_code", +) + + +class EpisodeError(ValueError): + """A task that can't run as an episode, or an episode asked for what it can't give.""" + + +@dataclass(frozen=True) +class TurnSpec: + step_id: str + prompt: str + system_prompt: str | None + + +@dataclass(frozen=True) +class AgentSpec: + role: str | None + turns: tuple[TurnSpec, ...] + + +class EpisodeSpec: + """A task read as an episode: the agents played from outside, their turns, and the copy of the task that runs.""" + + def __init__( + self, *, copy: Task, agents: dict[str, AgentSpec], reward_verifier_id: str, seed: dict[str, str] | None, + ): + self._copy = copy + self.agents = agents + self._reward_verifier_id = reward_verifier_id + self._seed = seed + + @classmethod + def from_task( + cls, task: Task, *, agents: Sequence[str] | None = None, seed: Mapping[str, str] | None = None, + ) -> EpisodeSpec: + seed = dict(seed) if seed else None + seeded = substitute_seed_in_steps(task.steps, seed or {}) + outside = _outside_agents(task, seeded, agents) + problems = _problems(seeded, outside) + verifiers = _top_level_verifiers(seeded) + if len(verifiers) != 1: + problems.append( + f"the reward is the score of the one verifier no aggregate_verifiers step lists, and the task has " + f"{len(verifiers)}" + (f": {[step.id for step in verifiers]}" if verifiers else "") + ) + if problems: + raise EpisodeError(f"task {task.id!r} can't run as an episode: " + "; ".join(problems)) + deploys = { + step.agent_name: step for step in task.steps if step.type == "deploy_agent" and step.agent_name in outside + } + return cls( + copy=Task(id=task.id, version=task.version, steps=[_stand_in(step, deploys) for step in task.steps]), + agents={name: _agent_spec(seeded, name) for name in outside}, + reward_verifier_id=verifiers[0].verifier_id, + seed=seed, + ) + + +def _outside_agents(task: Task, steps: list[TaskStep], agents: Sequence[str] | None) -> list[str]: + prompted = list(dict.fromkeys(step.agent_name for step in steps if step.type == "prompt_agent")) + if agents is None: + if len(prompted) != 1: + raise EpisodeError( + f"task {task.id!r} prompts {prompted}; name the ones to play with agents=" if prompted + else f"task {task.id!r} prompts no agent" + ) + return prompted + outside = list(dict.fromkeys(agents)) + if not outside: + raise EpisodeError(f"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 _problems(steps: list[TaskStep], outside: list[str]) -> list[str]: + """Why each step, and each outside agent, can't be run with the agent played from outside.""" + problems = [] + for name in outside: + deploys = [step for step in steps if step.type == "deploy_agent" and step.agent_name == name] + if len(deploys) != 1: + problems.append(f"agent {name!r} is played from outside, so it needs exactly one deploy_agent step, " + f"not {len(deploys)}") + outside_prompts = { + step.id: step.prompt_id for step in steps if step.type == "prompt_agent" and step.agent_name in outside + } + for step in steps: + problems.extend(f"step {step.id!r}: {why}" for why in _step_problems(step, outside, outside_prompts)) + return problems + + +def _step_problems(step: TaskStep, outside: list[str], outside_prompts: dict[str, str]) -> list[str]: + why = [] + if step.retry_config is not None: + why.append("retry_config can replay an outside turn") + agent = getattr(step, "agent_name", None) + if step.type == "prompt_agent" and agent in outside: + why.extend( + 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 + ) + elif step.type == "prompt_agent" and step.max_conversation_turns > 1 and step.user_agent_name in outside: + why.append(f"its user simulator {step.user_agent_name!r} is played from outside") + elif step.type == "deploy_agent" and agent in outside: + why.extend( + 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 + ) + elif step.type in _AGENT_DEPLOYERS and agent in outside: + why.append(f"{step.type} deploys {agent!r}, which is played from outside") + elif step.type in _AGENT_TARGETING and agent in outside: + why.append(f"{step.type} needs {agent!r} as a real agent, and it is played from outside") + elif step.type == "collect_artifacts": + if agent in outside and not (step.env_id or step.container_name or step.sandbox_name): + why.append(f"collects from {agent!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}") + elif step.type == "register_env_triggers" and step.executor_agent_name in outside: + why.append(f"its executor {step.executor_agent_name!r} is played from outside") + elif step.type == "peer_agents": + why.extend( + f"peers {name!r}, which is played from outside" for name in dict.fromkeys( + name for peering in step.peerings for name in (peering.source_agent_name, *peering.peer_agent_names) + ) if name in outside + ) + elif step.type == "rubrics_verifier": + 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 agent in outside: + why.append(f"judges in {agent!r}'s sandbox, and it is played from outside") + return why + + +def _top_level_verifiers(steps: list[TaskStep]) -> list[TaskStep]: + aggregated = {vid for step in steps if step.type == "aggregate_verifiers" for vid in step.verifier_ids} + return [step for step in steps if getattr(step, "verifier_id", None) and step.verifier_id not in aggregated] + + +def _agent_spec(steps: list[TaskStep], name: str) -> AgentSpec: + deploy = next(step for step in steps if step.type == "deploy_agent" and step.agent_name == name) + return AgentSpec(role=deploy.role, turns=tuple( + TurnSpec(step_id=step.id, prompt=step.prompt, system_prompt=step.system_prompt or deploy.system_prompt) + for step in steps if step.type == "prompt_agent" and step.agent_name == name + )) + + +def _stand_in(step: TaskStep, deploys: dict[str, TaskStep]) -> TaskStep: + """The step the episode's copy of the task runs in place of ``step``; ``deploys`` are the outside agents'.""" + if step.type == "deploy_agent" and step.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 step.type == "prompt_agent" and step.agent_name in deploys: + timeout = step.timeout_seconds + 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[step.agent_name].system_prompt, + timeout_seconds=PromptAgentTaskStep.DEFAULT_PROMPT_TIMEOUT_SECONDS if timeout is None else timeout, + ) + return step + + +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: + raise NotImplementedError + + +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: int, 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: + raise NotImplementedError diff --git a/tst/unit/task/episode_test.py b/tst/unit/task/episode_test.py new file mode 100644 index 00000000..92b93961 --- /dev/null +++ b/tst/unit/task/episode_test.py @@ -0,0 +1,244 @@ +"""An episode: a run of a task with some of its agents played from outside. What ``EpisodeSpec.from_task`` reads off +a task (the outside agents, their turns, the reward), what it refuses, and the copy of the task it runs.""" + +import pytest + +from agent_env.task import AgentSpec, EpisodeError, EpisodeSpec, Task, TurnSpec +from agent_env.task.episode import _EpisodeAgentStep, _EpisodeTurnStep +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, agents=None): + with pytest.raises(EpisodeError) as raised: + EpisodeSpec.from_task(task, agents=agents) + return str(raised.value) + + +def test_a_task_with_one_prompted_agent_is_played_as_that_agent_with_its_turns_seeded(): + spec = EpisodeSpec.from_task(_portsim(), seed=SEED) + + assert spec.agents == {"default-agent": AgentSpec(role="planner@Valencia", turns=( + TurnSpec(step_id="play", prompt="Plan week 3.", system_prompt="You plan Valencia."), + ))} + assert spec._reward_verifier_id == "berth" + + +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.agents["default-agent"].turns[0].system_prompt == "Plan only the berths." + + +def test_each_prompt_of_an_outside_agent_is_a_turn_in_task_order(): + again = {**PLAY, "id": "replan", "prompt_id": "replan", "prompt": "Replan.", "depends_on": ["play"]} + spec = EpisodeSpec.from_task(_portsim(again)) + + assert [turn.step_id for turn in spec.agents["default-agent"].turns] == ["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 agents=" + assert list(EpisodeSpec.from_task(task, agents=["bob"]).agents) == ["bob"] + assert list(EpisodeSpec.from_task(task, agents=["bob", "default-agent"]).agents) == ["bob", "default-agent"] + + +@pytest.mark.parametrize(("agents", "message"), [ + (["carol"], "task 'dock' doesn't prompt ['carol']; it prompts ['default-agent']"), + ([], "agents= names no agent; task 'dock' prompts ['default-agent']"), +]) +def test_only_a_prompted_agent_can_be_played(agents, message): + assert _refusal(_portsim(), agents=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, 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_an_inside_agent_keeps_its_own_steps(): + copy = EpisodeSpec.from_task(_portsim(BOB, ASK_BOB), 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"] + + +def test_the_reward_is_the_verifier_no_aggregate_lists(): + graded = {"id": "graded", "type": "agent_prompt_response_verifier", "prompt_id": "play", "verifier_id": "reply", + "criteria": [{"id": "c", "expected": "berths"}]} + overall = {"id": "overall", "type": "aggregate_verifiers", "verifier_ids": ["berth", "reply"], + "verifier_id": "overall"} + + assert EpisodeSpec.from_task(_portsim(graded, overall))._reward_verifier_id == "overall" + + +@pytest.mark.parametrize(("extra", "replace", "count"), [ + ((), {"verify": None}, "0"), + (({**VERIFY, "id": "verify-2", "verifier_id": "crane"},), {}, "2: ['verify', 'verify-2']"), +]) +def test_a_task_without_exactly_one_top_level_verifier_has_no_reward(extra, replace, count): + assert _refusal(_portsim(*extra, **replace)) == ( + "task 'dock' can't run as an episode: the reward is the score of the one verifier no aggregate_verifiers " + f"step lists, and the task has {count}") + + +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"), +} + + +@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), 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": [{"id": "c", "expected": "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)) From 9aaf645906fcddcf33fd8427648915728dcb171b Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Mon, 12 Oct 2026 16:43:07 -0700 Subject: [PATCH 02/17] feat(task): run an episode, waiting at each outside turn EpisodeSpec.start() runs one Task.run of the episode's copy of the task and returns once its record exists. The run waits at each outside turn: next_turn() hands it out with the agent's role, env MCP URLs and the headers each needs, and turn.end() records the reply as prompt_agent would. result() is the reward verifier's score once the run and its teardown are over; abandon() cancels the run, waits for teardown and records it cancelled. A caller's cancellation never reaches the run. The seed pass rebuilds each step with its own class, so the stand-ins take a seed, and deploy_agent's env binding rule is a module function both deploy_agent and episode_agent call. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/task/__init__.py | 4 +- src/agent_env/task/episode.py | 180 +++++++- src/agent_env/task/task.py | 4 +- .../task_step/task_steps/deploy_agent.py | 55 +-- tst/unit/task/episode_test.py | 398 +++++++++++++++++- .../test_parallel_deploy_env_binding.py | 16 +- 6 files changed, 613 insertions(+), 44 deletions(-) diff --git a/src/agent_env/task/__init__.py b/src/agent_env/task/__init__.py index bc39df0d..0cce40a4 100644 --- a/src/agent_env/task/__init__.py +++ b/src/agent_env/task/__init__.py @@ -1,7 +1,7 @@ """Task module for AgentEnv.""" from .task import Task -from .episode import AgentSpec, EpisodeError, EpisodeSpec, TurnSpec +from .episode import AgentSpec, Episode, EpisodeError, EpisodeSpec, Turn, TurnSpec 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, @@ -11,8 +11,10 @@ __all__ = [ "Task", "AgentSpec", + "Episode", "EpisodeError", "EpisodeSpec", + "Turn", "TurnSpec", "TaskStore", "TaskQuery", diff --git a/src/agent_env/task/episode.py b/src/agent_env/task/episode.py index b433fd42..ba8f0e4d 100644 --- a/src/agent_env/task/episode.py +++ b/src/agent_env/task/episode.py @@ -5,21 +5,33 @@ from __future__ import annotations +import asyncio +import logging +import uuid +from collections import deque from collections.abc import Mapping, Sequence from dataclasses import dataclass from typing import Any, ClassVar, Optional from agent_env.entity_refs import EntityRef +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, substitute_seed_in_steps -from agent_env.task_step.context import TaskStepContext +from agent_env.task.teardown import teardown_run +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 _live_mcp_url, bound_deployment from agent_env.task_step.task_steps.prompt_agent import PromptAgentTaskStep +logger = logging.getLogger(__name__) + _AGENT_DEPLOYERS = ("install_agent", "deploy_human_agent") _AGENT_TARGETING = ( "add_skills", "register_agent_triggers", "snapshot_agent_state", "load_artifact", "verify_sandbox", "run_code", ) +_live: dict[str, Episode] = {} + class EpisodeError(ValueError): """A task that can't run as an episode, or an episode asked for what it can't give.""" @@ -75,6 +87,158 @@ def from_task( seed=seed, ) + async def start(self, *, run_group_id: str | None = None) -> 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, "agents": list(self.agents)}} + if self._seed: + metadata["seed"] = dict(self._seed) + if run_group_id is not None: + metadata["run_group_id"] = run_group_id + episode = Episode(self, episode_id, TaskStepContext(metadata=metadata)) + _live[episode_id] = episode + running = asyncio.Event() + episode._run = asyncio.create_task(episode._drive(running)) + await running.wait() + return episode + + +class Turn: + """An outside agent's turn: what it is asked, the envs it acts in and the headers each needs, until ``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() + + 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 not self._reply.done(): + self._reply.set_result(response or "") + + +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, spec: EpisodeSpec, id: str, context: TaskStepContext): + self.id = id + self._spec = spec + 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.Event() + 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: + self._changed.clear() + while self._waiting and self._abandoned is None: + turn = self._waiting.popleft() + if not turn._reply.done(): + return turn + if self._run.done(): + self._raise_unless_completed() + return None + await self._wait() + + async def result(self) -> float: + """The reward, once the run and its teardown are over: the score of the task's one top-level verifier.""" + while not self._run.done(): + self._changed.clear() + if open_turn := next((turn for turn in self._turns if not turn._reply.done()), 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() + self._raise_unless_completed() + verifier_id = self._spec._reward_verifier_id + score = ((self._context.metadata.get("verifications") or {}).get(verifier_id) or {}).get("score") + if not isinstance(score, (int, float)): + raise EpisodeError(f"verifier {verifier_id!r} left no score") + return float(score) + + async def abandon(self, reason: str) -> None: + """Stop the run and wait for its teardown; it is recorded cancelled with ``reason``. 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): + self._abandoned = reason + self._run.cancel(msg=reason) + await asyncio.wait((self._run,)) + + async def _drive(self, running: asyncio.Event) -> None: + running.set() + try: + await self._spec._copy.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) -> None: + """Until a turn arrives or the run ends.""" + changed = asyncio.ensure_future(self._changed.wait()) + try: + await asyncio.wait((self._run, changed), return_when=asyncio.FIRST_COMPLETED) + finally: + changed.cancel() + + 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 + + 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() + try: + return await asyncio.wait_for(turn._reply, 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: + episode = _live.get((context.metadata.get("episode") or {}).get("id")) + if episode is None: + raise RuntimeError("episode_agent and episode_turn steps run only in the run of an episode") + return episode + def _outside_agents(task: Task, steps: list[TaskStep], agents: Sequence[str] | None) -> list[str]: prompted = list(dict.fromkeys(step.agent_name for step in steps if step.type == "prompt_agent")) @@ -217,7 +381,12 @@ def from_dict(cls, data: dict[str, Any]) -> _EpisodeAgentStep: env_step_id=data["env_step_id"], role=data["role"]) async def execute(self, context: TaskStepContext) -> TaskStepContext: - raise NotImplementedError + mcp_urls = {} + for env_id in self.env_ids: + deployed = bound_deployment(context, env_id, self) + mcp_urls[env_id] = deployed.mcp_url if deployed is not None else _live_mcp_url(env_id) + _episode_of(context)._bindings[self.agent_name] = (self.role, mcp_urls) + return context class _EpisodeTurnStep(TaskStep): @@ -249,4 +418,9 @@ def from_dict(cls, data: dict[str, Any]) -> _EpisodeTurnStep: timeout_seconds=data["timeout_seconds"]) async def execute(self, context: TaskStepContext) -> TaskStepContext: - raise NotImplementedError + 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/task.py b/src/agent_env/task/task.py index 30c54cd2..ad7fb5a7 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,11 +110,10 @@ 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 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..f39a3aea 100644 --- a/src/agent_env/task_step/task_steps/deploy_agent.py +++ b/src/agent_env/task_step/task_steps/deploy_agent.py @@ -100,6 +100,34 @@ 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, step: TaskStep) -> Optional[DeployedEnv]: + """The deployment of `env_id` the agent `step` deploys is bound to, by its `env_step_id`, or None if it isn't + deployed here.""" + env_step_id = step.env_step_id + 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 +304,7 @@ 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, self) 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: @@ -456,31 +484,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/unit/task/episode_test.py b/tst/unit/task/episode_test.py index 92b93961..972c0071 100644 --- a/tst/unit/task/episode_test.py +++ b/tst/unit/task/episode_test.py @@ -1,10 +1,20 @@ """An episode: a run of a task with some of its agents played from outside. What ``EpisodeSpec.from_task`` reads off -a task (the outside agents, their turns, the reward), what it refuses, and the copy of the task it runs.""" +a task (the outside agents, their turns, the reward), what it refuses, and the copy of the task it runs; then the run +itself: its turns, its reward, abandoning it, teardown and its record.""" + +import asyncio import pytest +import pytest_asyncio +import agent_env.task.episode as episode_module +from agent_env.config.runtime import Config +from agent_env.env.env import DeployedEnv from agent_env.task import AgentSpec, EpisodeError, EpisodeSpec, Task, TurnSpec from agent_env.task.episode import _EpisodeAgentStep, _EpisodeTurnStep +from agent_env.task.task import substitute_seed_in_steps +from agent_env.task.teardown import TeardownReport +from agent_env.task_step.task_step import TaskStep from agent_env.task_step.task_steps.deploy_agent import DeployAgentTaskStep from agent_env.task_step.task_steps.prompt_agent import PromptAgentTaskStep @@ -123,7 +133,7 @@ def test_an_inside_agent_keeps_its_own_steps(): def test_the_reward_is_the_verifier_no_aggregate_lists(): graded = {"id": "graded", "type": "agent_prompt_response_verifier", "prompt_id": "play", "verifier_id": "reply", - "criteria": [{"id": "c", "expected": "berths"}]} + "criteria": [{"type": "response_contains", "needles": ["berths"]}]} overall = {"id": "overall", "type": "aggregate_verifiers", "verifier_ids": ["berth", "reply"], "verifier_id": "overall"} @@ -228,7 +238,7 @@ def test_every_refusal_is_reported_at_once(): ALLOWED = { "a reply verifier on an outside prompt": ( {"verify": {"id": "graded", "type": "agent_prompt_response_verifier", "prompt_id": "play", - "criteria": [{"id": "c", "expected": "berths"}]}},), + "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}},), @@ -242,3 +252,385 @@ def test_every_refusal_is_reported_at_once(): @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)) + + +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() + + +@pytest.fixture +def teardowns(monkeypatch, local_stores): + 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_module, "teardown_run", teardowns) + return teardowns + + +@pytest_asyncio.fixture(autouse=True) +async def _no_episode_outlives_its_test(): + yield + for episode in list(episode_module._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): + return EpisodeSpec.from_task(_task(*steps), **kwargs) + + +async def _ended(episode): + await asyncio.wait((episode._run,)) + + +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() == 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, "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_module, "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"}} + + +@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"]}, 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() == 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()) == [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_agents_env_binding_and_role(teardowns): + episode = await _spec( + {**ITEMS, "env_id": ""}, {**ALICE, "env_ids": [""], "role": ""}, TURN, REPLIES, + seed={"env": "pantry", "who": "bob@example.com"}, + ).start() + + turn = await episode.next_turn() + + assert (turn.mcp_urls, turn.role) == ({"pantry": "http://127.0.0.1:20031/pantry/mcp"}, "bob@example.com") + + +def test_the_seed_pass_rebuilds_each_step_as_its_own_class(): + class _Unregistered(_Replies): + pass + + steps = [*_portsim().steps, _Unregistered(id="extra", version=None, verifier_id="")] + + 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].verifier_id) == ("planner@Valencia", "Plan week 3.", "Valencia") + + +@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.05}, 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.05s"): + await episode.result() + record = Task.get_instance(episode.id) + assert record.status == "failed" and "wasn't ended within 0.05s" 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): + episode = await _spec(ITEMS, ALICE, {**TURN, "timeout_seconds": 0.05, "fail_task_on_error": False}, REPLIES).start() + await _ended(episode) + + assert await episode.result() == 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() == 1.0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("reply", "reward"), [("Added a crate.", 1.0), (None, 0.0)]) +async def test_the_reply_is_graded_as_a_prompt_agents_would_be(teardowns, reply, reward): + 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() == reward + + +@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() == 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() == 1.0 + + +@pytest.mark.asyncio +async def test_a_run_group_id_groups_the_record(teardowns): + episode = await _spec(ITEMS, ALICE, TURN, REPLIES).start(run_group_id="grpo-step-7") + + await (await episode.next_turn()).end("done") + await episode.result() + + assert Task.get_instance(episode.id).context["metadata"]["run_group_id"] == "grpo-step-7" 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..883fbb5e 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" @@ -80,8 +80,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_deployment(context, ENV_ID, _agent_bound_to("env-a")) + bound_b = bound_deployment(context, ENV_ID, _agent_bound_to("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 +96,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_deployment(context, ENV_ID, _agent_bound_to(None)) @pytest.mark.asyncio @@ -108,7 +108,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_deployment(context, ENV_ID, _agent_bound_to(None)) assert bound.sandbox_id == "sb-ONLY" @@ -134,7 +134,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_deployment(context, ENV_ID, _agent_bound_to("env-typo")) @pytest.mark.asyncio @@ -150,7 +150,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_deployment(context, ENV_ID, _agent_bound_to("env-typo")) @pytest.mark.asyncio @@ -161,7 +161,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_deployment(context, "oc-halt-ask-human", _agent_bound_to("env-a")) is None def test_deploy_agent_env_step_id_round_trips(): From bc36efdc7e6d66aa22db2e7ad8a2cda688a7d8bc Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Mon, 12 Oct 2026 16:47:31 -0700 Subject: [PATCH 03/17] test(task): a task's agent played as an episode on local containers The items env deployed and loaded and the echo agent prompted around an outside agent, which plays its turn through ToolSession as its role and is never deployed. A played episode scores 1.0 and an unplayed one 0.0, side by side; an abandoned one is recorded cancelled; none leaves a container or a work folder. local_stack now gives each test its own state root and closes the grant server, as the other local fixtures do. Co-Authored-By: Claude Opus 5.5 --- .../env/envs/server_env_local_test.py | 85 ++++++++++++++++++- 1 file changed, 82 insertions(+), 3 deletions(-) diff --git a/tst/integration/env/envs/server_env_local_test.py b/tst/integration/env/envs/server_env_local_test.py index 69fa4c6c..2e415075 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,72 @@ 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 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, the reward + the verifier's score; and every episode, played, not played or abandoned, 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") + spec = EpisodeSpec.from_task(Task.from_dict({"id": f"crate-{uid}", "steps": [ + {"id": "env", "type": "deploy_env", "env_id": env.id}, + {"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": "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 crate."}, + {"id": "check", "type": "env_outcome_verifier", "env_id": env.id, "file_artifact_id": verify.id, + "verifier_id": "crate"}, + ]}), agents=["policy"]) + + async def play(episode, add): + turn = await episode.next_turn() + assert [agent["agent_name"] for agent in Task.get_instance(episode.id).context["deployed_agents"]] == ["echo"] + if add: + async with ToolSession(mcp_url=turn.mcp_urls[env.id], role=turn.role, headers=turn.headers[env.id]) as own: + await own.call_tool("items_add_item", {"item": "crate"}) + assert json.loads((await own.call_tool("items_whoami")).text)["role"] == "alice@example.com" + await turn.end("Added a crate." if add else None) + return await episode.result() + + played, unplayed = await spec.start(), await spec.start() + assert await asyncio.gather(play(played, True), play(unplayed, False)) == [1.0, 0.0] + record = Task.get_instance(played.id) + assert record.status == "completed" and record.context["metadata"]["episode"]["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() + await abandoned.abandon("test") + assert (Task.get_instance(abandoned.id).status, Task.get_instance(abandoned.id).error) == ("cancelled", "test") + + assert list((tmp_path / "sandboxes").glob("agent-env-*")) == [] + for episode in (played, unplayed, abandoned): + 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" From 2d61225a4ffa4b54ee3142db72ff1f315861fa09 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Mon, 12 Oct 2026 16:47:32 -0700 Subject: [PATCH 04/17] docs: running a task from your own loop as an episode Co-Authored-By: Claude Opus 5.5 --- AGENTS.md | 8 +++++++- README.md | 19 +++++++++++++++++++ 2 files changed, 26 insertions(+), 1 deletion(-) diff --git a/AGENTS.md b/AGENTS.md index be1458e6..ceabf735 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.py` 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.py`) is one `Task.run()` of a copy of the task in which + each outside agent's `deploy_agent` and `prompt_agent` steps 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..c85849ee 100644 --- a/README.md +++ b/README.md @@ -44,6 +44,25 @@ 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 agentenv_protocol import ToolSession +from agent_env.task import EpisodeSpec, Task + +spec = EpisodeSpec.from_task(Task.get("my-task")) +episode = await spec.start() +turn = await episode.next_turn() +async with ToolSession(mcp_url=turn.mcp_urls["my-env"], role=turn.role, headers=turn.headers["my-env"]) as tools: + ... # your model reads turn.prompt and calls tools.call_tool(name, arguments) +await turn.end(response="Done.") +reward = await episode.result() +``` + +The reward is the score of the task's one top-level verifier, `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. From ef4231bddc73e6c90164b6aeac65c957c909632b Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Mon, 12 Oct 2026 17:27:35 -0700 Subject: [PATCH 05/17] fix(task): what the review of the episode runtime found - An outside agent named by a seed placeholder is replaced too: the stand-ins are judged by each step as the seed makes it. Before, the real agent was deployed and prompted, and its score was the reward. - A cancelled start() cancels its run, so no episode is left running that nobody holds. - Abandoning a run a step has already failed ends its open turns, so the run's error and failed record stand instead of reading cancelled. - verify_a2a_* validator steps are refused: they check the task's one deployed agent, which an episode plays from outside. - A turn's wake-up is a fresh future per change, taken before the state is read, so a concurrent caller can't swallow it. Tests for each, a no-score reward, a failure elsewhere surfacing once the turn ends, a cancelled abandon() caller, and the integration test now checks for leftovers after each ending. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/task/episode.py | 69 ++++---- .../env/envs/server_env_local_test.py | 6 +- tst/unit/task/episode_test.py | 157 ++++++++++++++++-- 3 files changed, 185 insertions(+), 47 deletions(-) diff --git a/src/agent_env/task/episode.py b/src/agent_env/task/episode.py index ba8f0e4d..5e82abf6 100644 --- a/src/agent_env/task/episode.py +++ b/src/agent_env/task/episode.py @@ -78,10 +78,12 @@ def from_task( if problems: raise EpisodeError(f"task {task.id!r} can't run as an episode: " + "; ".join(problems)) deploys = { - step.agent_name: step for step in task.steps if step.type == "deploy_agent" and step.agent_name in outside + 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) for step, as_run in zip(task.steps, seeded)] return cls( - copy=Task(id=task.id, version=task.version, steps=[_stand_in(step, deploys) for step in task.steps]), + copy=Task(id=task.id, version=task.version, steps=steps), agents={name: _agent_spec(seeded, name) for name in outside}, reward_verifier_id=verifiers[0].verifier_id, seed=seed, @@ -100,7 +102,12 @@ async def start(self, *, run_group_id: str | None = None) -> Episode: _live[episode_id] = episode running = asyncio.Event() episode._run = asyncio.create_task(episode._drive(running)) - await running.wait() + try: + await running.wait() + except asyncio.CancelledError: + episode._abandoned = "start() was cancelled" + episode._run.cancel() + raise return episode @@ -141,7 +148,7 @@ def __init__(self, spec: EpisodeSpec, id: str, context: TaskStepContext): self._bindings: dict[str, tuple[str | None, dict[str, str]]] = {} self._waiting: deque[Turn] = deque() self._turns: set[Turn] = set() - self._changed = asyncio.Event() + self._changed: asyncio.Future[None] = asyncio.get_running_loop().create_future() self._tearing_down = False self._abandoned: str | None = None @@ -149,7 +156,7 @@ 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: - self._changed.clear() + changed = self._changed while self._waiting and self._abandoned is None: turn = self._waiting.popleft() if not turn._reply.done(): @@ -157,17 +164,17 @@ async def next_turn(self) -> Turn | None: if self._run.done(): self._raise_unless_completed() return None - await self._wait() + await self._wait(changed) async def result(self) -> float: """The reward, once the run and its teardown are over: the score of the task's one top-level verifier.""" while not self._run.done(): - self._changed.clear() + changed = self._changed if open_turn := next((turn for turn in self._turns if not turn._reply.done()), 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() + await self._wait(changed) self._raise_unless_completed() verifier_id = self._spec._reward_verifier_id score = ((self._context.metadata.get("verifications") or {}).get(verifier_id) or {}).get("score") @@ -176,11 +183,17 @@ async def result(self) -> float: return float(score) async def abandon(self, reason: str) -> None: - """Stop the run and wait for its teardown; it is recorded cancelled with ``reason``. Once the run has - ended, or while it tears down, this only waits.""" + """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): - self._abandoned = reason - self._run.cancel(msg=reason) + if any(failure.get("is_fatal") for failure in self._context.metadata.get("failed_steps") or []): + for turn in self._turns: + if not turn._reply.done(): + turn._reply.set_exception(EpisodeError(f"episode abandoned: {reason}")) + else: + self._abandoned = reason + self._run.cancel(msg=reason) await asyncio.wait((self._run,)) async def _drive(self, running: asyncio.Event) -> None: @@ -198,13 +211,9 @@ async def _drive(self, running: asyncio.Event) -> None: finally: del _live[self.id] - async def _wait(self) -> None: - """Until a turn arrives or the run ends.""" - changed = asyncio.ensure_future(self._changed.wait()) - try: - await asyncio.wait((self._run, changed), return_when=asyncio.FIRST_COMPLETED) - finally: - changed.cancel() + 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: @@ -222,7 +231,8 @@ async def _play(self, step: _EpisodeTurnStep) -> str: role=role, mcp_urls=dict(mcp_urls)) self._waiting.append(turn) self._turns.add(turn) - self._changed.set() + self._changed.set_result(None) + self._changed = asyncio.get_running_loop().create_future() try: return await asyncio.wait_for(turn._reply, step.timeout_seconds) except TimeoutError: @@ -234,10 +244,7 @@ async def _play(self, step: _EpisodeTurnStep) -> str: def _episode_of(context: TaskStepContext) -> Episode: - episode = _live.get((context.metadata.get("episode") or {}).get("id")) - if episode is None: - raise RuntimeError("episode_agent and episode_turn steps run only in the run of an episode") - return episode + return _live[context.metadata["episode"]["id"]] def _outside_agents(task: Task, steps: list[TaskStep], agents: Sequence[str] | None) -> list[str]: @@ -301,6 +308,8 @@ def _step_problems(step: TaskStep, outside: list[str], outside_prompts: dict[str ) elif step.type in _AGENT_DEPLOYERS and agent in outside: why.append(f"{step.type} deploys {agent!r}, which is played from outside") + elif step.type.startswith("verify_a2a_"): + why.append("validates the task's deployed A2A agent, and an episode plays one from outside") elif step.type in _AGENT_TARGETING and agent in outside: why.append(f"{step.type} needs {agent!r} as a real agent, and it is played from outside") elif step.type == "collect_artifacts": @@ -337,19 +346,21 @@ def _agent_spec(steps: list[TaskStep], name: str) -> AgentSpec: )) -def _stand_in(step: TaskStep, deploys: dict[str, TaskStep]) -> TaskStep: - """The step the episode's copy of the task runs in place of ``step``; ``deploys`` are the outside agents'.""" - if step.type == "deploy_agent" and step.agent_name in deploys: +def _stand_in(step: TaskStep, as_run: TaskStep, deploys: dict[str, TaskStep]) -> 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 step.type == "prompt_agent" and step.agent_name in deploys: + if as_run.type == "prompt_agent" and as_run.agent_name in deploys: timeout = step.timeout_seconds 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[step.agent_name].system_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/tst/integration/env/envs/server_env_local_test.py b/tst/integration/env/envs/server_env_local_test.py index 2e415075..bae78410 100644 --- a/tst/integration/env/envs/server_env_local_test.py +++ b/tst/integration/env/envs/server_env_local_test.py @@ -241,8 +241,10 @@ async def play(episode, add): 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)) == [1.0, 0.0] + assert list(sandboxes.glob("agent-env-*")) == [] record = Task.get_instance(played.id) assert record.status == "completed" and record.context["metadata"]["episode"]["agents"] == ["policy"] assert [agent["agent_name"] for agent in record.context["deployed_agents"]] == ["echo"] @@ -251,10 +253,10 @@ async def play(episode, add): 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((tmp_path / "sandboxes").glob("agent-env-*")) == [] + assert list(sandboxes.glob("agent-env-*")) == [] for episode in (played, unplayed, abandoned): context = Task.get_instance(episode.id).context sandbox_ids = [entry["sandbox_id"] for key in ("deployed_envs", "deployed_agents") diff --git a/tst/unit/task/episode_test.py b/tst/unit/task/episode_test.py index 972c0071..6cc5a779 100644 --- a/tst/unit/task/episode_test.py +++ b/tst/unit/task/episode_test.py @@ -57,6 +57,18 @@ def test_a_task_with_one_prompted_agent_is_played_as_that_agent_with_its_turns_s assert spec._reward_verifier_id == "berth" +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 list(spec.agents) == ["planner"] + assert [step.type for step in spec._copy.steps] == [ + "deploy_env", "load_artifact", "episode_agent", "episode_turn", "env_outcome_verifier"] + assert spec.agents["planner"].turns[0].system_prompt == "You plan Valencia." + + def test_a_prompts_own_system_prompt_wins_over_the_agents(): spec = EpisodeSpec.from_task(_portsim(play={**PLAY, "system_prompt": "Plan only the berths."})) @@ -99,7 +111,8 @@ def test_the_copy_stands_in_for_the_outside_agents_steps_and_keeps_everything_el "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] + 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) @@ -117,7 +130,8 @@ def test_the_stand_ins_carry_what_the_run_needs_unseeded_and_round_trip(): 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)]) +@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] @@ -163,7 +177,8 @@ def test_a_task_without_exactly_one_top_level_verifier_has_no_reward(extra, repl "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)"), + "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"]}}, @@ -209,6 +224,8 @@ def test_a_task_without_exactly_one_top_level_verifier_has_no_reward(extra, repl "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"), } @@ -344,19 +361,20 @@ async def __call__(self, context, **_): return TeardownReport() -@pytest.fixture -def teardowns(monkeypatch, local_stores): +@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_module, "teardown_run", teardowns) - return teardowns - - -@pytest_asyncio.fixture(autouse=True) -async def _no_episode_outlives_its_test(): - yield + yield teardowns + for gate in GATES.values(): + gate.set() + if teardowns.hold is not None: + teardowns.hold.set() for episode in list(episode_module._live.values()): await episode.abandon("the test ended") @@ -471,7 +489,8 @@ class _Unregistered(_Replies): 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].verifier_id) == ("planner@Valencia", "Plan week 3.", "Valencia") + assert (seeded[2].role, seeded[3].prompt, seeded[-1].verifier_id) == ( + "planner@Valencia", "Plan week 3.", "Valencia") @pytest.mark.asyncio @@ -498,22 +517,23 @@ async def test_a_turn_not_taken_by_its_deadline_fails_the_run_and_is_never_hande @pytest.mark.asyncio async def test_a_turn_ended_after_its_deadline_changes_nothing(teardowns): - episode = await _spec(ITEMS, ALICE, {**TURN, "timeout_seconds": 0.05}, REPLIES).start() + 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.05s"): + 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.05s" in record.error + 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): - episode = await _spec(ITEMS, ALICE, {**TURN, "timeout_seconds": 0.05, "fail_task_on_error": False}, REPLIES).start() + 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() == 0.0 @@ -634,3 +654,108 @@ async def test_a_run_group_id_groups_the_record(teardowns): await episode.result() assert Task.get_instance(episode.id).context["metadata"]["run_group_id"] == "grpo-step-7" + + +@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() == 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_a_reward_verifier_that_left_no_score_is_no_reward(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).start() + await (await episode.next_turn()).end("done") + + with pytest.raises(EpisodeError, match="verifier 'graded' left no score"): + await episode.result() + + +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 _ended(episode) + + 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 + await _until(lambda: not episode_module._live) + + assert teardowns.finished == teardowns.started and len(teardowns.finished) == 1 From ba18cb1686035b088f953d0780fc995ff873205e Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Mon, 12 Oct 2026 22:49:28 -0700 Subject: [PATCH 06/17] fix: draw record ids and a state root's first host port from secrets The instance ids of envs, tasks, A2A agents and env states, and the first host port a fresh state root hands out, were drawn from the global random. A training loop reseeds it, so ids repeated: a GRPO run on episodes lost 32 of 352 env-instance records to DuplicateKeyError, which registration logs and swallows. Fresh state roots started at the same port too, which the random start exists to prevent. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/a2a_agent/store.py | 4 ++-- src/agent_env/env/store.py | 4 ++-- src/agent_env/providers/env_state/store.py | 4 ++-- .../providers/sandbox_providers/local_sandbox.py | 4 ++-- src/agent_env/task/store.py | 4 ++-- tst/unit/env/store_test.py | 13 +++++++++++++ .../local_sandbox_host_ports_test.py | 14 +++++++++++++- 7 files changed, 36 insertions(+), 11 deletions(-) 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/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/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/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] = [] From d415329453e76f2de55342a57354442aabd223cb Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Mon, 12 Oct 2026 22:49:30 -0700 Subject: [PATCH 07/17] feat(task): an episode's turn deadline can be set for the loop playing it EpisodeSpec.from_task(turn_timeout=) replaces every outside turn's deadline, since a prompt's timeout_seconds budgets an agent's run: a training loop needs its own. EpisodeError is a RuntimeError, as most of what it reports happens at run time. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/task/episode.py | 13 ++++++++----- tst/unit/task/episode_test.py | 7 +++++++ 2 files changed, 15 insertions(+), 5 deletions(-) diff --git a/src/agent_env/task/episode.py b/src/agent_env/task/episode.py index 5e82abf6..9b21b230 100644 --- a/src/agent_env/task/episode.py +++ b/src/agent_env/task/episode.py @@ -33,7 +33,7 @@ _live: dict[str, Episode] = {} -class EpisodeError(ValueError): +class EpisodeError(RuntimeError): """A task that can't run as an episode, or an episode asked for what it can't give.""" @@ -64,7 +64,10 @@ def __init__( @classmethod def from_task( cls, task: Task, *, agents: Sequence[str] | None = None, seed: Mapping[str, str] | None = None, + turn_timeout: float | None = None, ) -> EpisodeSpec: + """``turn_timeout`` is every outside 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 = _outside_agents(task, seeded, agents) @@ -81,7 +84,7 @@ def from_task( 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) for step, as_run in zip(task.steps, seeded)] + 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), agents={name: _agent_spec(seeded, name) for name in outside}, @@ -346,7 +349,7 @@ def _agent_spec(steps: list[TaskStep], name: str) -> AgentSpec: )) -def _stand_in(step: TaskStep, as_run: TaskStep, deploys: dict[str, TaskStep]) -> TaskStep: +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.""" @@ -356,7 +359,7 @@ def _stand_in(step: TaskStep, as_run: TaskStep, deploys: dict[str, TaskStep]) -> 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 + 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, @@ -408,7 +411,7 @@ class _EpisodeTurnStep(TaskStep): def __init__( self, id: str, version: Optional[int], *, agent_name: str, prompt_id: str, prompt: str, - system_prompt: str | None, timeout_seconds: int, depends_on: Optional[list] = None, + 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) diff --git a/tst/unit/task/episode_test.py b/tst/unit/task/episode_test.py index 6cc5a779..7120f007 100644 --- a/tst/unit/task/episode_test.py +++ b/tst/unit/task/episode_test.py @@ -138,6 +138,13 @@ def test_a_turns_deadline_is_its_prompts_timeout(timeout, deadline): 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), agents=["bob"])._copy From 2ac79f8f9a00832314ba5c4febf0462a3545fdb1 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Mon, 12 Oct 2026 22:50:32 -0700 Subject: [PATCH 08/17] fix(task): a cancelled start() reports once its run is torn down The caller holds no handle on the run, so it couldn't tell when cleanup ended. The README's sample now plays every turn before reading the reward, and shows turn_timeout. Co-Authored-By: Claude Opus 5.5 --- README.md | 12 ++++++------ src/agent_env/task/episode.py | 1 + tst/unit/task/episode_test.py | 2 +- 3 files changed, 8 insertions(+), 7 deletions(-) diff --git a/README.md b/README.md index c85849ee..5f71ffad 100644 --- a/README.md +++ b/README.md @@ -52,16 +52,16 @@ An episode is one run of a task in which some of its agents are played from outs from agentenv_protocol import ToolSession from agent_env.task import EpisodeSpec, Task -spec = EpisodeSpec.from_task(Task.get("my-task")) +spec = EpisodeSpec.from_task(Task.get("my-task"), turn_timeout=900) episode = await spec.start() -turn = await episode.next_turn() -async with ToolSession(mcp_url=turn.mcp_urls["my-env"], role=turn.role, headers=turn.headers["my-env"]) as tools: - ... # your model reads turn.prompt and calls tools.call_tool(name, arguments) -await turn.end(response="Done.") +while (turn := await episode.next_turn()) is not None: + async with ToolSession(mcp_url=turn.mcp_urls["my-env"], role=turn.role, headers=turn.headers["my-env"]) as tools: + ... # your model reads turn.prompt and calls tools.call_tool(name, arguments) + await turn.end(response="Done.") reward = await episode.result() ``` -The reward is the score of the task's one top-level verifier, `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. +The reward is the score of the task's one top-level verifier, `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 diff --git a/src/agent_env/task/episode.py b/src/agent_env/task/episode.py index 9b21b230..fe2b3edc 100644 --- a/src/agent_env/task/episode.py +++ b/src/agent_env/task/episode.py @@ -110,6 +110,7 @@ async def start(self, *, run_group_id: str | None = None) -> Episode: except asyncio.CancelledError: episode._abandoned = "start() was cancelled" episode._run.cancel() + await asyncio.wait((episode._run,)) raise return episode diff --git a/tst/unit/task/episode_test.py b/tst/unit/task/episode_test.py index 7120f007..1a134a41 100644 --- a/tst/unit/task/episode_test.py +++ b/tst/unit/task/episode_test.py @@ -763,6 +763,6 @@ async def test_a_cancelled_start_takes_its_run_down_with_it(teardowns): starting.cancel() with pytest.raises(asyncio.CancelledError): await starting - await _until(lambda: not episode_module._live) + assert not episode_module._live assert teardowns.finished == teardowns.started and len(teardowns.finished) == 1 From c834356a759aee6ebdb302d586f041196a2e0ac3 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Tue, 13 Oct 2026 09:02:51 -0700 Subject: [PATCH 09/17] refactor(task): episode rules in their own module, each piece of episode state with one writer - The v1 rules move to task/episode_rules.py: CHECKS maps a step type to its check, UNAFFECTED lists the types an outside agent leaves alone, and a test holds the step registry to the two, so a new step type is classified before an episode runs it. The refusals and their messages are unchanged. - Episode takes its task and reward verifier at construction; _launch() owns starting and cancelling the run, _bind() the agent bindings, and only Turn touches its reply future. - episode_test.py splits into episode_spec_test.py (from_task) and episode_run_test.py (the run); the run tests wait for the end on the run's record rather than on the episode's task. - deploy_agent's live_mcp_url is public and bound_deployment takes env_step_id and step_id by keyword. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/task/episode.py | 155 +++------- src/agent_env/task/episode_rules.py | 160 ++++++++++ .../task_step/task_steps/deploy_agent.py | 22 +- .../{episode_test.py => episode_run_test.py} | 292 +----------------- tst/unit/task/episode_spec_test.py | 291 +++++++++++++++++ .../test_deploy_agent_live_endpoint.py | 4 +- .../test_parallel_deploy_env_binding.py | 24 +- 7 files changed, 533 insertions(+), 415 deletions(-) create mode 100644 src/agent_env/task/episode_rules.py rename tst/unit/task/{episode_test.py => episode_run_test.py} (51%) create mode 100644 tst/unit/task/episode_spec_test.py diff --git a/src/agent_env/task/episode.py b/src/agent_env/task/episode.py index fe2b3edc..3d2ab9b3 100644 --- a/src/agent_env/task/episode.py +++ b/src/agent_env/task/episode.py @@ -15,21 +15,17 @@ from agent_env.entity_refs import EntityRef from agent_env.providers.sandbox_providers.sandbox_provider import sandbox_request_headers_for_url +from agent_env.task.episode_rules import problems from agent_env.task.store import record_task_cancelled from agent_env.task.task import Task, substitute_seed_in_steps from agent_env.task.teardown import teardown_run 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 _live_mcp_url, bound_deployment +from agent_env.task_step.task_steps.deploy_agent import bound_deployment, live_mcp_url from agent_env.task_step.task_steps.prompt_agent import PromptAgentTaskStep logger = logging.getLogger(__name__) -_AGENT_DEPLOYERS = ("install_agent", "deploy_human_agent") -_AGENT_TARGETING = ( - "add_skills", "register_agent_triggers", "snapshot_agent_state", "load_artifact", "verify_sandbox", "run_code", -) - _live: dict[str, Episode] = {} @@ -71,15 +67,15 @@ def from_task( seed = dict(seed) if seed else None seeded = substitute_seed_in_steps(task.steps, seed or {}) outside = _outside_agents(task, seeded, agents) - problems = _problems(seeded, outside) + found = problems(seeded, outside) verifiers = _top_level_verifiers(seeded) if len(verifiers) != 1: - problems.append( + found.append( f"the reward is the score of the one verifier no aggregate_verifiers step lists, and the task has " f"{len(verifiers)}" + (f": {[step.id for step in verifiers]}" if verifiers else "") ) - if problems: - raise EpisodeError(f"task {task.id!r} can't run as an episode: " + "; ".join(problems)) + if found: + 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 @@ -101,17 +97,8 @@ async def start(self, *, run_group_id: str | None = None) -> Episode: metadata["seed"] = dict(self._seed) if run_group_id is not None: metadata["run_group_id"] = run_group_id - episode = Episode(self, episode_id, TaskStepContext(metadata=metadata)) - _live[episode_id] = episode - running = asyncio.Event() - episode._run = asyncio.create_task(episode._drive(running)) - try: - await running.wait() - except asyncio.CancelledError: - episode._abandoned = "start() was cancelled" - episode._run.cancel() - await asyncio.wait((episode._run,)) - raise + episode = Episode(episode_id, self._copy, self._reward_verifier_id, TaskStepContext(metadata=metadata)) + await episode._launch() return episode @@ -134,9 +121,20 @@ def __init__( 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 not self._reply.done(): + 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 Episode: """One run of an ``EpisodeSpec``'s task, which waits at each outside turn for the caller to play it. @@ -144,9 +142,10 @@ class Episode: 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, spec: EpisodeSpec, id: str, context: TaskStepContext): + def __init__(self, id: str, task: Task, reward_verifier_id: str, context: TaskStepContext): self.id = id - self._spec = spec + self._task = task + self._reward_verifier_id = reward_verifier_id self._context = context self._run: asyncio.Task[None] self._bindings: dict[str, tuple[str | None, dict[str, str]]] = {} @@ -163,7 +162,7 @@ async def next_turn(self) -> Turn | None: changed = self._changed while self._waiting and self._abandoned is None: turn = self._waiting.popleft() - if not turn._reply.done(): + if turn._open: return turn if self._run.done(): self._raise_unless_completed() @@ -174,16 +173,16 @@ async def result(self) -> float: """The reward, once the run and its teardown are over: the score of the task's one top-level verifier.""" while not self._run.done(): changed = self._changed - if open_turn := next((turn for turn in self._turns if not turn._reply.done()), None): + 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() - verifier_id = self._spec._reward_verifier_id - score = ((self._context.metadata.get("verifications") or {}).get(verifier_id) or {}).get("score") + verifications = self._context.metadata.get("verifications") or {} + score = (verifications.get(self._reward_verifier_id) or {}).get("score") if not isinstance(score, (int, float)): - raise EpisodeError(f"verifier {verifier_id!r} left no score") + raise EpisodeError(f"verifier {self._reward_verifier_id!r} left no score") return float(score) async def abandon(self, reason: str) -> None: @@ -193,17 +192,29 @@ async def abandon(self, reason: str) -> None: 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: - if not turn._reply.done(): - turn._reply.set_exception(EpisodeError(f"episode abandoned: {reason}")) + 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._spec._copy.run(context=self._context, instance_id=self.id) + await self._task.run(context=self._context, instance_id=self.id) finally: self._tearing_down = True try: @@ -225,6 +236,9 @@ def _raise_unless_completed(self) -> None: 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( @@ -238,7 +252,7 @@ async def _play(self, step: _EpisodeTurnStep) -> str: self._changed.set_result(None) self._changed = asyncio.get_running_loop().create_future() try: - return await asyncio.wait_for(turn._reply, step.timeout_seconds) + 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" @@ -268,75 +282,6 @@ def _outside_agents(task: Task, steps: list[TaskStep], agents: Sequence[str] | N return outside -def _problems(steps: list[TaskStep], outside: list[str]) -> list[str]: - """Why each step, and each outside agent, can't be run with the agent played from outside.""" - problems = [] - for name in outside: - deploys = [step for step in steps if step.type == "deploy_agent" and step.agent_name == name] - if len(deploys) != 1: - problems.append(f"agent {name!r} is played from outside, so it needs exactly one deploy_agent step, " - f"not {len(deploys)}") - outside_prompts = { - step.id: step.prompt_id for step in steps if step.type == "prompt_agent" and step.agent_name in outside - } - for step in steps: - problems.extend(f"step {step.id!r}: {why}" for why in _step_problems(step, outside, outside_prompts)) - return problems - - -def _step_problems(step: TaskStep, outside: list[str], outside_prompts: dict[str, str]) -> list[str]: - why = [] - if step.retry_config is not None: - why.append("retry_config can replay an outside turn") - agent = getattr(step, "agent_name", None) - if step.type == "prompt_agent" and agent in outside: - why.extend( - 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 - ) - elif step.type == "prompt_agent" and step.max_conversation_turns > 1 and step.user_agent_name in outside: - why.append(f"its user simulator {step.user_agent_name!r} is played from outside") - elif step.type == "deploy_agent" and agent in outside: - why.extend( - 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 - ) - elif step.type in _AGENT_DEPLOYERS and agent in outside: - why.append(f"{step.type} deploys {agent!r}, which is played from outside") - elif step.type.startswith("verify_a2a_"): - why.append("validates the task's deployed A2A agent, and an episode plays one from outside") - elif step.type in _AGENT_TARGETING and agent in outside: - why.append(f"{step.type} needs {agent!r} as a real agent, and it is played from outside") - elif step.type == "collect_artifacts": - if agent in outside and not (step.env_id or step.container_name or step.sandbox_name): - why.append(f"collects from {agent!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}") - elif step.type == "register_env_triggers" and step.executor_agent_name in outside: - why.append(f"its executor {step.executor_agent_name!r} is played from outside") - elif step.type == "peer_agents": - why.extend( - f"peers {name!r}, which is played from outside" for name in dict.fromkeys( - name for peering in step.peerings for name in (peering.source_agent_name, *peering.peer_agent_names) - ) if name in outside - ) - elif step.type == "rubrics_verifier": - 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 agent in outside: - why.append(f"judges in {agent!r}'s sandbox, and it is played from outside") - return why - - def _top_level_verifiers(steps: list[TaskStep]) -> list[TaskStep]: aggregated = {vid for step in steps if step.type == "aggregate_verifiers" for vid in step.verifier_ids} return [step for step in steps if getattr(step, "verifier_id", None) and step.verifier_id not in aggregated] @@ -398,9 +343,9 @@ def from_dict(cls, data: dict[str, Any]) -> _EpisodeAgentStep: async def execute(self, context: TaskStepContext) -> TaskStepContext: mcp_urls = {} for env_id in self.env_ids: - deployed = bound_deployment(context, env_id, self) - mcp_urls[env_id] = deployed.mcp_url if deployed is not None else _live_mcp_url(env_id) - _episode_of(context)._bindings[self.agent_name] = (self.role, mcp_urls) + 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 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_step/task_steps/deploy_agent.py b/src/agent_env/task_step/task_steps/deploy_agent.py index f39a3aea..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 @@ -104,10 +104,11 @@ 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, step: TaskStep) -> Optional[DeployedEnv]: - """The deployment of `env_id` the agent `step` deploys is bound to, by its `env_step_id`, or None if it isn't - deployed here.""" - env_step_id = step.env_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. @@ -115,13 +116,13 @@ def bound_deployment(context: TaskStepContext, env_id: str, step: TaskStep) -> O 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"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"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." ) @@ -304,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: bound_deployment(context, env_id, self) 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: @@ -388,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) diff --git a/tst/unit/task/episode_test.py b/tst/unit/task/episode_run_test.py similarity index 51% rename from tst/unit/task/episode_test.py rename to tst/unit/task/episode_run_test.py index 1a134a41..c5e8fe0c 100644 --- a/tst/unit/task/episode_test.py +++ b/tst/unit/task/episode_run_test.py @@ -1,6 +1,4 @@ -"""An episode: a run of a task with some of its agents played from outside. What ``EpisodeSpec.from_task`` reads off -a task (the outside agents, their turns, the reward), what it refuses, and the copy of the task it runs; then the run -itself: its turns, its reward, abandoning it, teardown and its record.""" +"""The run of an episode, on fake steps: its turns, its reward, abandoning it, teardown and its record.""" import asyncio @@ -10,273 +8,9 @@ import agent_env.task.episode as episode_module from agent_env.config.runtime import Config from agent_env.env.env import DeployedEnv -from agent_env.task import AgentSpec, EpisodeError, EpisodeSpec, Task, TurnSpec -from agent_env.task.episode import _EpisodeAgentStep, _EpisodeTurnStep -from agent_env.task.task import substitute_seed_in_steps +from agent_env.task import EpisodeError, EpisodeSpec, Task from agent_env.task.teardown import TeardownReport from agent_env.task_step.task_step import TaskStep -from agent_env.task_step.task_steps.deploy_agent import DeployAgentTaskStep -from agent_env.task_step.task_steps.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, agents=None): - with pytest.raises(EpisodeError) as raised: - EpisodeSpec.from_task(task, agents=agents) - return str(raised.value) - - -def test_a_task_with_one_prompted_agent_is_played_as_that_agent_with_its_turns_seeded(): - spec = EpisodeSpec.from_task(_portsim(), seed=SEED) - - assert spec.agents == {"default-agent": AgentSpec(role="planner@Valencia", turns=( - TurnSpec(step_id="play", prompt="Plan week 3.", system_prompt="You plan Valencia."), - ))} - assert spec._reward_verifier_id == "berth" - - -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 list(spec.agents) == ["planner"] - assert [step.type for step in spec._copy.steps] == [ - "deploy_env", "load_artifact", "episode_agent", "episode_turn", "env_outcome_verifier"] - assert spec.agents["planner"].turns[0].system_prompt == "You plan Valencia." - - -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.agents["default-agent"].turns[0].system_prompt == "Plan only the berths." - - -def test_each_prompt_of_an_outside_agent_is_a_turn_in_task_order(): - again = {**PLAY, "id": "replan", "prompt_id": "replan", "prompt": "Replan.", "depends_on": ["play"]} - spec = EpisodeSpec.from_task(_portsim(again)) - - assert [turn.step_id for turn in spec.agents["default-agent"].turns] == ["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 agents=" - assert list(EpisodeSpec.from_task(task, agents=["bob"]).agents) == ["bob"] - assert list(EpisodeSpec.from_task(task, agents=["bob", "default-agent"]).agents) == ["bob", "default-agent"] - - -@pytest.mark.parametrize(("agents", "message"), [ - (["carol"], "task 'dock' doesn't prompt ['carol']; it prompts ['default-agent']"), - ([], "agents= names no agent; task 'dock' prompts ['default-agent']"), -]) -def test_only_a_prompted_agent_can_be_played(agents, message): - assert _refusal(_portsim(), agents=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, 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), 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"] - - -def test_the_reward_is_the_verifier_no_aggregate_lists(): - graded = {"id": "graded", "type": "agent_prompt_response_verifier", "prompt_id": "play", "verifier_id": "reply", - "criteria": [{"type": "response_contains", "needles": ["berths"]}]} - overall = {"id": "overall", "type": "aggregate_verifiers", "verifier_ids": ["berth", "reply"], - "verifier_id": "overall"} - - assert EpisodeSpec.from_task(_portsim(graded, overall))._reward_verifier_id == "overall" - - -@pytest.mark.parametrize(("extra", "replace", "count"), [ - ((), {"verify": None}, "0"), - (({**VERIFY, "id": "verify-2", "verifier_id": "crane"},), {}, "2: ['verify', 'verify-2']"), -]) -def test_a_task_without_exactly_one_top_level_verifier_has_no_reward(extra, replace, count): - assert _refusal(_portsim(*extra, **replace)) == ( - "task 'dock' can't run as an episode: the reward is the score of the one verifier no aggregate_verifiers " - f"step lists, and the task has {count}") - - -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), 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)) - class _Env(TaskStep): """Deploys ``env_id`` with no sandbox, or fails as an env that doesn't come up.""" @@ -395,11 +129,12 @@ async def teardowns(monkeypatch, local_stores): def _spec(*steps, **kwargs): - return EpisodeSpec.from_task(_task(*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 asyncio.wait((episode._run,)) + await _until(lambda: Task.get_instance(episode.id).status != "running") async def _until(holds): @@ -487,19 +222,6 @@ async def test_a_seed_reaches_the_agents_env_binding_and_role(teardowns): assert (turn.mcp_urls, turn.role) == ({"pantry": "http://127.0.0.1:20031/pantry/mcp"}, "bob@example.com") -def test_the_seed_pass_rebuilds_each_step_as_its_own_class(): - class _Unregistered(_Replies): - pass - - steps = [*_portsim().steps, _Unregistered(id="extra", version=None, verifier_id="")] - - 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].verifier_id) == ( - "planner@Valencia", "Plan week 3.", "Valencia") - - @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() @@ -748,9 +470,9 @@ async def test_cancelling_the_caller_of_abandon_still_abandons_the_run(teardowns abandoning.cancel() with pytest.raises(asyncio.CancelledError): await abandoning - await _ended(episode) + 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 (Task.get_instance(episode.id).status,Task.get_instance(episode.id).error) == ("cancelled", "stop") assert teardowns.finished == [episode.id] diff --git a/tst/unit/task/episode_spec_test.py b/tst/unit/task/episode_spec_test.py new file mode 100644 index 00000000..293c0325 --- /dev/null +++ b/tst/unit/task/episode_spec_test.py @@ -0,0 +1,291 @@ +"""What ``EpisodeSpec.from_task`` reads off a task (the outside agents, their turns, the reward), what it refuses, and +the copy of the task it runs.""" + +import pytest + +from agent_env.task import AgentSpec, EpisodeError, EpisodeSpec, Task, TurnSpec +from agent_env.task.episode import _EpisodeAgentStep, _EpisodeTurnStep +from agent_env.task.episode_rules import CHECKS, UNAFFECTED +from agent_env.task.task import substitute_seed_in_steps +from agent_env.task_step.registry import get_task_step_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, agents=None): + with pytest.raises(EpisodeError) as raised: + EpisodeSpec.from_task(task, agents=agents) + return str(raised.value) + + +def test_a_task_with_one_prompted_agent_is_played_as_that_agent_with_its_turns_seeded(): + spec = EpisodeSpec.from_task(_portsim(), seed=SEED) + + assert spec.agents == {"default-agent": AgentSpec(role="planner@Valencia", turns=( + TurnSpec(step_id="play", prompt="Plan week 3.", system_prompt="You plan Valencia."), + ))} + assert spec._reward_verifier_id == "berth" + + +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 list(spec.agents) == ["planner"] + assert [step.type for step in spec._copy.steps] == [ + "deploy_env", "load_artifact", "episode_agent", "episode_turn", "env_outcome_verifier"] + assert spec.agents["planner"].turns[0].system_prompt == "You plan Valencia." + + +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.agents["default-agent"].turns[0].system_prompt == "Plan only the berths." + + +def test_each_prompt_of_an_outside_agent_is_a_turn_in_task_order(): + again = {**PLAY, "id": "replan", "prompt_id": "replan", "prompt": "Replan.", "depends_on": ["play"]} + spec = EpisodeSpec.from_task(_portsim(again)) + + assert [turn.step_id for turn in spec.agents["default-agent"].turns] == ["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 agents=" + assert list(EpisodeSpec.from_task(task, agents=["bob"]).agents) == ["bob"] + assert list(EpisodeSpec.from_task(task, agents=["bob", "default-agent"]).agents) == ["bob", "default-agent"] + + +@pytest.mark.parametrize(("agents", "message"), [ + (["carol"], "task 'dock' doesn't prompt ['carol']; it prompts ['default-agent']"), + ([], "agents= names no agent; task 'dock' prompts ['default-agent']"), +]) +def test_only_a_prompted_agent_can_be_played(agents, message): + assert _refusal(_portsim(), agents=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, 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), 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"] + + +def test_the_reward_is_the_verifier_no_aggregate_lists(): + graded = {"id": "graded", "type": "agent_prompt_response_verifier", "prompt_id": "play", "verifier_id": "reply", + "criteria": [{"type": "response_contains", "needles": ["berths"]}]} + overall = {"id": "overall", "type": "aggregate_verifiers", "verifier_ids": ["berth", "reply"], + "verifier_id": "overall"} + + assert EpisodeSpec.from_task(_portsim(graded, overall))._reward_verifier_id == "overall" + + +@pytest.mark.parametrize(("extra", "replace", "count"), [ + ((), {"verify": None}, "0"), + (({**VERIFY, "id": "verify-2", "verifier_id": "crane"},), {}, "2: ['verify', 'verify-2']"), +]) +def test_a_task_without_exactly_one_top_level_verifier_has_no_reward(extra, replace, count): + assert _refusal(_portsim(*extra, **replace)) == ( + "task 'dock' can't run as an episode: the reward is the score of the one verifier no aggregate_verifiers " + f"step lists, and the task has {count}") + + +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), 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_registered_step_type_is_checked_or_unaffected(): + registered, checked = set(get_task_step_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 883fbb5e..5b95f743 100644 --- a/tst/unit/task_step/test_parallel_deploy_env_binding.py +++ b/tst/unit/task_step/test_parallel_deploy_env_binding.py @@ -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 = bound_deployment(context, ENV_ID, _agent_bound_to("env-a")) - bound_b = bound_deployment(context, ENV_ID, _agent_bound_to("env-b")) + 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"): - bound_deployment(context, ENV_ID, _agent_bound_to(None)) + _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 = bound_deployment(context, ENV_ID, _agent_bound_to(None)) + 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"): - bound_deployment(context, ENV_ID, _agent_bound_to("env-typo")) + _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"): - bound_deployment(context, ENV_ID, _agent_bound_to("env-typo")) + _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 bound_deployment(context, "oc-halt-ask-human", _agent_bound_to("env-a")) is None + assert _bound(context, "oc-halt-ask-human", "env-a") is None def test_deploy_agent_env_step_id_round_trips(): From 335ff87d125639aa69595db27018b53957582afc Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Tue, 13 Oct 2026 09:35:22 -0700 Subject: [PATCH 10/17] test(task): hold only the built-in step types to the episode rules A plugin or a [task_steps] impls step is in the configured registry but not in the rules, which don't inspect plugin types, so the guard failed on a valid local setup. Co-Authored-By: Claude Opus 5.5 --- tst/unit/task/episode_spec_test.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tst/unit/task/episode_spec_test.py b/tst/unit/task/episode_spec_test.py index 293c0325..752a7cf7 100644 --- a/tst/unit/task/episode_spec_test.py +++ b/tst/unit/task/episode_spec_test.py @@ -7,7 +7,7 @@ from agent_env.task.episode import _EpisodeAgentStep, _EpisodeTurnStep from agent_env.task.episode_rules import CHECKS, UNAFFECTED from agent_env.task.task import substitute_seed_in_steps -from agent_env.task_step.registry import get_task_step_registry +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 @@ -270,8 +270,8 @@ def test_every_refusal_is_reported_at_once(): def test_steps_that_need_only_what_an_outside_player_leaves_are_allowed(replace): EpisodeSpec.from_task(_portsim(**replace)) -def test_every_registered_step_type_is_checked_or_unaffected(): - registered, checked = set(get_task_step_registry()), set(CHECKS) +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) == [] From 407768635d65ddc27d2fcf1139a93c6ea7cf3dec Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Tue, 13 Oct 2026 15:59:40 -0700 Subject: [PATCH 11/17] refactor(task): an episode's turns exist only at run time, and its result is every score From the API review of the episode: - AgentSpec and TurnSpec are gone: everything an agent is asked arrives on a Turn from next_turn(), so prompt turns and the message turns to come share one model. EpisodeSpec exposes only the agents the loop plays. - agents is trainer_agents throughout: from_task's parameter, the spec's attribute, the record's episode tag and the error messages. - start() takes no run_group_id: nothing in core reads it back from an episode's record. - result() returns every score the run recorded, by step id, as agent-env run reports them, and the caller makes the reward. The rule that a task have exactly one top-level verifier is gone. The step-id mapping moves from bundle run's scores to task.recorded_scores, which both use. Co-Authored-By: Claude Opus 5.5 --- AGENTS.md | 8 +- README.md | 4 +- src/agent_env/bundle/run.py | 9 +- src/agent_env/task/__init__.py | 4 +- src/agent_env/task/episode.py | 98 +++++-------------- src/agent_env/task/task.py | 12 +++ .../env/envs/server_env_local_test.py | 10 +- tst/unit/task/episode_run_test.py | 66 +++++++------ tst/unit/task/episode_spec_test.py | 69 +++++-------- 9 files changed, 111 insertions(+), 169 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index ceabf735..ed08d5fa 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -156,10 +156,10 @@ orchestrators call it the same way. What callers rely on: - 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.py`) is one `Task.run()` of a copy of the task in which - each outside agent's `deploy_agent` and `prompt_agent` steps 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 + 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 5f71ffad..9b4c9b05 100644 --- a/README.md +++ b/README.md @@ -58,10 +58,10 @@ while (turn := await episode.next_turn()) is not None: async with ToolSession(mcp_url=turn.mcp_urls["my-env"], role=turn.role, headers=turn.headers["my-env"]) as tools: ... # your model reads turn.prompt and calls tools.call_tool(name, arguments) await turn.end(response="Done.") -reward = await episode.result() +scores = await episode.result() # every score the run recorded, by step id ``` -The reward is the score of the task's one top-level verifier, `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. +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 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/task/__init__.py b/src/agent_env/task/__init__.py index 0cce40a4..f7d2f7c2 100644 --- a/src/agent_env/task/__init__.py +++ b/src/agent_env/task/__init__.py @@ -1,7 +1,7 @@ """Task module for AgentEnv.""" from .task import Task -from .episode import AgentSpec, Episode, EpisodeError, EpisodeSpec, Turn, TurnSpec +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, @@ -10,12 +10,10 @@ __all__ = [ "Task", - "AgentSpec", "Episode", "EpisodeError", "EpisodeSpec", "Turn", - "TurnSpec", "TaskStore", "TaskQuery", "get_task_store", diff --git a/src/agent_env/task/episode.py b/src/agent_env/task/episode.py index 3d2ab9b3..e97692a4 100644 --- a/src/agent_env/task/episode.py +++ b/src/agent_env/task/episode.py @@ -1,6 +1,6 @@ """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 the reward is the score ``agent-env run`` would give. +AgentEnv runs everything else exactly as ``Task.run`` does, so its scores are the ones ``agent-env run`` would report. """ from __future__ import annotations @@ -10,14 +10,13 @@ import uuid from collections import deque from collections.abc import Mapping, Sequence -from dataclasses import dataclass from typing import Any, ClassVar, Optional from agent_env.entity_refs import EntityRef from agent_env.providers.sandbox_providers.sandbox_provider import sandbox_request_headers_for_url from agent_env.task.episode_rules import problems from agent_env.task.store import record_task_cancelled -from agent_env.task.task import Task, substitute_seed_in_steps +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 PromptResponse, TaskStepContext from agent_env.task_step.task_step import TaskStep @@ -33,71 +32,43 @@ class EpisodeError(RuntimeError): """A task that can't run as an episode, or an episode asked for what it can't give.""" -@dataclass(frozen=True) -class TurnSpec: - step_id: str - prompt: str - system_prompt: str | None - - -@dataclass(frozen=True) -class AgentSpec: - role: str | None - turns: tuple[TurnSpec, ...] - - class EpisodeSpec: - """A task read as an episode: the agents played from outside, their turns, and the copy of the task that runs.""" + """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, agents: dict[str, AgentSpec], reward_verifier_id: str, seed: dict[str, str] | None, - ): + def __init__(self, *, copy: Task, trainer_agents: tuple[str, ...], seed: dict[str, str] | None): self._copy = copy - self.agents = agents - self._reward_verifier_id = reward_verifier_id + self.trainer_agents = trainer_agents self._seed = seed @classmethod def from_task( - cls, task: Task, *, agents: Sequence[str] | None = None, seed: Mapping[str, str] | None = None, + cls, task: Task, *, trainer_agents: Sequence[str] | None = None, seed: Mapping[str, str] | None = None, turn_timeout: float | None = None, ) -> EpisodeSpec: - """``turn_timeout`` is every outside 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.""" + """``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 = _outside_agents(task, seeded, agents) - found = problems(seeded, outside) - verifiers = _top_level_verifiers(seeded) - if len(verifiers) != 1: - found.append( - f"the reward is the score of the one verifier no aggregate_verifiers step lists, and the task has " - f"{len(verifiers)}" + (f": {[step.id for step in verifiers]}" if verifiers else "") - ) - if found: + 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), - agents={name: _agent_spec(seeded, name) for name in outside}, - reward_verifier_id=verifiers[0].verifier_id, - seed=seed, - ) + return cls(copy=Task(id=task.id, version=task.version, steps=steps), trainer_agents=tuple(outside), seed=seed) - async def start(self, *, run_group_id: str | None = None) -> Episode: + 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, "agents": list(self.agents)}} + metadata: dict[str, Any] = {"episode": {"id": episode_id, "trainer_agents": list(self.trainer_agents)}} if self._seed: metadata["seed"] = dict(self._seed) - if run_group_id is not None: - metadata["run_group_id"] = run_group_id - episode = Episode(episode_id, self._copy, self._reward_verifier_id, TaskStepContext(metadata=metadata)) + episode = Episode(episode_id, self._copy, TaskStepContext(metadata=metadata)) await episode._launch() return episode @@ -142,10 +113,9 @@ class Episode: 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, reward_verifier_id: str, context: TaskStepContext): + def __init__(self, id: str, task: Task, context: TaskStepContext): self.id = id self._task = task - self._reward_verifier_id = reward_verifier_id self._context = context self._run: asyncio.Task[None] self._bindings: dict[str, tuple[str | None, dict[str, str]]] = {} @@ -169,8 +139,9 @@ async def next_turn(self) -> Turn | None: return None await self._wait(changed) - async def result(self) -> float: - """The reward, once the run and its teardown are over: the score of the task's one top-level verifier.""" + 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): @@ -179,11 +150,7 @@ async def result(self) -> float: ) await self._wait(changed) self._raise_unless_completed() - verifications = self._context.metadata.get("verifications") or {} - score = (verifications.get(self._reward_verifier_id) or {}).get("score") - if not isinstance(score, (int, float)): - raise EpisodeError(f"verifier {self._reward_verifier_id!r} left no score") - return float(score) + return recorded_scores(self._task.steps, 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 @@ -265,36 +232,23 @@ def _episode_of(context: TaskStepContext) -> Episode: return _live[context.metadata["episode"]["id"]] -def _outside_agents(task: Task, steps: list[TaskStep], agents: Sequence[str] | None) -> list[str]: +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 agents is None: + if named is None: if len(prompted) != 1: raise EpisodeError( - f"task {task.id!r} prompts {prompted}; name the ones to play with agents=" if prompted + 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(agents)) + outside = list(dict.fromkeys(named)) if not outside: - raise EpisodeError(f"agents= names no agent; task {task.id!r} prompts {prompted}") + 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 _top_level_verifiers(steps: list[TaskStep]) -> list[TaskStep]: - aggregated = {vid for step in steps if step.type == "aggregate_verifiers" for vid in step.verifier_ids} - return [step for step in steps if getattr(step, "verifier_id", None) and step.verifier_id not in aggregated] - - -def _agent_spec(steps: list[TaskStep], name: str) -> AgentSpec: - deploy = next(step for step in steps if step.type == "deploy_agent" and step.agent_name == name) - return AgentSpec(role=deploy.role, turns=tuple( - TurnSpec(step_id=step.id, prompt=step.prompt, system_prompt=step.system_prompt or deploy.system_prompt) - for step in steps if step.type == "prompt_agent" and step.agent_name == name - )) - - 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 diff --git a/src/agent_env/task/task.py b/src/agent_env/task/task.py index ad7fb5a7..46960671 100644 --- a/src/agent_env/task/task.py +++ b/src/agent_env/task/task.py @@ -117,6 +117,18 @@ def _sub(value): 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/tst/integration/env/envs/server_env_local_test.py b/tst/integration/env/envs/server_env_local_test.py index bae78410..a958677f 100644 --- a/tst/integration/env/envs/server_env_local_test.py +++ b/tst/integration/env/envs/server_env_local_test.py @@ -211,8 +211,8 @@ async def verify(mcp_url): @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 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, the reward - the verifier's score; and every episode, played, not played or abandoned, leaving nothing behind.""" + which plays its turn through the env's tools as its role and is never deployed; its reply recorded and scored + by the verifier; and every episode, played, not played or abandoned, 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") @@ -229,7 +229,7 @@ async def test_a_tasks_agent_is_played_from_a_loop_of_ones_own_as_an_episode(loc {"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"}, - ]}), agents=["policy"]) + ]}), trainer_agents=["policy"]) async def play(episode, add): turn = await episode.next_turn() @@ -243,10 +243,10 @@ async def play(episode, add): sandboxes = tmp_path / "sandboxes" played, unplayed = await spec.start(), await spec.start() - assert await asyncio.gather(play(played, True), play(unplayed, False)) == [1.0, 0.0] + 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"]["agents"] == ["policy"] + 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:") diff --git a/tst/unit/task/episode_run_test.py b/tst/unit/task/episode_run_test.py index c5e8fe0c..6330fa07 100644 --- a/tst/unit/task/episode_run_test.py +++ b/tst/unit/task/episode_run_test.py @@ -1,4 +1,4 @@ -"""The run of an episode, on fake steps: its turns, its reward, abandoning it, teardown and its record.""" +"""The run of an episode, on fake steps: its turns, its scores, abandoning it, teardown and its record.""" import asyncio @@ -154,12 +154,12 @@ async def test_the_run_waits_at_an_outside_turn_and_carries_on_once_it_ends(tear 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() == 1.0 + 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, "agents": ["alice"]} + 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"] == [] @@ -181,7 +181,7 @@ async def test_two_outside_agents_turns_are_each_handed_out_once(teardowns): 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"]}, agents=["alice", "bob"], + {**REPLIES, "depends_on": ["turn", "bob-turn"]}, trainer_agents=["alice", "bob"], ).start() first = await episode.next_turn() @@ -190,7 +190,7 @@ async def test_two_outside_agents_turns_are_each_handed_out_once(teardowns): await first.end("one") await second.end("two") - assert await episode.result() == 1.0 + assert await episode.result() == {"replies": 1.0} @pytest.mark.asyncio @@ -204,21 +204,24 @@ async def play(n): played = dict(await asyncio.gather(*(play(n) for n in range(8)))) - assert list(played.values()) == [1.0] * 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_agents_env_binding_and_role(teardowns): +async def test_a_seed_reaches_the_turns_prompt_env_binding_and_role(teardowns): episode = await _spec( - {**ITEMS, "env_id": ""}, {**ALICE, "env_ids": [""], "role": ""}, TURN, REPLIES, + {**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") @@ -265,7 +268,7 @@ async def test_a_tolerant_turn_past_its_deadline_is_scored_without_a_reply(teard episode = await _spec(ITEMS, ALICE, tolerant, REPLIES).start() await _ended(episode) - assert await episode.result() == 0.0 + 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"] @@ -280,19 +283,19 @@ async def test_the_result_is_refused_while_a_turn_waits_or_is_open(teardowns): with pytest.raises(EpisodeError, match="isn't ended"): await episode.result() await turn.end("done") - assert await episode.result() == 1.0 + assert await episode.result() == {"replies": 1.0} @pytest.mark.asyncio -@pytest.mark.parametrize(("reply", "reward"), [("Added a crate.", 1.0), (None, 0.0)]) -async def test_the_reply_is_graded_as_a_prompt_agents_would_be(teardowns, reply, reward): +@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() == reward + assert await episode.result() == {"graded": score} @pytest.mark.asyncio @@ -349,7 +352,7 @@ async def test_abandoning_during_teardown_waits_for_it_and_leaves_the_runs_outco await abandoning assert teardowns.finished == [episode.id] - assert await episode.result() == 1.0 + assert await episode.result() == {"replies": 1.0} assert Task.get_instance(episode.id).status == "completed" @@ -372,17 +375,7 @@ async def test_cancelling_a_caller_cancels_only_its_wait(teardowns): with pytest.raises(asyncio.CancelledError): await scoring _gate("after").set() - assert await episode.result() == 1.0 - - -@pytest.mark.asyncio -async def test_a_run_group_id_groups_the_record(teardowns): - episode = await _spec(ITEMS, ALICE, TURN, REPLIES).start(run_group_id="grpo-step-7") - - await (await episode.next_turn()).end("done") - await episode.result() - - assert Task.get_instance(episode.id).context["metadata"]["run_group_id"] == "grpo-step-7" + assert await episode.result() == {"replies": 1.0} @pytest.mark.asyncio @@ -395,7 +388,7 @@ async def test_an_agent_named_by_the_seed_is_played_from_outside_in_the_run(tear assert turn.agent == "bob" await turn.end("done") - assert await episode.result() == 1.0 + assert await episode.result() == {"replies": 1.0} @pytest.mark.asyncio @@ -411,15 +404,28 @@ async def test_a_turn_ended_twice_keeps_its_first_reply(teardowns): @pytest.mark.asyncio -async def test_a_reward_verifier_that_left_no_score_is_no_reward(teardowns): +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_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).start() + episode = await _spec(ITEMS, ALICE, TURN, unscored, REPLIES).start() await (await episode.next_turn()).end("done") - with pytest.raises(EpisodeError, match="verifier 'graded' left no score"): - await episode.result() + assert await episode.result() == {"replies": 1.0} BREAKS_WHILE_THE_TURN_IS_OPEN = ( diff --git a/tst/unit/task/episode_spec_test.py b/tst/unit/task/episode_spec_test.py index 752a7cf7..edf153e1 100644 --- a/tst/unit/task/episode_spec_test.py +++ b/tst/unit/task/episode_spec_test.py @@ -1,9 +1,9 @@ -"""What ``EpisodeSpec.from_task`` reads off a task (the outside agents, their turns, the reward), what it refuses, and -the copy of the task it runs.""" +"""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 AgentSpec, EpisodeError, EpisodeSpec, Task, TurnSpec +from agent_env.task import EpisodeError, EpisodeSpec, Task from agent_env.task.episode import _EpisodeAgentStep, _EpisodeTurnStep from agent_env.task.episode_rules import CHECKS, UNAFFECTED from agent_env.task.task import substitute_seed_in_steps @@ -35,19 +35,14 @@ def _portsim(*extra, **replace): return _task(*(step for step in steps.values() if step is not None), *extra) -def _refusal(task, agents=None): +def _refusal(task, trainer_agents=None): with pytest.raises(EpisodeError) as raised: - EpisodeSpec.from_task(task, agents=agents) + 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_with_its_turns_seeded(): - spec = EpisodeSpec.from_task(_portsim(), seed=SEED) - - assert spec.agents == {"default-agent": AgentSpec(role="planner@Valencia", turns=( - TurnSpec(step_id="play", prompt="Plan week 3.", system_prompt="You plan Valencia."), - ))} - assert spec._reward_verifier_id == "berth" +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(): @@ -56,39 +51,40 @@ def test_an_agent_named_by_the_seed_is_the_one_played_from_outside(): seed={**SEED, "who": "planner"}, ) - assert list(spec.agents) == ["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"] - assert spec.agents["planner"].turns[0].system_prompt == "You plan Valencia." 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.agents["default-agent"].turns[0].system_prompt == "Plan only the berths." + assert spec._copy.steps[3].system_prompt == "Plan only the berths." -def test_each_prompt_of_an_outside_agent_is_a_turn_in_task_order(): +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 [turn.step_id for turn in spec.agents["default-agent"].turns] == ["play", "replan"] + 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 agents=" - assert list(EpisodeSpec.from_task(task, agents=["bob"]).agents) == ["bob"] - assert list(EpisodeSpec.from_task(task, agents=["bob", "default-agent"]).agents) == ["bob", "default-agent"] + 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(("agents", "message"), [ +@pytest.mark.parametrize(("trainer_agents", "message"), [ (["carol"], "task 'dock' doesn't prompt ['carol']; it prompts ['default-agent']"), - ([], "agents= names no agent; task 'dock' prompts ['default-agent']"), + ([], "trainer_agents= names no agent; task 'dock' prompts ['default-agent']"), ]) -def test_only_a_prompted_agent_can_be_played(agents, message): - assert _refusal(_portsim(), agents=agents) == message +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(): @@ -97,7 +93,7 @@ def test_a_task_prompting_no_agent_has_nothing_to_play(): 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, agents=["default-agent"])._copy + 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] == [ @@ -139,31 +135,12 @@ def test_a_turn_timeout_replaces_every_outside_prompts_own(): def test_an_inside_agent_keeps_its_own_steps(): - copy = EpisodeSpec.from_task(_portsim(BOB, ASK_BOB), agents=["bob"])._copy + 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"] -def test_the_reward_is_the_verifier_no_aggregate_lists(): - graded = {"id": "graded", "type": "agent_prompt_response_verifier", "prompt_id": "play", "verifier_id": "reply", - "criteria": [{"type": "response_contains", "needles": ["berths"]}]} - overall = {"id": "overall", "type": "aggregate_verifiers", "verifier_ids": ["berth", "reply"], - "verifier_id": "overall"} - - assert EpisodeSpec.from_task(_portsim(graded, overall))._reward_verifier_id == "overall" - - -@pytest.mark.parametrize(("extra", "replace", "count"), [ - ((), {"verify": None}, "0"), - (({**VERIFY, "id": "verify-2", "verifier_id": "crane"},), {}, "2: ['verify', 'verify-2']"), -]) -def test_a_task_without_exactly_one_top_level_verifier_has_no_reward(extra, replace, count): - assert _refusal(_portsim(*extra, **replace)) == ( - "task 'dock' can't run as an episode: the reward is the score of the one verifier no aggregate_verifiers " - f"step lists, and the task has {count}") - - 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 " @@ -241,7 +218,7 @@ 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), agents=["default-agent"]) + _portsim(BOB, ask), trainer_agents=["default-agent"]) def test_every_refusal_is_reported_at_once(): From 411fb20d885b8531cb1cea07dcb1b87b87bffd67 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Tue, 13 Oct 2026 16:36:06 -0700 Subject: [PATCH 12/17] fix(task): key an episode's scores by step when the seed names a verifier The run seeds every step field, verifier_id included, so a score lands under the seeded id; mapping it with the unseeded copy left it under that raw key instead of its step's id. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/task/episode.py | 3 ++- tst/unit/task/episode_run_test.py | 9 +++++++++ 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/src/agent_env/task/episode.py b/src/agent_env/task/episode.py index e97692a4..ea2b7e67 100644 --- a/src/agent_env/task/episode.py +++ b/src/agent_env/task/episode.py @@ -150,7 +150,8 @@ async def result(self) -> dict[str, float]: ) await self._wait(changed) self._raise_unless_completed() - return recorded_scores(self._task.steps, self._context) + 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 diff --git a/tst/unit/task/episode_run_test.py b/tst/unit/task/episode_run_test.py index 6330fa07..d50df579 100644 --- a/tst/unit/task/episode_run_test.py +++ b/tst/unit/task/episode_run_test.py @@ -417,6 +417,15 @@ def graded(step_id, needle): 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", From dbce463e39178c1cbc076c4f043fc8072a934672 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Tue, 13 Oct 2026 21:02:51 -0700 Subject: [PATCH 13/17] feat(task): a turn's tools, called as its agent's role with each URL's headers turn.tools() opens one ToolSession per env the agent acts in, as its role and with the headers each URL needs. list() names every tool across those envs, and call() routes a call to the env that offers it. A tool two envs offer is refused when the client opens, and nothing is left open. A loop no longer pairs mcp_urls, headers and role itself; the raw fields stay for one that brings its own MCP client. The README sample and the episode integration test use it. Co-Authored-By: Claude Opus 5.5 --- README.md | 5 +- src/agent_env/task/episode.py | 53 +++++++++++- .../env/envs/server_env_local_test.py | 7 +- tst/unit/task/episode_run_test.py | 83 +++++++++++++++++++ 4 files changed, 141 insertions(+), 7 deletions(-) diff --git a/README.md b/README.md index 9b4c9b05..b7595571 100644 --- a/README.md +++ b/README.md @@ -49,14 +49,13 @@ source .venv/bin/activate 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 agentenv_protocol import ToolSession 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 ToolSession(mcp_url=turn.mcp_urls["my-env"], role=turn.role, headers=turn.headers["my-env"]) as tools: - ... # your model reads turn.prompt and calls tools.call_tool(name, arguments) + 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 ``` diff --git a/src/agent_env/task/episode.py b/src/agent_env/task/episode.py index ea2b7e67..06ad3866 100644 --- a/src/agent_env/task/episode.py +++ b/src/agent_env/task/episode.py @@ -12,6 +12,8 @@ from collections.abc import Mapping, Sequence from typing import Any, ClassVar, Optional +from agentenv_protocol import EnvironmentTool, ToolResult, ToolSession + from agent_env.entity_refs import EntityRef from agent_env.providers.sandbox_providers.sandbox_provider import sandbox_request_headers_for_url from agent_env.task.episode_rules import problems @@ -74,7 +76,8 @@ async def start(self) -> Episode: class Turn: - """An outside agent's turn: what it is asked, the envs it acts in and the headers each needs, until ``end()``.""" + """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, @@ -89,6 +92,11 @@ def __init__( 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.""" @@ -107,6 +115,49 @@ def _fail(self, error: Exception) -> None: 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. diff --git a/tst/integration/env/envs/server_env_local_test.py b/tst/integration/env/envs/server_env_local_test.py index a958677f..4b7d42ae 100644 --- a/tst/integration/env/envs/server_env_local_test.py +++ b/tst/integration/env/envs/server_env_local_test.py @@ -235,9 +235,10 @@ async def play(episode, add): turn = await episode.next_turn() assert [agent["agent_name"] for agent in Task.get_instance(episode.id).context["deployed_agents"]] == ["echo"] if add: - async with ToolSession(mcp_url=turn.mcp_urls[env.id], role=turn.role, headers=turn.headers[env.id]) as own: - await own.call_tool("items_add_item", {"item": "crate"}) - assert json.loads((await own.call_tool("items_whoami")).text)["role"] == "alice@example.com" + 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() diff --git a/tst/unit/task/episode_run_test.py b/tst/unit/task/episode_run_test.py index d50df579..8950a6cd 100644 --- a/tst/unit/task/episode_run_test.py +++ b/tst/unit/task/episode_run_test.py @@ -4,6 +4,7 @@ import pytest import pytest_asyncio +from agentenv_protocol import EnvironmentTool, ToolResult import agent_env.task.episode as episode_module from agent_env.config.runtime import Config @@ -102,6 +103,37 @@ async def __call__(self, context, **_): 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 @@ -175,6 +207,57 @@ async def test_a_turn_carries_the_headers_each_url_needs(teardowns, monkeypatch) 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_module, "ToolSession", sessions) + monkeypatch.setattr(episode_module, "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"]} From 434ad400b3a309dbe71443c65829b82c00e0f5f0 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Tue, 13 Oct 2026 21:29:26 -0700 Subject: [PATCH 14/17] refactor(task): the episode as a package, task/episode/ The episode module is split by what each part does, with no change in behaviour: spec.py (EpisodeSpec and the copy of the task), run.py (Episode, Turn, TurnTools and EpisodeError), steps.py (episode_agent and episode_turn) and rules.py (the v1 rules, formerly task/episode_rules.py). agent_env.task.episode re-exports the public names, and agent_env.task still exports them. The unit tests move to tst/unit/task/episode/ to match. Co-Authored-By: Claude Opus 5.5 --- AGENTS.md | 4 +- src/agent_env/task/episode/__init__.py | 9 + .../{episode_rules.py => episode/rules.py} | 0 .../task/{episode.py => episode/run.py} | 168 +----------------- src/agent_env/task/episode/spec.py | 91 ++++++++++ src/agent_env/task/episode/steps.py | 80 +++++++++ tst/unit/task/episode/__init__.py | 0 .../run_test.py} | 14 +- .../spec_test.py} | 4 +- 9 files changed, 196 insertions(+), 174 deletions(-) create mode 100644 src/agent_env/task/episode/__init__.py rename src/agent_env/task/{episode_rules.py => episode/rules.py} (100%) rename src/agent_env/task/{episode.py => episode/run.py} (54%) create mode 100644 src/agent_env/task/episode/spec.py create mode 100644 src/agent_env/task/episode/steps.py create mode 100644 tst/unit/task/episode/__init__.py rename tst/unit/task/{episode_run_test.py => episode/run_test.py} (97%) rename tst/unit/task/{episode_spec_test.py => episode/spec_test.py} (99%) diff --git a/AGENTS.md b/AGENTS.md index ed08d5fa..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/episode.py` runs a task with some of its agents played from outside, as a training loop does (`EpisodeSpec`). | +| `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,7 +155,7 @@ 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.py`) is one `Task.run()` of a copy of the task in which +- 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 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 similarity index 100% rename from src/agent_env/task/episode_rules.py rename to src/agent_env/task/episode/rules.py diff --git a/src/agent_env/task/episode.py b/src/agent_env/task/episode/run.py similarity index 54% rename from src/agent_env/task/episode.py rename to src/agent_env/task/episode/run.py index 06ad3866..8ee4ddfd 100644 --- a/src/agent_env/task/episode.py +++ b/src/agent_env/task/episode/run.py @@ -1,29 +1,20 @@ -"""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 __future__ import annotations import asyncio import logging -import uuid from collections import deque -from collections.abc import Mapping, Sequence -from typing import Any, ClassVar, Optional +from typing import TYPE_CHECKING, Any from agentenv_protocol import EnvironmentTool, ToolResult, ToolSession -from agent_env.entity_refs import EntityRef from agent_env.providers.sandbox_providers.sandbox_provider import sandbox_request_headers_for_url -from agent_env.task.episode_rules import problems 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 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 -from agent_env.task_step.task_steps.prompt_agent import PromptAgentTaskStep +from agent_env.task_step.context import TaskStepContext + +if TYPE_CHECKING: + from agent_env.task.episode.steps import _EpisodeTurnStep logger = logging.getLogger(__name__) @@ -34,47 +25,6 @@ class EpisodeError(RuntimeError): """A task that can't run as an episode, or an episode asked for what it can't give.""" -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 - - 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()``.""" @@ -282,111 +232,3 @@ async def _play(self, step: _EpisodeTurnStep) -> str: def _episode_of(context: TaskStepContext) -> Episode: return _live[context.metadata["episode"]["id"]] - - -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 - - -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/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/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 similarity index 97% rename from tst/unit/task/episode_run_test.py rename to tst/unit/task/episode/run_test.py index 8950a6cd..4cd5ccdb 100644 --- a/tst/unit/task/episode_run_test.py +++ b/tst/unit/task/episode/run_test.py @@ -6,7 +6,7 @@ import pytest_asyncio from agentenv_protocol import EnvironmentTool, ToolResult -import agent_env.task.episode as episode_module +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 @@ -142,13 +142,13 @@ async def teardowns(monkeypatch, local_stores): monkeypatch.setattr(Config, "task_step_registry", lambda self: steps) GATES.clear() teardowns = _Teardowns() - monkeypatch.setattr(episode_module, "teardown_run", 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_module._live.values()): + for episode in list(episode_run._live.values()): await episode.abandon("the test ended") @@ -199,7 +199,7 @@ async def test_the_run_waits_at_an_outside_turn_and_carries_on_once_it_ends(tear @pytest.mark.asyncio async def test_a_turn_carries_the_headers_each_url_needs(teardowns, monkeypatch): - monkeypatch.setattr(episode_module, "sandbox_request_headers_for_url", lambda url: {"X-Sandbox": url}) + 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() @@ -214,8 +214,8 @@ async def test_a_turn_carries_the_headers_each_url_needs(teardowns, monkeypatch) def _tool_sessions(monkeypatch, offered): sessions = _ToolSessions(offered) - monkeypatch.setattr(episode_module, "ToolSession", sessions) - monkeypatch.setattr(episode_module, "sandbox_request_headers_for_url", lambda url: {"X-Sandbox": url}) + monkeypatch.setattr(episode_run, "ToolSession", sessions) + monkeypatch.setattr(episode_run, "sandbox_request_headers_for_url", lambda url: {"X-Sandbox": url}) return sessions @@ -584,5 +584,5 @@ async def test_a_cancelled_start_takes_its_run_down_with_it(teardowns): with pytest.raises(asyncio.CancelledError): await starting - assert not episode_module._live + 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 similarity index 99% rename from tst/unit/task/episode_spec_test.py rename to tst/unit/task/episode/spec_test.py index edf153e1..14e4d4f2 100644 --- a/tst/unit/task/episode_spec_test.py +++ b/tst/unit/task/episode/spec_test.py @@ -4,8 +4,8 @@ import pytest from agent_env.task import EpisodeError, EpisodeSpec, Task -from agent_env.task.episode import _EpisodeAgentStep, _EpisodeTurnStep -from agent_env.task.episode_rules import CHECKS, UNAFFECTED +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 f506e4c9f44a29b6554c8354f4432c24dff14096 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Tue, 13 Oct 2026 22:11:25 -0700 Subject: [PATCH 15/17] test(task): a seeded episode on local Docker, and one forgotten past its turn's deadline The episode integration test now runs a seeded task (the env, the role, the prompt and the verifier's id are placeholders) and checks the seed reaches the turn and the scores stay keyed by step. It also forgets one episode's turn: the run fails on the deadline and nothing it deployed is left running. Co-Authored-By: Claude Opus 5.5 --- .../env/envs/server_env_local_test.py | 37 ++++++++++++------- 1 file changed, 24 insertions(+), 13 deletions(-) diff --git a/tst/integration/env/envs/server_env_local_test.py b/tst/integration/env/envs/server_env_local_test.py index 4b7d42ae..2481076c 100644 --- a/tst/integration/env/envs/server_env_local_test.py +++ b/tst/integration/env/envs/server_env_local_test.py @@ -210,29 +210,33 @@ async def verify(mcp_url): @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 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 - by the verifier; and every episode, played, not played or abandoned, leaving nothing behind.""" + """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") - spec = EpisodeSpec.from_task(Task.from_dict({"id": f"crate-{uid}", "steps": [ - {"id": "env", "type": "deploy_env", "env_id": env.id}, - {"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"}, + 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 crate."}, - {"id": "check", "type": "env_outcome_verifier", "env_id": env.id, "file_artifact_id": verify.id, - "verifier_id": "crate"}, - ]}), trainer_agents=["policy"]) + {"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: @@ -258,7 +262,14 @@ async def play(episode, add): 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-*")) == [] - for episode in (played, unplayed, abandoned): + + 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")] From 5964af337beaec67707f209c4c22958651c50a45 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Tue, 13 Oct 2026 22:21:03 -0700 Subject: [PATCH 16/17] test(task): an episode whose env the task deploys on Modal A loop on this machine plays the agent over the env's Modal tunnel as its role, and the verifier scores the env's state there; played, abandoned and forgotten-past-the-deadline episodes each leave their sandbox stopped, as Modal's API reports it. The env's image is a build context Modal builds, so no registry Modal can pull from is needed. Gated like the file's other Modal test. Co-Authored-By: Claude Opus 5.5 --- .../env/envs/server_env_modal_test.py | 79 +++++++++++++++++++ 1 file changed, 79 insertions(+) diff --git a/tst/integration/env/envs/server_env_modal_test.py b/tst/integration/env/envs/server_env_modal_test.py index 9895b746..79aa4db5 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,82 @@ 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) + + played = await spec.start() + 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 spec.start() + await abandoned.next_turn() + await abandoned.abandon("test") + assert Task.get_instance(abandoned.id).status == "cancelled" + + forgotten = await EpisodeSpec.from_task(task, turn_timeout=3).start() + 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) + + +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" From 3dcb0b08e2beaf5c7015e4a71a3a496f74bae3ac Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Tue, 13 Oct 2026 22:46:05 -0700 Subject: [PATCH 17/17] test(task): abandon every episode the Modal test started, however it ends A failed call or assertion before an episode ended left its Modal sandbox running until the deploy's TTL; the test now abandons each episode it started on the way out, which does nothing to one that already ended. Co-Authored-By: Claude Opus 5.5 --- .../env/envs/server_env_modal_test.py | 50 +++++++++++-------- 1 file changed, 29 insertions(+), 21 deletions(-) diff --git a/tst/integration/env/envs/server_env_modal_test.py b/tst/integration/env/envs/server_env_modal_test.py index 79aa4db5..ed7ef874 100644 --- a/tst/integration/env/envs/server_env_modal_test.py +++ b/tst/integration/env/envs/server_env_modal_test.py @@ -89,30 +89,38 @@ async def test_an_episodes_env_on_modal_is_played_from_this_machine_and_no_endin "verifier_id": "crate"}, ]}) spec = EpisodeSpec.from_task(task) + started = [] - played = await spec.start() - 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 spec.start() - await abandoned.next_turn() - await abandoned.abandon("test") - assert Task.get_instance(abandoned.id).status == "cancelled" - - forgotten = await EpisodeSpec.from_task(task, turn_timeout=3).start() - await forgotten.next_turn() - with pytest.raises(TimeoutError, match="wasn't ended within 3s"): - await forgotten.next_turn() + async def start(spec): + started.append(episode := await spec.start()) + return episode - 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) + 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)