fix(openenv): a scored episode answers at once, and a role-less agent's toolset keys the default role's access

src/agent_env/a2a_agent/protocol.py

@@ -2,6 +2,7 @@
 from __future__ import annotations
 
 import asyncio
+import contextlib
 import logging
 import re
 import time
@@ -25,6 +26,12 @@
 # What a proxy in front of the agent answers when it can't reach it; any other status is the agent answering.
 _GATEWAY_STATUSES = frozenset({502, 503, 504})
 _CANCEL_TIMEOUT_SECONDS = 10
+# The A2A tasks ``poll_a2a_task`` is polling in this process, or that ``wake_poll`` woke before their poll began, each
+# with the event ``wake_poll`` sets. A task its waker just ended can show as working for a moment longer, so a woken
+# poll asks again every _WOKEN_POLL_SECONDS for up to _WOKEN_FOR_SECONDS before going back to its interval.
+_polls: dict[str, asyncio.Event] = {}
+_WOKEN_POLL_SECONDS = 0.05
+_WOKEN_FOR_SECONDS = 2.0
 
 
 class AgentUnreachableError(TimeoutError):
@@ -150,18 +157,42 @@ async def poll_a2a_task(
     *,
     sandbox_id: Optional[str] = None,
 ) -> dict:
-    """POST /a2a tasks/get until the task reaches a terminal state, and return it.
+    """POST /a2a tasks/get until the task reaches a terminal state, and return it. ``wake_poll(task_id)`` has it ask
+    at once rather than at its next interval.
 
     Out of time, it raises ``TimeoutError``. ``sandbox_id`` names the sandbox the agent runs on: then it raises
     ``AgentUnreachableError``, naming it, once the agent has gone UNREACHABLE_AFTER_SECONDS without answering,
     rather than waiting out the timeout. Either way it first asks the agent to cancel the task, so an agent still
     working on it stops."""
+    woken = _polls.setdefault(task_id, asyncio.Event())
+    try:
+        return await _poll(a2a_url, task_id, timeout_seconds, poll_interval_seconds, sandbox_id, woken)
+    finally:
+        if _polls.get(task_id) is woken:
+            del _polls[task_id]
+
+
+def wake_poll(task_id: str) -> None:
+    """Have this process's poll of ``task_id`` ask for its state now, or as it begins: for whoever just ended the task
+    here, as an episode ending a turn its proxy holds does, so the step waiting on it goes on without a poll interval."""
+    _polls.setdefault(task_id, asyncio.Event()).set()
+
+
+async def _poll(
+    a2a_url: str, task_id: str, timeout_seconds: int, poll_interval_seconds: int, sandbox_id: Optional[str],
+    woken: asyncio.Event,
+) -> dict:
     deadline = time.monotonic() + timeout_seconds
     consecutive_failures = 0
     answered_at, unanswered, last_failure = time.monotonic(), 0, ""
+    quick_until = 0.0
     while time.monotonic() < deadline:
         backoff = min(poll_interval_seconds + consecutive_failures * 5, 60)
-        await asyncio.sleep(backoff)
+        if time.monotonic() < quick_until:
+            backoff = min(backoff, _WOKEN_POLL_SECONDS)
+        if await _nap(woken, backoff):
+            quick_until = time.monotonic() + _WOKEN_FOR_SECONDS
+        woken.clear()
         try:
             async with httpx.AsyncClient() as client:
                 resp = await client.post(f"{a2a_url}/a2a", json={
@@ -204,6 +235,15 @@ async def poll_a2a_task(
     raise TimeoutError(f"A2A task {task_id} did not complete within {timeout_seconds}s")
 
 
+async def _nap(woken: asyncio.Event, seconds: float) -> bool:
+    """Whether ``woken`` was set before ``seconds`` passed."""
+    with contextlib.suppress(TimeoutError):
+        async with asyncio.timeout(seconds):
+            await woken.wait()
+            return True
+    return False
+
+
 async def cancel_a2a_task(a2a_url: str, task_id: str) -> None:
     """POST /a2a tasks/cancel, so the agent stops working on ``task_id``. Best effort: an agent that can't be
     reached, or won't cancel, is only logged."""

src/agent_env/openenv/environment.py

@@ -25,6 +25,7 @@
 )
 from openenv.core.env_server.types import Observation, State
 
+from agent_env.env.gateway.constants import DEFAULT_ROLE
 from agent_env.task import EpisodeError, EpisodeSpec, Task
 from agent_env.task.episode import Episode, Turn
 from agent_env.task.episode.rules import DEPLOYS, works_in_sandbox
