diff --git a/packages/agentenv-protocol/tests/test_computer_extensions.py b/packages/agentenv-protocol/tests/test_computer_extensions.py index a76b766f..cbd0a9c8 100644 --- a/packages/agentenv-protocol/tests/test_computer_extensions.py +++ b/packages/agentenv-protocol/tests/test_computer_extensions.py @@ -164,8 +164,13 @@ async def test_a_background_apps_output_is_drained_after_the_call_returns(monkey r = await http.post("/agentenv/ext/computer-exec/run", json=_run( argv=["sh", "-c", "echo now; (sleep 0.4; echo later) &"])) draining = len(linux._draining) - await asyncio.sleep(1.0) - assert r.json()["stdout"] == "now\n" and draining >= 1 and not linux._draining + + async def drained(): + while linux._draining: + await asyncio.sleep(0.05) + + await asyncio.wait_for(drained(), timeout=10) + assert r.json()["stdout"] == "now\n" and draining >= 1 @pytest.mark.asyncio diff --git a/src/agent_env/openenv/environment.py b/src/agent_env/openenv/environment.py index 0d979341..e1021a9e 100644 --- a/src/agent_env/openenv/environment.py +++ b/src/agent_env/openenv/environment.py @@ -1,6 +1,7 @@ """A bundle's tasks as an OpenEnv env: each session plays one episode of a task at a time (``EpisodeSpec``), as OpenEnv's own tool envs do. Reset deploys and opens the first turn, steps list and call the env's tools, and the -``end_turn`` tool hands the reply back; an episode that can't go on ends in an observation, never an exception.""" +``end_turn`` tool hands the reply back. An episode that can't go on ends in an observation, never an exception; a step +before the session's first reset raises, as OpenEnv's imported envs do, so a tool call never starts an episode.""" from __future__ import annotations @@ -34,6 +35,7 @@ input_schema={"type": "object", "properties": {"response": {"type": "string", "description": "Your reply."}}}, ) ENDED = "the episode has ended; call reset" +NOT_RESET = "no episode has started; call reset before stepping" @dataclass(frozen=True) @@ -146,6 +148,8 @@ async def step_async(self, action: Any, timeout_s: float | None = None, **kwargs return await self._step(action, timeout_s) async def _step(self, action: Any, timeout_s: float | None) -> Observation: + if self._loop is None: + raise RuntimeError(NOT_RESET) if isinstance(action, ListToolsAction): if (stopped := await self._why_stopped()) is not None: return ListToolsObservation(tools=[], done=True, reward=None, metadata=stopped) @@ -162,7 +166,7 @@ async def _step(self, action: Any, timeout_s: float | None) -> Observation: def step(self, action: Any, timeout_s: float | None = None, **kwargs: Any) -> Observation: """OpenEnv's web playground steps from a worker thread; the episode lives on the loop its reset ran on.""" if self._loop is None: - return asyncio.run(self.step_async(action, timeout_s)) + raise RuntimeError(NOT_RESET) return asyncio.run_coroutine_threadsafe(self.step_async(action, timeout_s), self._loop).result() @property @@ -199,9 +203,9 @@ async def _call(self, action: CallToolAction, timeout_s: float | None) -> CallTo "content": result.content, "structured_content": None, "meta": None, "data": result.text, "is_error": False}) async def _end_turn(self, name: str, response: str) -> CallToolObservation: - await self._turn.end(response) - await self._close_tools() try: + await self._turn.end(response) + await self._close_tools() turn = await self._episode.next_turn() if turn is not None: await self._open(turn) diff --git a/tst/unit/openenv/app_test.py b/tst/unit/openenv/app_test.py index 4c51cb87..d2d7aff4 100644 --- a/tst/unit/openenv/app_test.py +++ b/tst/unit/openenv/app_test.py @@ -1,6 +1,6 @@ """``build_app`` over bundle folders, served in process: the splits it serves and the tasks it leaves out, the reward -and bundles it refuses, an episode played over OpenEnv's own WebSocket, and namespace routing for each of the app's -runs.""" +and bundles it refuses, an episode played over OpenEnv's own WebSocket after a step before reset is refused, and +namespace routing for each of the app's runs.""" import json import logging @@ -12,6 +12,7 @@ from agent_env.bundle import BundleError from agent_env.openenv import build_app +from agent_env.openenv.environment import NOT_RESET from agent_env.store.routing import namespace_routing, namespace_routing_enabled from agent_env.task import Task from tst.unit.bundle._support import layout @@ -74,6 +75,8 @@ def test_an_episode_is_played_over_openenvs_websocket_and_the_session_closing_te with TestClient(app) as client: assert namespace_routing_enabled() with client.websocket_connect("/ws") as ws: + ws.send_json({"type": "step", "data": {"type": "list_tools"}}) + refused = ws.receive_json() ws.send_json({"type": "reset", "data": {"split": "all", "index": 0}}) reset = ws.receive_json()["data"] ws.send_json({"type": "step", "data": {"type": "call_tool", "tool_name": "end_turn", @@ -84,6 +87,7 @@ def test_an_episode_is_played_over_openenvs_websocket_and_the_session_closing_te ws.send_json({"type": "state"}) open_episode = ws.receive_json()["data"]["episode_id"] + assert (refused["type"], refused["data"]["message"]) == ("error", NOT_RESET) assert reset["observation"]["metadata"]["prompt"] == "Add a crate." assert (ended["done"], ended["reward"]) == (True, 1.0) assert not namespace_routing_enabled() diff --git a/tst/unit/openenv/environment_test.py b/tst/unit/openenv/environment_test.py index 9083bc48..547c4024 100644 --- a/tst/unit/openenv/environment_test.py +++ b/tst/unit/openenv/environment_test.py @@ -1,6 +1,6 @@ """A bundle task as an OpenEnv env, on the Episode API's fake steps: the Task API, an episode reset, stepped and scored, -its tool errors, every way it ends in an error observation, calls that overlap, and how a session's episode is torn -down.""" +its tool errors, every way it ends in an error observation, a step before any reset, calls that overlap, and how a +session's episode is torn down.""" import asyncio import threading @@ -13,8 +13,8 @@ import agent_env.task.episode.run as episode_run from agent_env.openenv.app import _score_at -from agent_env.openenv.environment import ENDED, AgentEnvEnvironment, Served, ServedTask -from agent_env.task import EpisodeSpec, Task +from agent_env.openenv.environment import ENDED, NOT_RESET, AgentEnvEnvironment, Served, ServedTask +from agent_env.task import EpisodeError, EpisodeSpec, Task from tst.unit.task.episode.run_test import ALICE, ITEMS, REPLIES, TURN, _gate, _until AGAIN = {**TURN, "id": "again", "prompt_id": "again", "prompt": "Add another.", "depends_on": ["turn"]} @@ -133,7 +133,8 @@ async def test_an_episode_resets_to_its_turn_steps_through_the_envs_tools_and_is @pytest.mark.asyncio async def test_a_step_from_a_worker_thread_runs_on_the_episodes_loop_as_openenvs_playground_steps(teardowns, tools): env = _env(ITEMS, ALICE, TURN, REPLIES) - assert (await asyncio.to_thread(env.step, ListToolsAction())).metadata["error"] == ENDED + with pytest.raises(RuntimeError, match=NOT_RESET): + await asyncio.to_thread(env.step, ListToolsAction()) await env.reset_async() listed = await asyncio.to_thread(env.step, ListToolsAction()) @@ -143,6 +144,19 @@ async def test_a_step_from_a_worker_thread_runs_on_the_episodes_loop_as_openenvs assert (ended.done, ended.reward) == (True, 1.0) +@pytest.mark.asyncio +@pytest.mark.parametrize("action", [ListToolsAction(), CallToolAction(tool_name="add_item", arguments={})]) +async def test_a_step_before_the_first_reset_raises_and_starts_no_episode(teardowns, tools, action): + env = _env(ITEMS, ALICE, TURN, REPLIES) + + with pytest.raises(RuntimeError, match=NOT_RESET): + await env.step_async(action) + + assert (teardowns.started, env._served.live, env.state.episode_id) == ([], set(), None) + await env.reset_async(split="all", index=1) + assert (await env.step_async(action)).metadata["error"] == ENDED + + @pytest.mark.asyncio async def test_a_task_with_two_turns_hands_over_the_next_prompt_then_the_reward(teardowns, tools): env = _env(ITEMS, ALICE, TURN, AGAIN, REPLIES) @@ -207,6 +221,40 @@ async def test_a_run_that_fails_mid_turn_ends_the_episode_at_the_next_step(teard assert (await _call(env, "add_item")).metadata["error"] == ENDED +@pytest.mark.asyncio +async def test_an_end_turn_the_proxy_refuses_ends_the_episode_in_an_error_observation(teardowns, tools, monkeypatch): + env = _env(ITEMS, ALICE, TURN, REPLIES) + await env.reset_async() + + async def over(response=None): + raise EpisodeError("turn 'turn' of agent 'alice' is over") + + monkeypatch.setattr(env._turn, "end", over) + ended = await _call(env, "end_turn", response="Added a crate.") + + assert (ended.done, ended.reward, ended.metadata["status"], ended.metadata["error"]) == ( + True, None, "error", "EpisodeError: turn 'turn' of agent 'alice' is over") + assert (Task.get_instance(env.state.episode_id).status, teardowns.finished) == ("cancelled", [env.state.episode_id]) + + +@pytest.mark.asyncio +async def test_a_teardown_that_fails_ends_the_episode_in_an_error_observation_and_frees_the_session( + teardowns, tools, monkeypatch): + async def failing(context, **_): + raise RuntimeError("the sandbox didn't stop") + + monkeypatch.setattr(episode_run, "teardown_run", failing) + env = _env(ITEMS, ALICE, TURN, REPLIES) + await env.reset_async() + ended = await _call(env, "end_turn", response="Added a crate.") + await env.reset_async() + await asyncio.to_thread(env.close) + + assert (ended.done, ended.reward, ended.metadata["status"]) == (True, None, "error") + assert ended.metadata["error"] == "RuntimeError: the sandbox didn't stop" + assert (env._served.live, list(episode_run._live)) == (set(), []) + + @pytest.mark.asyncio @pytest.mark.parametrize("tool", ["state", "end_turn"]) async def test_an_env_offering_a_name_openenv_reserves_is_refused_at_reset_and_torn_down(teardowns, tools, tool):