fix: reconnect when the WebSocket was closed by the far end

[sergiopaniego]

Summary

EnvClient keeps one persistent WebSocket. _connect_async only dropped the cached socket when it came from a different event loop, so a socket closed by the far end stayed cached: self._ws is still not None, and every later reset/step raised ConnectionClosed for the rest of the process's life. This drops the socket when its state is CLOSING or CLOSED, and has _receive ensure the connection instead of asserting it exists.

Hit this training against two env Spaces: a transient network blip closed both clients, and neither recovered. The Spaces answered POST /reset in 0.3s from outside while the jobs could not reach them for the rest of the run. Two runs lost, ~4.5 GPU-hours burnt on groups that scored zero.

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_reconnects_after_close closes the socket from under the client and asserts the next call still works. It fails on main with ConnectionClosedOK and passes with the fix.

tests/envs/test_websockets.py    19 passed, 2 skipped
tests/test_core/                281 passed, 11 skipped
ruff check                       All checks passed

Only a demonstrably closed socket is dropped, so one still CONNECTING is left alone and the existing same-loop no-op behaviour is unchanged.

Claude Code Review

N/A

---

Note

Medium Risk
Touches persistent WebSocket lifecycle for all env clients; behavior change is narrowly scoped to demonstrably closed sockets, but long-running training jobs depend on this path.

Overview
EnvClient no longer treats a dead WebSocket as a live session. When _connect_async sees a cached socket in CLOSING or CLOSED, it clears _ws and reconnects instead of reusing it and failing every later call with ConnectionClosed. Sockets still CONNECTING or on the same event loop keep the existing short-circuit behavior.

_receive now calls _ensure_connected() (like _send) so a receive path can trigger that reconnect after the far end drops the connection.

Adds integration test test_protocol_client_reconnects_after_close: force-close the socket under a sync client and assert reset / call_tool still work.

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

[bot-ci-comment]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.