fix(openenv): a reset or end_turn past its deadline ends the episode, and CI runs TRL

.github/workflows/trl.yml

@@ -0,0 +1,56 @@
+name: TRL contract
+
+# One GRPOTrainer step on CPU through agent_env.training.trl (tst/trl), against TRL 1.16 and the latest. TRL marks
+# environment_factory experimental, so a release can break the helper with nothing else in CI noticing. TRL 1.15 can't
+# train on CPU (its log-prob path is a Triton GPU kernel), so 1.16 is the oldest this job can run. It runs on pull requests that touch the helper, the OpenEnv server it trains against or the tier, weekly,
+# and on demand; it isn't a required check.
+
+on:
+  pull_request:
+    paths:
+      - "src/agent_env/training/**"
+      - "src/agent_env/openenv/**"
+      - "tst/trl/**"
+      - "tst/unit/training/**"
+      - "pyproject.toml"
+      - "uv.lock"
+      - ".github/workflows/trl.yml"
+  schedule:
+    - cron: "23 7 * * 1"
+  workflow_dispatch:
+
+permissions:
+  contents: read
+
+concurrency:
+  group: trl-${{ github.event_name == 'pull_request' && github.ref || github.sha }}
+  cancel-in-progress: ${{ github.event_name == 'pull_request' }}
+
+env:
+  PIP_DISABLE_PIP_VERSION_CHECK: "1"
+
+jobs:
+  grpo:
+    name: grpo (${{ matrix.trl }})
+    runs-on: ubuntu-latest
+    timeout-minutes: 20
+    strategy:
+      fail-fast: false
+      matrix:
+        trl: ["trl==1.16.*", "trl"]
+    steps:
+      - uses: actions/checkout@11d5960a326750d5838078e36cf38b85af677262 # v4.4.0
+      - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
+        with:
+          python-version: "3.12"
+      - name: Install from uv.lock, then TRL and CPU torch
+        run: |
+          python -m pip install --quiet -r .github/ci-requirements.txt
+          uv sync --locked --extra dev --python "$(command -v python)"
+          uv pip install --python .venv/bin/python --torch-backend cpu torch "${{ matrix.trl }}"
+          .venv/bin/python -c "import torch, transformers, trl; print('trl', trl.__version__, 'transformers', transformers.__version__, 'torch', torch.__version__)"
+      - name: TRL contract tier
+        run: .venv/bin/python -m pytest tst/trl -rA --timeout=900 --junitxml=junit-trl.xml
+      - name: No skips
+        if: always()
+        run: .venv/bin/python .github/scripts/check_skip_policy.py junit-trl.xml

AGENTS.md

@@ -44,6 +44,7 @@ or one of them only where it is used.
 | `make int-test-fast` | `tst/integration` minus `int_test_slow`, in parallel; needs Docker and a local OCI registry (`docker run -d -p 5000:5000 public.ecr.aws/docker/library/registry:2`) |
 | `make int-test-slow` | the `int_test_slow` tests, serially (they build the same Docker tags from module fixtures and race under xdist) |
 | `make installer-test` | `tst/installer`: `plugin add` / `remove` through the real pip, uv and pipx, offline against wheels built from the checkout, and the container user journey (`python:3.12-slim`, `--network none`) with the PEP 668 refusal; needs uv, pipx and Docker |
+| `make trl-test` | `tst/trl`: one `GRPOTrainer` step on CPU through `agent_env.training.trl`, with a tiny model from the Hub; needs trl and CPU torch, which the `dev` extra leaves out |
 | `make clean-install-test` | the `clean-install` CI gate: both distributions built as the release builds them, installed into a fresh venv from public PyPI with an allowlisted environment, and `agent-env run hello` run twice by name and checked through the local store; needs Python 3.11, uv and git |
 | `npm ci && npm run typecheck && npm run lint && npm run test:smoke` | in `src/agent_env/explorer/ui`: the explorer UI checks the `ui` job runs (it also runs `npm run build:static`); needs Node 22 |
 | `pytest tst/unit/env/store_test.py -v` | one file; add `--log-cli-level=DEBUG` for debug logs |
