From 22ca1ab2a55960fc5255bd3b7d017cbb1c0f923a Mon Sep 17 00:00:00 2001 From: sergiopaniego Date: Fri, 30 Oct 2026 15:57:40 +0100 Subject: [PATCH 1/2] Raise SessionLostError when the WebSocket drops mid-episode Since #1103 the client replaces a socket closed by the far end. The server keeps one environment per WebSocket, so the replacement is a new session that this client never reset, and a step() after it ran on another game with no error. The client now remembers whether reset() opened an episode on the current socket and raises SessionLostError from step()/state() if that socket is gone. reset() still reconnects as before. MCP clients in production mode reattach to the same HTTP MCP session, so they keep the episode and are left as is. --- src/openenv/core/__init__.py | 4 ++- src/openenv/core/env_client.py | 46 +++++++++++++++++++++++++++++-- src/openenv/core/mcp_client.py | 6 +++++ tests/envs/test_websockets.py | 49 ++++++++++++++++++++++++++++++++++ 4 files changed, 102 insertions(+), 3 deletions(-) diff --git a/src/openenv/core/__init__.py b/src/openenv/core/__init__.py index b365f5ff73..1f62b60f6f 100644 --- a/src/openenv/core/__init__.py +++ b/src/openenv/core/__init__.py @@ -11,7 +11,7 @@ from .env_server import * # noqa: F403 if TYPE_CHECKING: - from .env_client import EnvClient + from .env_client import EnvClient, SessionLostError from .generic_client import GenericAction, GenericEnvClient from .llm_client import ( AnthropicClient, @@ -26,6 +26,7 @@ __all__ = [ "EnvClient", + "SessionLostError", "SyncEnvClient", "GenericEnvClient", "GenericAction", @@ -42,6 +43,7 @@ _LAZY_ATTRS = { "EnvClient": (".env_client", "EnvClient"), + "SessionLostError": (".env_client", "SessionLostError"), "SyncEnvClient": (".sync_client", "SyncEnvClient"), "GenericEnvClient": (".generic_client", "GenericEnvClient"), "GenericAction": (".generic_client", "GenericAction"), diff --git a/src/openenv/core/env_client.py b/src/openenv/core/env_client.py index 01abe1aeb9..cb2447a5dd 100644 --- a/src/openenv/core/env_client.py +++ b/src/openenv/core/env_client.py @@ -275,6 +275,15 @@ async def _best_effort_graceful_close(ws: ClientConnection) -> None: await _best_effort_close(ws) +class SessionLostError(ConnectionError): + """The connection to the server was lost in the middle of an episode. + + The server keeps one environment per WebSocket connection, so reconnecting + opens a new session with an environment this client never reset. Call + `reset()` to start a new episode. + """ + + class EnvClient(ABC, Generic[ActT, ObsT, StateT]): """ Async environment client for persistent sessions. @@ -386,6 +395,8 @@ def __init__( # Strong references for fire-and-forget socket closes (see _receive), # so the task isn't garbage-collected mid-close. Discarded on done. self._pending_close_tasks: set[asyncio.Task] = set() + # Whether the current socket carries an episode started by `reset()`. + self._episode_open = False if base_url is not None: self._set_base_url(base_url) @@ -653,6 +664,7 @@ async def _disconnect_async(self) -> None: # leave a stale socket cached for a later operation. self._ws = None self._ws_loop = None + self._episode_open = False same_loop = ws_loop is asyncio.get_running_loop() if same_loop: # Keep ownership of the handshake if this caller is cancelled. @@ -996,8 +1008,12 @@ async def _reset_async(self, **kwargs: Any) -> StepResult[ObsT]: "type": "reset", "data": kwargs, } + # A reset may reconnect: it starts a new episode on whatever session it lands on. + self._episode_open = False response = await self._send_and_receive(message) - return self._parse_result(response.get("data", {})) + result = self._parse_result(response.get("data", {})) + self._episode_open = not result.done + return result def step(self, action: ActT, **kwargs: Any) -> Any: return self._dispatch(lambda: self._step_async(action, **kwargs)) @@ -1019,8 +1035,12 @@ async def _step_async(self, action: ActT, **kwargs: Any) -> StepResult[ObsT]: "type": "step", "data": self._step_payload(action), } + self._raise_if_session_lost() response = await self._send_and_receive(message) - return self._parse_result(response.get("data", {})) + result = self._parse_result(response.get("data", {})) + if result.done: + self._episode_open = False + return result def state(self) -> Any: return self._dispatch(self._state_async) @@ -1033,9 +1053,31 @@ async def _state_async(self) -> StateT: State object with environment state information """ message = {"type": "state"} + self._raise_if_session_lost() response = await self._send_and_receive(message) return self._parse_state(response.get("data", {})) + def _raise_if_session_lost(self) -> None: + """Raise `SessionLostError` if the socket that carried the episode is gone. + + `_connect_async()` would replace it with a new socket, which the server + treats as a new session, and the call would run on an environment this + client never reset. + """ + if not self._episode_open: + return + ws = self._ws + if ( + ws is None + or ws.state in (State.CLOSING, State.CLOSED) + or self._ws_loop is not asyncio.get_running_loop() + ): + # Keep raising until reset(): a later step() would otherwise reconnect and run on a new session. + raise SessionLostError( + f"Lost the connection to {self._ws_url} in the middle of an episode. " + "Call reset() to start a new one." + ) + def close(self) -> Any: return self._dispatch(self._close_async) diff --git a/src/openenv/core/mcp_client.py b/src/openenv/core/mcp_client.py index 9002fabf9b..c7967a9262 100644 --- a/src/openenv/core/mcp_client.py +++ b/src/openenv/core/mcp_client.py @@ -255,6 +255,12 @@ async def _connect_async(self) -> EnvClient: return await super()._connect_async() + def _raise_if_session_lost(self) -> None: + # In production mode a reconnect attaches to the same HTTP MCP session, + # which outlives the WebSocket, so the episode carries over. + if not getattr(self, "use_production_mode", False): + super()._raise_if_session_lost() + async def _ensure_production_session(self) -> str: """Create and cache a persistent HTTP MCP session id if needed.""" async with self._production_session_lock: diff --git a/tests/envs/test_websockets.py b/tests/envs/test_websockets.py index dce2bbcb02..d1d6275889 100644 --- a/tests/envs/test_websockets.py +++ b/tests/envs/test_websockets.py @@ -339,6 +339,55 @@ def test_protocol_client_reconnects_after_close(self, echo_server): client.reset() assert client.call_tool("echo_message", message="after") == "after" + def test_protocol_client_raises_when_session_lost_mid_episode(self, echo_server): + """A step after the socket closed mid-episode must raise, not run on a new session.""" + from openenv.core import GenericEnvClient + from openenv.core.env_client import SessionLostError + + def echo(message): + return { + "type": "call_tool", + "tool_name": "echo_message", + "arguments": {"message": message}, + } + + # The server keeps one environment per WebSocket for this client. + with GenericEnvClient(base_url=echo_server).sync() as client: + client.reset() + client.step(echo("before")) + + inner = client._async + asyncio.run_coroutine_threadsafe(inner._ws.close(), client._loop).result( + timeout=10 + ) + + # A new socket would be a fresh environment that was never reset. + with pytest.raises(SessionLostError): + client.step(echo("lost")) + # Retrying the step does not reconnect into a new session either. + with pytest.raises(SessionLostError): + client.step(echo("lost again")) + + # The next episode reconnects as before. + client.reset() + client.step(echo("after")) + assert client.state()["step_count"] == 1 + + def test_protocol_client_keeps_production_session_after_close(self, echo_server): + """MCP clients reattach to the same HTTP MCP session, so the episode carries over.""" + from envs.echo_env.client import EchoEnv + + with EchoEnv(base_url=echo_server).sync() as client: + client.reset() + client.call_tool("echo_message", message="a") + + inner = client._async + asyncio.run_coroutine_threadsafe(inner._ws.close(), client._loop).result( + timeout=10 + ) + + assert client.state() is not None + def test_protocol_client_multiple_episodes(self, echo_server): """Test client can run multiple episodes.""" from envs.echo_env.client import EchoEnv From d27431a36ff5c31d2913b31b14a7483ee8460cd6 Mon Sep 17 00:00:00 2001 From: sergiopaniego Date: Fri, 30 Oct 2026 16:31:31 +0100 Subject: [PATCH 2/2] Bind the open episode to its socket Clients whose tests replace _send_and_receive never open a socket, so a boolean flag made step() raise SessionLostError on them. The client now remembers the socket reset() used and raises only if that socket was closed or replaced. --- src/openenv/core/env_client.py | 22 +++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/src/openenv/core/env_client.py b/src/openenv/core/env_client.py index cb2447a5dd..fb978e3625 100644 --- a/src/openenv/core/env_client.py +++ b/src/openenv/core/env_client.py @@ -395,8 +395,8 @@ def __init__( # Strong references for fire-and-forget socket closes (see _receive), # so the task isn't garbage-collected mid-close. Discarded on done. self._pending_close_tasks: set[asyncio.Task] = set() - # Whether the current socket carries an episode started by `reset()`. - self._episode_open = False + # The socket that carries the episode started by `reset()`, if one is open. + self._episode_ws: Optional[ClientConnection] = None if base_url is not None: self._set_base_url(base_url) @@ -664,7 +664,7 @@ async def _disconnect_async(self) -> None: # leave a stale socket cached for a later operation. self._ws = None self._ws_loop = None - self._episode_open = False + self._episode_ws = None same_loop = ws_loop is asyncio.get_running_loop() if same_loop: # Keep ownership of the handshake if this caller is cancelled. @@ -1009,10 +1009,10 @@ async def _reset_async(self, **kwargs: Any) -> StepResult[ObsT]: "data": kwargs, } # A reset may reconnect: it starts a new episode on whatever session it lands on. - self._episode_open = False + self._episode_ws = None response = await self._send_and_receive(message) result = self._parse_result(response.get("data", {})) - self._episode_open = not result.done + self._episode_ws = None if result.done else self._ws return result def step(self, action: ActT, **kwargs: Any) -> Any: @@ -1039,7 +1039,7 @@ async def _step_async(self, action: ActT, **kwargs: Any) -> StepResult[ObsT]: response = await self._send_and_receive(message) result = self._parse_result(response.get("data", {})) if result.done: - self._episode_open = False + self._episode_ws = None return result def state(self) -> Any: @@ -1062,17 +1062,17 @@ def _raise_if_session_lost(self) -> None: `_connect_async()` would replace it with a new socket, which the server treats as a new session, and the call would run on an environment this - client never reset. + client never reset. It keeps raising until `reset()`, so a retried call + does not land on the new session either. """ - if not self._episode_open: + ws = self._episode_ws + if ws is None: return - ws = self._ws if ( - ws is None + self._ws is not ws or ws.state in (State.CLOSING, State.CLOSED) or self._ws_loop is not asyncio.get_running_loop() ): - # Keep raising until reset(): a later step() would otherwise reconnect and run on a new session. raise SessionLostError( f"Lost the connection to {self._ws_url} in the middle of an episode. " "Call reset() to start a new one."