@@ -75,7 +76,7 @@ def toolset(self) -> Toolset:
                                  if hasattr(step, "env_version") and step.env_id in envs},
                        access=[{"env": step.env_id, "action": step.action, "tools": step.tools}
                                for step in self.task.steps
-                               if step.type == "modify_env_tool_access" and step.role == role],
+                               if step.type == "modify_env_tool_access" and step.role == (role or DEFAULT_ROLE)],
                        bash=works_in_sandbox(self.task.steps, agent))
 
 
@@ -253,15 +254,15 @@ async def _end_turn(self, name: str, response: str, timeout_s: float | None) ->
                     await self._open(turn)
                     return CallToolObservation(tool_name=name, done=False, reward=None, result={
                         "prompt": turn.prompt, "system_prompt": turn.system_prompt})
-                scores = await self._episode.result()
+                scores = await self._episode.result(teardown=False)
         except Exception as e:
             if deadline.expired():
                 observation = await self._end_with_error(f"{name} didn't finish within {timeout_s:g}s, and the episode "
                                                          "ended", wait=False)
             else:
                 observation = await self._end_with_error(e)
             return CallToolObservation(tool_name=name, done=True, reward=None, metadata=observation.metadata)
-        await self._abandon("the episode ended")
+        await self._abandon("the episode ended", wait=False)
         return CallToolObservation(tool_name=name, done=True, **self._scored(scores))
 
     def _scored(self, scores: dict[str, float]) -> dict[str, Any]:

src/agent_env/task/episode/run.py

@@ -15,6 +15,7 @@
     ProxyToolAction,
 )
 
+from agent_env.a2a_agent import protocol
 from agent_env.a2a_agent.a2a_agent import A2AAgent
 from agent_env.providers.sandbox_providers.sandbox_provider import sandbox_request_headers_for_url
 from agent_env.store.routing import namespace_routing, run_scope
@@ -81,6 +82,7 @@ async def end(self, response: str | None = None) -> None:
             return
         try:
             await self._checked(self._proxy.reply(self._task_id, parts=[ProxyTextPart(text=response or "")]))
+            protocol.wake_poll(self._task_id)
         except httpx.HTTPError as e:
             logger.warning("turn %r of agent %r: its reply didn't reach the proxy: %s", self.step_id, self.agent, e)
         self._open = False
@@ -90,6 +92,7 @@ async def _fail(self, error: Exception) -> None:
             self._open = False
             try:
                 await self._proxy.reply(self._task_id, error=str(error))
+                protocol.wake_poll(self._task_id)
             except httpx.HTTPError as e:
                 logger.warning("turn %r of agent %r: its failure didn't reach the proxy: %s", self.step_id, self.agent,
                                e)
@@ -167,6 +170,7 @@ def __init__(self, id: str, task: Task, context: TaskStepContext):
         self._waiting: deque[Turn] = deque()
         self._turns: set[Turn] = set()
         self._changed: asyncio.Future[None] = asyncio.get_running_loop().create_future()
+        self._ran = False
         self._tearing_down = False
         self._abandoned: str | None = None
         self._proxies: list[ProxyClient] = []
@@ -180,8 +184,8 @@ def finished(self) -> bool:
         return self._run.done()
 
     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."""
+        """The next outside turn the run reaches, or None once the run's steps have finished, its teardown aside.
+        Raises the run's error if it failed first. End the turns you hold before waiting on one that depends on them."""
         while True:
             changed = self._changed
             self._close_turns_of_failed_steps()
@@ -192,20 +196,24 @@ async def next_turn(self) -> Turn | None:
             if self._run.done():
                 self._raise_unless_completed()
                 return None
+            if self._ran:
+                return None
             await self._wait(changed)
 
-    async def result(self) -> dict[str, float]:
+    async def result(self, *, teardown: bool = True) -> 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():
+        over: what ``agent-env run`` reports. Making a reward of them is the caller's. Without ``teardown``, once the
+        run's steps have finished; the teardown goes on, and ``abandon()`` waits for it."""
+        while not (self._run.done() or (self._ran and not teardown)):
             changed = self._changed
             self._close_turns_of_failed_steps()
             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()
+        if self._run.done():
+            self._raise_unless_completed()
         return recorded_scores(self._seeded_steps(), self._context)
 
     async def abandon(self, reason: str) -> None:
@@ -243,6 +251,8 @@ async def _drive(self, running: asyncio.Event) -> None:
         with namespace_routing():
             try:
                 await self._task.run(context=self._context, instance_id=self.id)
+                self._ran = True
+                self._notify()
             finally:
                 self._tearing_down = True
                 try:
@@ -323,8 +333,11 @@ async def _drain(self, agent: str, proxy: ProxyClient) -> None:
                     )
                     self._waiting.append(turn)
                     self._turns.add(turn)