@@ -266,7 +267,7 @@ orchestrators call it the same way. What callers rely on:
 - Skips are declared capability gaps, decided at collection time: the gates in
   `tst/util/capabilities.py` produce `skipif` marks with the reason
   `agentenv-capability-missing: <name>` (`model_endpoint_configured`, `remote_sandbox`,
-  `default_a2a_agent`, `mcp_server_sources`, `gnu_stat`). CI rejects any other skip reason; a broken configuration must fail loudly, not skip.
+  `default_a2a_agent`, `mcp_server_sources`, `gnu_stat`, `trl_installed`). CI rejects any other skip reason; a broken configuration must fail loudly, not skip.
 - The suite is backend-agnostic: a test about one particular backend lives with that backend, not
   here. A custom store must pass `tst/store/{conformance,object_conformance,image_conformance,secret_conformance}.py`,
   the same kits the built-ins pass, and a custom `[runner]` must pass `tst/runner/conformance.py`, which
@@ -300,6 +301,9 @@ skips for the allowed capabilities pass, and the fast job allows no skips under
 touches the plugin CLI, `agent_env/plugins` or `agent_env/config`, the protocol package, the tier
 itself, `pyproject.toml`, `uv.lock`, the `Makefile` or the CI that runs the tier, and allows no
 skips; a `changes` job decides, and a skipped `installer` counts as passed.
