Raise SessionLostError when the WebSocket drops mid-episode

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"),

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()
+        # 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)
 
@@ -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_ws = None
             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_ws = None
         response = await self._send_and_receive(message)
-        return self._parse_result(response.get("data", {}))
+        result = self._parse_result(response.get("data", {}))
+        self._episode_ws = None if result.done else self._ws
+        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_ws = None
+        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. It keeps raising until `reset()`, so a retried call
+        does not land on the new session either.
+        """
+        ws = self._episode_ws
+        if ws is None:
+            return
+        if (
+            self._ws is not ws
+            or ws.state in (State.CLOSING, State.CLOSED)
+            or self._ws_loop is not asyncio.get_running_loop()
+        ):
+            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)
 

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:

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