-            self._changed.set_result(None)
-            self._changed = asyncio.get_running_loop().create_future()
+            self._notify()
+
+    def _notify(self) -> None:
+        self._changed.set_result(None)
+        self._changed = asyncio.get_running_loop().create_future()
 
     def _step_of(self, agent: str, context_id: str) -> str:
         """The ``prompt_agent`` step a message to ``agent`` on ``context_id`` comes from: of its prompts in that

tst/integration/openenv/build_app_test.py

@@ -236,7 +236,7 @@ async def test_each_session_holds_its_own_deployment_and_every_way_it_ends_tears
     with _serving(build_app(str(root), max_concurrent_envs=2)) as url:
         rewards = await asyncio.gather(_play(url, add=True), _play(url, add=False))
         assert [ended.reward for _, ended in rewards] == [1.0, 0.0]
-        assert _left(items.sandboxes) == []
+        await _until(lambda: _left(items.sandboxes) == [])
 
         held = GenericEnvClient(base_url=url, message_timeout_s=300)
         await held.connect()

tst/unit/a2a_agent/protocol_test.py

@@ -4,6 +4,7 @@
 telemetry and — when the agent ran with output_format — the typed structured_output.
 """
 
+import asyncio
 import json
 import logging
 from types import SimpleNamespace
@@ -19,6 +20,7 @@
     extract_terminal_response,
     poll_a2a_task,
     raise_for_extension_status,
+    wake_poll,
 )
 from agentenv_protocol.a2a_agent import TaskResult, Usage
 from agentenv_protocol.a2a_agent.framework import _to_a2a_parts
@@ -205,7 +207,8 @@ def __init__(self, monkeypatch, answer):
         client = httpx.AsyncClient
         monkeypatch.setattr(protocol.httpx, "AsyncClient", lambda **kw: client(transport=transport, **kw))
         monkeypatch.setattr(protocol, "time", SimpleNamespace(monotonic=lambda: self.now))
-        monkeypatch.setattr(protocol, "asyncio", SimpleNamespace(sleep=self._sleep))
+        monkeypatch.setattr(protocol, "asyncio", SimpleNamespace(sleep=self._sleep, Event=asyncio.Event))
+        monkeypatch.setattr(protocol, "_nap", lambda woken, seconds: self._sleep(seconds))
 
     async def _sleep(self, seconds):
         self.now += seconds