+`.github/workflows/trl.yml` runs `tst/trl` against TRL 1.16 (1.15 can't train on CPU) and the latest, on pull
+requests that touch `agent_env/training`, `agent_env/openenv` or the tier, weekly, and on demand; it allows no skips
+and isn't a required check.
 `.github/workflows/plugin-api.yml` runs the `plugin-api` job on every pull
 request, title edits included, and on `main`: `.github/scripts/check_plugin_api.py` compares the
 plugin surface (its `BASES` and `USED`) between `HEAD^1` and `HEAD` with griffe, pinned in the `dev` extra, and fails on a break the

CONTRIBUTING.md

@@ -43,6 +43,7 @@ reinstalling.
 | `make int-test-slow` | tests that build images or deploy sandboxes, minutes each | Docker, and a sandbox backend for some |
 | `make clean-install-test` | both distributions built as the release builds them, installed into a fresh venv from public PyPI, and `agent-env run hello` run twice by name | Python 3.11, uv and git |
 | `make installer-test` | `plugin add` and `remove` through the real pip, uv and pipx, offline against wheels built from the checkout, and the new-user journey in containers with no network | uv, pipx and Docker; the first run downloads the dependencies once |
+| `make trl-test` | one TRL `GRPOTrainer` step on CPU through the training helper, with a tiny model | trl and CPU torch installed, and the Hugging Face Hub |
 
 CI runs the same tiers as the `unit`, `integration-local`, `integration-local-slow` and
 `integration-local-gateway` jobs (the last runs the slow tier's gateway tests beside the rest);

Makefile

@@ -1,4 +1,4 @@
-.PHONY: help install lint unit-test int-test int-test-fast int-test-slow installer-test clean-install-test test clean
+.PHONY: help install lint unit-test int-test int-test-fast int-test-slow installer-test trl-test clean-install-test test clean
 
 VENV := .venv
 VENV_BIN := $(abspath $(VENV))/bin
@@ -64,6 +64,11 @@ int-test-slow: ## Run only the @int_test_slow integration tests (serial to avoid
 installer-test: ## Run plugin add/remove through the real pip, uv and pipx, and the container journey (needs uv, pipx, Docker)
 	$(PYTHON) -m pytest tst/installer -n auto --timeout=900 $(PYTEST_ARGS)
 
+# The TRL contract tier trains one GRPOTrainer step on CPU through agent_env.training.trl, with a tiny model from the
+# Hub. trl and torch aren't in the dev extra: install them first, as .github/workflows/trl.yml does.
+trl-test: ## Run one GRPOTrainer step on CPU through the TRL helper (needs trl and CPU torch, and the Hub)
+	$(PYTHON) -m pytest tst/trl --timeout=900 $(PYTEST_ARGS)
+
 # The clean-install gate CI runs, on the same Python. The script gives every step an allowlisted
 # environment, so nothing in this shell reaches the install it checks. A failed run keeps its folder.
 clean-install-test: ## Build both distributions, install them clean from public PyPI and run hello twice (needs python3.11, uv, git)

src/agent_env/openenv/environment.py

@@ -36,8 +36,9 @@
     input_schema={"type": "object", "properties": {"response": {"type": "string", "description": "Your reply."}}},
 )
 ENDED = "the episode has ended; call reset"
-#: The key of a tool call's deadline in seconds in its action's metadata. A call that runs past it ends the episode, since
-#: cancelling the call doesn't stop what it does, and the session goes on.
+#: The key of a call's deadline in seconds in its action's metadata, ``end_turn``'s included; ``reset_async`` takes one as
+#: ``timeout_s``. A call that runs past it ends the episode, since cancelling the call doesn't stop what it does, and the
+#: session goes on; so no call holds a session longer than its client waits.
 TIMEOUT_METADATA = "timeout_s"
 NOT_RESET = "no episode has started; call reset before stepping"
 
@@ -137,13 +138,13 @@ def get_task_range(self, split: str, start: int | None = None, stop: int | None
     # The episode
 
     async def reset_async(self, seed: int | None = None, episode_id: str | None = None, *, split: str | None = None,
-                          index: int | None = None, **kwargs: Any) -> Observation:
+                          index: int | None = None, timeout_s: float | None = None, **kwargs: Any) -> Observation:
         """Plays task ``index`` of ``split``, or without either the first task served, as OpenEnv's ORS and verifiers
-        importers do."""
+        importers do, and ends the episode if its first turn hasn't opened ``timeout_s`` seconds in."""
         async with self._lock:
-            return await self._reset(split, index)
+            return await self._reset(split, index, timeout_s)
 
-    async def _reset(self, split: str | None, index: int | None) -> Observation:
+    async def _reset(self, split: str | None, index: int | None, timeout_s: float | None) -> Observation:
         await self._abandon("reset")
         self._loop = asyncio.get_running_loop()
         self._served_task, self._state = None, State(episode_id=None, step_count=0, done=False)
@@ -158,14 +159,19 @@ async def _reset(self, split: str | None, index: int | None) -> Observation:
         self._state = State(episode_id=None, step_count=0, task_id=self._served_task.task.id,
                             task_name=self._served_task.name, split=split, index=index, turn=None, done=False)
         self._served.live.add(self)
+        deadline = asyncio.timeout(timeout_s)
         try:
-            self._episode = await self._served_task.spec.start(sandbox=self._served.sandbox)
-            self._state.episode_id = self._episode.id
-            turn = await self._episode.next_turn()
-            if turn is None:
-                raise EpisodeError("the run ended before the agent's first turn")
-            await self._open(turn)
+            async with deadline:
+                self._episode = await self._served_task.spec.start(sandbox=self._served.sandbox)
+                self._state.episode_id = self._episode.id
+                turn = await self._episode.next_turn()
+                if turn is None:
+                    raise EpisodeError("the run ended before the agent's first turn")
+                await self._open(turn)
         except Exception as e:
+            if deadline.expired():
+                return await self._end_with_error(f"the episode didn't open its first turn within {timeout_s:g}s",
+                                                  wait=False)
             return await self._end_with_error(e)
         return Observation(done=False, reward=None, metadata={
             "prompt": turn.prompt, "system_prompt": turn.system_prompt, "task": self._task_metadata()})
@@ -189,9 +195,11 @@ async def _step(self, action: Any, timeout_s: float | None) -> Observation:
         if (stopped := await self._why_stopped()) is not None:
             return CallToolObservation(tool_name=action.tool_name, done=True, reward=None, metadata=stopped)
         self._state.step_count += 1
+        if timeout_s is None:
+            timeout_s = _deadline_of(action)
         if action.tool_name == END_TURN.name:
-            return await self._end_turn(action.tool_name, action.arguments.get("response") or "")
-        return await self._call(action, timeout_s if timeout_s is not None else _deadline_of(action))
+            return await self._end_turn(action.tool_name, action.arguments.get("response") or "", timeout_s)
+        return await self._call(action, timeout_s)
 
     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."""
@@ -225,7 +233,7 @@ async def _call(self, action: CallToolAction, timeout_s: float | None) -> CallTo
             return _tool_error(name, ToolErrorType.TOOL_NOT_FOUND, e)
         except TimeoutError:
             ended = await self._end_with_error(f"{name} didn't answer within {timeout_s:g}s, and the episode ended, "
-                                               "since what the call does may still be under way")
+                                               "since what the call does may still be under way", wait=False)
             return CallToolObservation(tool_name=name, done=True, reward=None, metadata=ended.metadata)
         except httpx.HTTPError as e:
             return _tool_error(name, ToolErrorType.TRANSPORT_ERROR, e)
@@ -234,18 +242,24 @@ async def _call(self, action: CallToolAction, timeout_s: float | None) -> CallTo
         return CallToolObservation(tool_name=name, result={
             "content": result.content, "structured_content": None, "meta": None, "data": result.text, "is_error": False})
 
-    async def _end_turn(self, name: str, response: str) -> CallToolObservation:
+    async def _end_turn(self, name: str, response: str, timeout_s: float | None) -> CallToolObservation:
+        deadline = asyncio.timeout(timeout_s)
         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)
-                return CallToolObservation(tool_name=name, done=False, reward=None, result={
-                    "prompt": turn.prompt, "system_prompt": turn.system_prompt})
-            scores = await self._episode.result()
+            async with deadline:
+                await self._turn.end(response)
+                await self._close_tools()
+                turn = await self._episode.next_turn()
+                if turn is not 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()
         except Exception as e:
-            observation = await self._end_with_error(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")
         return CallToolObservation(tool_name=name, done=True, **self._scored(scores))
@@ -282,9 +296,9 @@ async def _why_stopped(self) -> dict[str, Any] | None:
             error = _described(e)
         return (await self._end_with_error(error)).metadata
 
-    async def _end_with_error(self, error: Exception | str) -> Observation:
+    async def _end_with_error(self, error: Exception | str, *, wait: bool = True) -> Observation:
         message = error if isinstance(error, str) else _described(error)
-        await self._abandon(message)
+        await self._abandon(message, wait=wait)
         return self._error_observation(message)
 
     def _error_observation(self, error: str) -> Observation:
@@ -296,17 +310,18 @@ async def _abandon_when_free(self, reason: str) -> None:
         async with self._lock:
             await self._abandon(reason)
 
-    async def _abandon(self, reason: str) -> None:
-        """Abandon the episode and close the turn's tools, waiting for the teardown; the caller holds the lock. The
-        abandon is tracked before anything is awaited, so a server shutting down meanwhile waits for it too, and a
-        caller cancelled while it waits leaves it running for the next reset to wait for."""
+    async def _abandon(self, reason: str, *, wait: bool = True) -> None:
+        """Abandon the episode and close the turn's tools, waiting for the teardown unless ``wait`` is false, as a call
+        past its deadline answers at once; the caller holds the lock. The abandon is tracked before anything is
+        awaited, so a server shutting down meanwhile waits for it too, and one not waited for, or whose caller is
+        cancelled while it waits, is left running for the next reset or close to wait for."""
         episode, self._episode = self._episode, None
         if episode is not None:
             self._ending = self._served.end_later(episode.abandon(reason))
         self._served.live.discard(self)
         await self._close_tools()
         self._turn = None
-        if (ending := self._ending) is not None:
+        if wait and (ending := self._ending) is not None:
             try:
                 await asyncio.shield(ending)
             finally: