feat(task): run a task with some of its agents played from your own loop

[earakely-scale]

What

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 everything else exactly as agent-env run does (env deploys, setup, the other agents, scoring), so the reward is the score agent-env run would give. It's the core piece a trainer adapter (TRL's GRPOTrainer first) builds on.

spec = EpisodeSpec.from_task(task, agents=None, seed=None, turn_timeout=None)   # reads the task, runs nothing
spec.agents          # {"default-agent": AgentSpec(role, turns=(TurnSpec(step_id, prompt, system_prompt),))}
episode = await spec.start(run_group_id=None)                # returns once the run's record exists
turn = await episode.next_turn()                             # Turn, or None once the run has finished
turn.prompt, turn.system_prompt, turn.role, turn.mcp_urls, turn.headers
await turn.end(response="...")
reward = await episode.result()                              # the one top-level verifier's score
await episode.abandon("a tool call failed")                  # recorded cancelled

How it works

• from_task builds a copy of the task in which each outside agent's deploy_agent becomes episode_agent (deploys nothing; binds the agent to its envs with deploy_agent's own rule) and each of its prompt_agent steps becomes episode_turn (hands the turn to the caller, waits for end() within the prompt's timeout_seconds, records the reply as a PromptResponse). Ids, edges and fail_task_on_error are kept; every other step is untouched.
• start() runs one Task.run() of that copy as a task on the caller's loop, tagged context.metadata["episode"]. However the run ends, teardown_run takes down what it deployed before any caller sees the end.
• Callers wait with asyncio.wait, never by awaiting the run, so cancelling a caller of next_turn(), result() or abandon() never reaches the run. A cancelled start() takes its run down with it, and returns once that run is torn down.
• A turn's deadline is its prompt's timeout_seconds, or turn_timeout for every outside turn: a prompt budgets an agent's run, and the loop playing it may need its own. EpisodeError is a RuntimeError.

What from_task refuses (v1)

The rules are in task/episode_rules.py. Each built-in step type is either checked there (CHECKS, type to check) or listed as untouched by an outside agent (UNAFFECTED). A unit test fails when a registered type is in neither, so a new step type has to be classified before an episode can run it. Plugin step types aren't inspected.

It reports every problem at once, naming the step:

• an outside agent without exactly one deploy_agent, or deployed by install_agent / deploy_human_agent;
• an outside prompt with parts, a user simulator, output_format, agentenv_tools or snapshot_config;
• an outside deploy_agent with skills or agent state (snapshot files, changelog);
• steps that need the outside agent to be a real A2A agent: add_skills, register_agent_triggers, snapshot_agent_state, load_artifact / verify_sandbox / run_code into its sandbox, collect_artifacts from it, a trigger executor, peer_agents, a rubric that reads its trajectory or judges in its sandbox, a user simulator played from outside, and verify_a2a_* validators;
• any retry_config, since a rollback could replay an outside turn;
• anything but exactly one top-level verifier (the reward).

Deliberately not handled, with no defensive code: user simulators, trajectory-based rubric judging and peer tools aimed at an outside agent, retries, and per-agent rewards.

Other changes

• substitute_seed_in_steps rebuilds each step with its own class (type(step).from_dict) instead of looking it up in the registry. That's the same class for every registered type, and its docstring already said so; the episode's unregistered stand-ins can now take a seed.
• deploy_agent's env binding rule is a module function, bound_deployment(context, env_id, *, env_step_id, step_id), shared by deploy_agent and episode_agent, and live_mcp_url is public for the same reason. Their behaviour is unchanged.
• Record ids (env, task, A2A agent and env-state instances) and a fresh state root's first host port are drawn from secrets, not the global random. A training loop reseeds that, and ids repeated: a GRPO run lost about one env-instance record in ten to a swallowed DuplicateKeyError.
• local_stack in server_env_local_test.py now gives each test its own state root, removes every container its sandbox root names, and closes the grant server, as the task_step/*_local fixtures do.
• README: a short "Run a task from your own loop" section. AGENTS.md: the episode in the Task.run() embedder contract.

Tests

• tst/unit/task/episode_spec_test.py, 52 tests on from_task: every refusal, table-driven; the reward pick; seeded turns and a seeded agent name; the copy; the registry held to CHECKS and UNAFFECTED.
• tst/unit/task/episode_run_test.py, 25 tests on the runtime, on fake steps: turns, two outside agents, 8 concurrent episodes, deadlines, the reply graded by a real agent_prompt_response_verifier, abandon at every stage, cancelled callers and start(), and records.
• They pass on Python 3.11, 3.12 and 3.14, and repeatedly under -n 6. The full unit tier passes.
• tst/integration/env/envs/server_env_local_test.py (fast tier, local Docker), with the items env and the echo agent:
  • a played episode (ToolSession as the agent's role) scores 1.0 side by side with an unplayed one at 0.0;
  • an abandoned one is recorded cancelled;
  • no agent is deployed for the outside player, and nothing is left behind after each ending.
  • Three other fast-tier tests already deploy the echo agent, so it stays in the fast tier.

Validated end to end

A prototype adapter drove TRL 1.15's GRPOTrainer (environment_factory, vLLM colocate) over this API on PortSim's dock tasks, each rollout an episode deployed by the task's own deploy_env on local Docker:

• Zero-cost replay through the adapter (no model):
  • the optimal plan scores 1.0, a naive one 0.20, and no plan 0.0;
  • a broken connection leaves the rollout unscored, with its episode abandoned.
• 20 steps of 16 rollouts:
  • 320 episodes, all scored;
  • 1,660 tool calls, none failed;
  • about 10 s from start to turn with 16 at once;
  • nothing left running.

That run motivated turn_timeout and found the id collisions.

🤖 Generated with Claude Code

Confidence Score: 5/5

The PR appears safe to merge; the previous test finding is fixed and no new blocking issue remains.

Summary

This PR adds episodes that let a caller play selected agents while AgentEnv runs the remaining steps, scoring, and cleanup.

• Tasks can hand selected agents’ turns to a caller.
• Callers play task turns and receive the verifier’s score.
• Generated IDs and fresh sandbox ports ignore Python’s random seed.
• Agent and episode steps share the same environment-binding rules.

Diagram

sequenceDiagram
  participant Caller
  participant Episode
  participant Task
  Caller->>Episode: start()
  Episode->>Task: Run task copy
  Task->>Episode: Open outside turn
  Caller->>Episode: next_turn()
  Episode-->>Caller: Prompt and environment access
  Caller->>Episode: turn.end(response)
  Episode-->>Task: Record response and continue
  Task-->>Episode: Finish scoring
  Episode->>Episode: Tear down deployed resources
  Caller->>Episode: result()
  Episode-->>Caller: Reward

Loading

_{Reviews (4) · Last reviewed commit: "test(task): hold only the built-in step ..." · Reviewed by Greptile}