[sergiopaniego]

Summary

Since #1103, EnvClient replaces a WebSocket closed by the far end and reconnects on the next call. The server keeps one environment per WebSocket (websocket_endpoint calls _create_session() for each connection and destroys it on disconnect), so the new socket is a new session with an environment this client never reset. If the socket closes in the middle of an episode, the next step() runs on that fresh environment and returns observations from another game, with no error.

The client now remembers the socket that reset() used. If that socket was closed or replaced when step() or state() is called, it raises SessionLostError (a ConnectionError), and keeps raising on every step() or state() until reset(), which reconnects as before and never raises it. Otherwise a caller that retries the step (a trainer returning the error to the model as a tool result, for example) would land on the new session again.

MCP clients in production mode reattach the new socket to the same HTTP MCP session (MCPClientBase._connect_async), which outlives the WebSocket, so the episode carries over and they are left as is.

How we hit it

GRPO with TRL environment_factory against a private Sudoku Space (textarena_env, MAX_CONCURRENT_ENVS=64), 64 clients that live for the whole run, one reset() and 15 to 25 step() calls per game.

The Space closed all 64 WebSockets about every 2 hours (server logs: 64 connection closed, then 64 new WebSocket /ws [accepted], with no error in the app). When a close landed while games were in progress, the model kept playing the board in its context while the server applied the moves to another board. That step: 0% solved and 7 moves per game, against 39% and 17 moves in the steps around it. Steps 74 and 75 were hit the same way later in the run. When a close landed between games nothing happened, because the reconnect happened on reset(). Nothing in the client or trainer logs showed the reconnects. We found them in the Space logs.

Changes

• src/openenv/core/env_client.py: SessionLostError, an _episode_ws reference to the socket the episode started on (set by reset(), cleared on done, disconnect and close), and _raise_if_session_lost() called by step() and state() before sending.
• src/openenv/core/mcp_client.py: skip the check in production mode.
• src/openenv/core/__init__.py: export SessionLostError.
• tests/envs/test_websockets.py: two tests next to test_protocol_client_reconnects_after_close.

Type of Change

• Bug fix
• New feature
• Breaking change
• Documentation
• New environment
• Refactoring

Alignment Checklist

Before submitting, verify:

• I have read .claude/docs/PRINCIPLES.md and this PR aligns with our principles
• I have checked .claude/docs/INVARIANTS.md and no invariants are violated
• I have run /pre-submit-pr (or bash .claude/hooks/lint.sh and tests) and addressed all issues

RFC Status

• Not required (bug fix, docs, minor refactoring)
• RFC exists: #___
• RFC needed (will create before merge)

Test Plan

test_protocol_client_raises_when_session_lost_mid_episode closes the socket from under a GenericEnvClient in the middle of an episode, asserts the next step() raises SessionLostError, and a retry too, then checks that reset() and a step work on the new session. On main the step goes through on a new session and the test fails with DID NOT RAISE. test_protocol_client_keeps_production_session_after_close checks that an EchoEnv (production mode) keeps working after the same close. test_protocol_client_reconnects_after_close still passes.

tests/envs/test_websockets.py + tests/test_core/   610 passed, 16 skipped, 3 failed
ruff check, ruff format, usort                      clean on the changed files

The 3 failures (test_protocol_client_state, test_protocol_client_multiple_episodes, test_fastapi_command_available) fail the same way on main locally.

---

Note

Medium Risk
Changes runtime behavior when connections drop during an active episode: callers must handle SessionLostError and call reset(), but this prevents incorrect cross-session steps in long-running training clients.

Overview
Fixes silent mid-episode reconnects that could apply step() / state() to a brand-new server environment after the WebSocket dropped, because each connection gets its own session until reset().

EnvClient now tracks the socket used for the current episode (_episode_ws set after reset(), cleared on done or disconnect). If that socket is gone before step() or state(), it raises SessionLostError (exported from openenv.core) instead of auto-reconnecting, and keeps raising on retries until reset() starts a new episode. MCP clients in production mode skip this check because reconnect reuses the same HTTP MCP session.

Adds websocket tests for the raise-on-lost behavior and for production MCP continuity after a close.

^{Reviewed by Cursor Bugbot for commit d27431a. Bugbot is set up for automated code reviews on this repo. Configure here.}