fix(openenv): raise on a step before reset, and end an episode whose end_turn is refused in an observation

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

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)

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()

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):