diff --git a/src/openenv/core/__init__.py b/src/openenv/core/__init__.py index b365f5ff7..1f62b60f6 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 01abe1aeb..fb978e362 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() + # 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) diff --git a/src/openenv/core/mcp_client.py b/src/openenv/core/mcp_client.py index 9002fabf9..c7967a926 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 dce2bbcb0..d1d627588 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