@@ -316,6 +319,38 @@ async def test_an_agent_answering_with_errors_is_waited_on_until_the_timeout(mon
     assert agent.now >= 600
 
 
+@pytest.mark.asyncio
+async def test_a_poll_woken_for_its_task_asks_at_once_and_again_until_it_shows_as_ended(monkeypatch):
+    asked = []
+
+    def answer(request):
+        asked.append(json.loads(request.content)["params"]["id"])
+        return _task("completed" if len(asked) > 1 else "working")
+
+    client = httpx.AsyncClient
+    monkeypatch.setattr(protocol.httpx, "AsyncClient", lambda **kw: client(transport=httpx.MockTransport(answer), **kw))
+    polling = asyncio.create_task(poll_a2a_task("http://agent", "t-1", 600, 60))
+    await asyncio.sleep(0)
+
+    wake_poll("t-1")
+
+    assert (await asyncio.wait_for(polling, 1))["status"]["state"] == "completed"
+    assert asked == ["t-1", "t-1"]
+    assert "t-1" not in protocol._polls
+
+
+@pytest.mark.asyncio
+async def test_a_wake_before_the_poll_begins_has_it_ask_at_once(monkeypatch):
+    client = httpx.AsyncClient
+    monkeypatch.setattr(protocol.httpx, "AsyncClient",
+                        lambda **kw: client(transport=httpx.MockTransport(lambda request: _task("completed")), **kw))
+
+    wake_poll("t-1")
+
+    assert (await asyncio.wait_for(poll_a2a_task("http://agent", "t-1", 600, 60), 1))["status"]["state"] == "completed"
+    assert "t-1" not in protocol._polls
+
+
 @pytest.mark.asyncio
 async def test_without_a_sandbox_a_silent_agent_is_waited_on_until_the_timeout(monkeypatch):
     agent = _Agent(monkeypatch, lambda method, request, agent: _refused(request, agent))

tst/unit/openenv/environment_test.py

@@ -97,6 +97,11 @@ async def test_the_task_api_lists_each_split_with_each_tasks_prompt_and_agent(te
     hidden = _task(ITEMS, ALICE, hide, {**hide, "id": "other", "role": "bob@example.com"}, TURN, REPLIES)
     assert ServedTask(hidden, "hidden", EpisodeSpec.from_task(hidden), _score_at("replies")).toolset["access"] == [
         {"env": "items", "action": "disable", "tools": ["add_item"]}]
+    roleless = {key: value for key, value in ALICE.items() if key != "role"}
+    defaulted = _task(ITEMS, roleless, {**hide, "role": "default"}, TURN, REPLIES)
+    assert ServedTask(defaulted, "defaulted", EpisodeSpec.from_task(defaulted), _score_at("replies")).toolset == {
+        "envs": ["items"], "versions": {}, "role": None, "access": [{"env": "items", "action": "disable",
+                                                                      "tools": ["add_item"]}], "bash": False}
     assert [task["index"] for task in env.get_task_range("all", 1)] == [1, 2]
     assert [task["id"] for task in env.list_tasks("smoke")] == ["dock-0"]
     with pytest.raises(IndexError, match="split 'all' has 3 tasks, so it has no task 3"):
@@ -109,9 +114,10 @@ async def test_the_task_api_lists_each_split_with_each_tasks_prompt_and_agent(te
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize(("reply", "reward"), [("Added a crate.", 1.0), ("", 0.0)])
-async def test_an_episode_resets_to_its_turn_steps_through_the_envs_tools_and_is_scored_on_end_turn(
+async def test_an_episode_resets_to_its_turn_steps_through_the_envs_tools_and_is_scored_on_end_turn_before_teardown(
         teardowns, tools, reply, reward):
     env = _env(ITEMS, ALICE, TURN, REPLIES)
+    teardowns.hold = asyncio.Event()
 
     reset = await env.reset_async(split="all", index=0)
     assert (reset.done, reset.reward) == (False, None)
@@ -135,7 +141,10 @@ async def test_an_episode_resets_to_its_turn_steps_through_the_envs_tools_and_is
     after = await _call(env, "add_item")
     assert (after.done, after.reward, after.metadata["status"], after.metadata["error"]) == (
         True, None, "error", ENDED)
-    assert teardowns.finished == [env.state.episode_id]
+    await _until(lambda: teardowns.started == [env.state.episode_id])
+    assert teardowns.finished == []
+    teardowns.hold.set()
+    await _until(lambda: teardowns.finished == [env.state.episode_id])
 
 
 @pytest.mark.asyncio
@@ -267,7 +276,7 @@ async def over(response=None):
 
 
 @pytest.mark.asyncio
-async def test_a_teardown_that_fails_ends_the_episode_in_an_error_observation_and_frees_the_session(
+async def test_a_teardown_that_fails_after_the_episode_is_scored_leaves_its_reward_and_frees_the_session(
         teardowns, tools, monkeypatch):
     async def failing(context, **_):
         raise RuntimeError("the sandbox didn't stop")
@@ -279,8 +288,7 @@ async def failing(context, **_):
     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 (ended.done, ended.reward, ended.metadata["scores"]) == (True, 1.0, {"replies": 1.0})
     assert (env._served.live, list(episode_run._live)) == (set(), [])
 
 

tst/unit/task/episode/run_test.py

@@ -198,17 +198,19 @@ async def _until(holds):
 
 
 @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()
+async def test_the_run_waits_at_an_outside_turn_and_carries_on_as_soon_as_it_ends(teardowns):
+    episode = await _spec(ITEMS, ALICE, {**TURN, "poll_interval_seconds": 30}, 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) == (
         "alice", "turn", "Add a crate.", "You keep the items.")
     assert not episode.finished
 
+    ending = time.monotonic()
     await turn.end("Added a crate.")
     assert await episode.result() == {"replies": 1.0}
+    assert time.monotonic() - ending < 2
     assert episode.finished
     assert await episode.next_turn() is None
     assert teardowns.finished == [episode.id]
@@ -471,11 +473,13 @@ async def test_a_turn_waiting_when_the_episode_is_abandoned_is_never_handed_out(
 
 
 @pytest.mark.asyncio
-async def test_abandoning_during_teardown_waits_for_it_and_leaves_the_runs_outcome(teardowns):
+async def test_during_teardown_the_scores_are_ready_and_abandoning_waits_for_it_leaving_the_runs_outcome(teardowns):
     teardowns.hold = asyncio.Event()
     episode = await _spec(ITEMS, ALICE, TURN, REPLIES).start()
     await (await episode.next_turn()).end("done")
+    assert await episode.result(teardown=False) == {"replies": 1.0}
     await _until(lambda: teardowns.started)
+    assert not episode.finished
 
     abandoning = asyncio.create_task(episode.abandon("too late"))
     await asyncio.sleep(0.05)

tst/unit/training/trl_test.py

@@ -84,6 +84,7 @@ def test_a_rollout_plays_its_rows_task_and_an_environment_plays_one_after_anothe
     environment.reset(**row)
     assert environment.get_reward() == 0.0
     assert environment.get_reward() is None
+    env.close()
     assert len(teardowns.finished) == 3