From 58b55000ce186bf3b30aa67a34c70c364625a1ca Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 16 Sep 2026 10:15:41 +0000 Subject: [PATCH] fix(mcp): share one production session across /ws and /mcp Resolve #1169 against current main by adopting the contributor's single-session connect (HTTP MCP create, then attach WebSocket with session_id) and teaching /ws to attach without destroying HTTP-owned sessions. Keeps sync-safe _close_async teardown from #1175. Co-authored-by: mugenkyou --- src/openenv/core/env_server/http_server.py | 57 +++++- src/openenv/core/mcp_client.py | 21 +- tests/core/test_mode_selection.py | 221 ++++++++++++++++----- 3 files changed, 230 insertions(+), 69 deletions(-) diff --git a/src/openenv/core/env_server/http_server.py b/src/openenv/core/env_server/http_server.py index 01954c5bd5..3f4fd9da46 100644 --- a/src/openenv/core/env_server/http_server.py +++ b/src/openenv/core/env_server/http_server.py @@ -1487,22 +1487,59 @@ async def websocket_endpoint(websocket: WebSocket): """ WebSocket endpoint for persistent environment sessions. - Each WebSocket connection gets its own environment instance. The client sends - WSResetMessage, WSStepMessage, WSStateMessage, or WSCloseMessage; the server - responds with WSObservationResponse, WSStateResponse, or WSErrorResponse. + Each WebSocket connection gets its own environment instance by default. + Clients may pass `?session_id=` to attach to an existing HTTP MCP + session instead of creating a second one (required for production-mode + clients that share Gym `/ws` and tool `/mcp` traffic). + + The client sends WSResetMessage, WSStepMessage, WSStateMessage, or + WSCloseMessage; the server responds with WSObservationResponse, + WSStateResponse, or WSErrorResponse. """ await websocket.accept() session_id = None session_env = None + owns_session = False try: - # Create session with dedicated environment - session_id, session_env = await self._create_session() - if session_env is None: - raise RuntimeError( - "Session environment not initialized for websocket" - ) + requested_session_id = websocket.query_params.get("session_id") + if requested_session_id: + # Attach to an HTTP MCP (or otherwise pre-created) session. + # Do not destroy it when this WebSocket disconnects — the + # HTTP owner remains responsible for session lifetime. + attached_env = self._sessions.get(requested_session_id, _MISSING) + if attached_env is _MISSING: + error_resp = WSErrorResponse( + data={ + "message": f"Unknown session_id: {requested_session_id}", + "code": WSErrorCode.SESSION_ERROR, + } + ) + await websocket.send_text(error_resp.model_dump_json()) + return + if attached_env is None: + error_resp = WSErrorResponse( + data={ + "message": ( + f"Session {requested_session_id} is still " + "initializing; retry shortly" + ), + "code": WSErrorCode.SESSION_ERROR, + } + ) + await websocket.send_text(error_resp.model_dump_json()) + return + session_id = requested_session_id + session_env = attached_env + else: + # Create session with dedicated environment + session_id, session_env = await self._create_session() + owns_session = True + if session_env is None: + raise RuntimeError( + "Session environment not initialized for websocket" + ) # Keep MCP session open for entire websocket lifetime # (avoids reconnect overhead on every message) @@ -1688,7 +1725,7 @@ async def websocket_endpoint(websocket: WebSocket): ) await websocket.send_text(error_resp.model_dump_json()) finally: - if session_id: + if session_id and owns_session: await self._destroy_session(session_id) try: await websocket.close() diff --git a/src/openenv/core/mcp_client.py b/src/openenv/core/mcp_client.py index 1afc6254b9..a6dac9c9bb 100644 --- a/src/openenv/core/mcp_client.py +++ b/src/openenv/core/mcp_client.py @@ -202,16 +202,23 @@ async def _connect_async(self) -> EnvClient: """ Establish connection to the server. - In production mode (`use_production_mode=True`), open the WebSocket used - by `reset` / `step` / `state` and create a persistent HTTP MCP session - for `list_tools` / `call_tool`. Tool calls bypass `step()` over `/mcp`, - but the Gym lifecycle still requires `/ws` until production routing - covers those methods end-to-end. + In production mode (`use_production_mode=True`), create an HTTP MCP + session first and connect the WebSocket with that `session_id` so + Gym (`reset` / `step` / `state`) and tool (`list_tools` / `call_tool`) + traffic share one server-side environment session. """ if getattr(self, "use_production_mode", False): try: - await super()._connect_async() - await self._ensure_production_session() + self._start_provider_if_needed() + session_id = await self._ensure_production_session() + original_ws_url = self._ws_url + if self._ws_url and "session_id=" not in self._ws_url: + sep = "&" if "?" in self._ws_url else "?" + self._ws_url = f"{self._ws_url}{sep}session_id={session_id}" + try: + await super()._connect_async() + finally: + self._ws_url = original_ws_url except Exception: await self.close() raise diff --git a/tests/core/test_mode_selection.py b/tests/core/test_mode_selection.py index 04aa7ef490..3161058364 100644 --- a/tests/core/test_mode_selection.py +++ b/tests/core/test_mode_selection.py @@ -256,16 +256,13 @@ async def test_production_mode_call_tool_uses_jsonrpc_protocol(self, clean_env): ) @pytest.mark.asyncio - async def test_production_mode_connect_opens_websocket_and_http_session( + async def test_production_mode_connect_creates_single_session_with_websocket( self, clean_env ): - """Production connect must open WebSocket (reset/step/state) and HTTP MCP session.""" + """Production connect creates one HTTP MCP session and attaches `/ws` to it.""" client = MCPToolClient(base_url="http://localhost:8000", mode="production") assert client.use_production_mode is True - mock_ws = MagicMock() - mock_ws.closed = False - with patch.object( client, "_production_mcp_request", @@ -275,14 +272,12 @@ async def test_production_mode_connect_opens_websocket_and_http_session( ], ) as mock_mcp_request: with patch( - "openenv.core.env_client.ws_connect", - new_callable=AsyncMock, - return_value=mock_ws, + "openenv.core.env_client.ws_connect", new_callable=AsyncMock ) as mock_ws_connect: await client.connect() mock_ws_connect.assert_called_once() - assert client._ws is mock_ws + assert "session_id=test-session" in mock_ws_connect.call_args[0][0] assert client._production_session_id == "test-session" mock_mcp_request.assert_called_once_with("openenv/session/create") @@ -300,65 +295,187 @@ async def test_production_mode_connect_opens_websocket_and_http_session( @pytest.mark.asyncio async def test_production_mode_connect_failure_cleans_up_resources(self, clean_env): - """Test that failure during production mode connect() triggers client.close() cleanup.""" + """Failure during production connect() must trigger client.close() cleanup.""" client = MCPToolClient(base_url="http://localhost:8000", mode="production") assert client.use_production_mode is True - mock_ws = MagicMock() - mock_ws.closed = False - mock_ws.close = AsyncMock() + with patch.object( + client, + "_ensure_production_session", + side_effect=RuntimeError("Session creation failed"), + ): + with patch.object(client, "close", wraps=client.close) as mock_close: + with pytest.raises(RuntimeError, match="Session creation failed"): + await client.connect() + + mock_close.assert_called_once() + + def test_production_mode_sync_close_closes_mcp_session(self, clean_env): + """Sync close must tear down the HTTP MCP session via `_close_async`.""" + client = MCPToolClient( + base_url="http://localhost:8000", mode="production" + ).sync() + client._async._production_session_id = "test-session-sync" + + mock_http_client = AsyncMock() + client._async._http_client = mock_http_client - with patch( - "openenv.core.env_client.ws_connect", + with patch.object( + client._async, + "_production_mcp_request", new_callable=AsyncMock, - return_value=mock_ws, - ): - with patch.object( - client, - "_ensure_production_session", - side_effect=RuntimeError("Session creation failed"), - ): - with patch.object(client, "close", wraps=client.close) as mock_close: - with pytest.raises(RuntimeError, match="Session creation failed"): - await client.connect() + return_value={"result": {"status": "closed"}}, + ) as mock_mcp_req: + client.close() + + mock_mcp_req.assert_awaited_once_with( + "openenv/session/close", + {"session_id": "test-session-sync"}, + ) + assert client._async._production_session_id is None + mock_http_client.aclose.assert_awaited_once() + assert client._async._http_client is None + + def test_production_mode_sync_context_manager_closes_mcp_session(self, clean_env): + """Sync context-manager exit must close the HTTP MCP session.""" + client = MCPToolClient( + base_url="http://localhost:8000", mode="production" + ).sync() + client._async._production_session_id = "test-session-context" + + mock_http_client = AsyncMock() + client._async._http_client = mock_http_client - mock_close.assert_called_once() + with patch.object( + client._async, + "_production_mcp_request", + new_callable=AsyncMock, + return_value={"result": {"status": "closed"}}, + ) as mock_mcp_req: + with patch.object(client._async, "_connect_async", new_callable=AsyncMock): + with client: + pass + + mock_mcp_req.assert_awaited_once_with( + "openenv/session/close", + {"session_id": "test-session-context"}, + ) + assert client._async._production_session_id is None + mock_http_client.aclose.assert_awaited_once() + assert client._async._http_client is None @pytest.mark.asyncio - async def test_production_mode_sync_close_closes_mcp_session(self, clean_env): - """Sync close must tear down the HTTP MCP session via `_close_async`.""" + async def test_production_mode_async_close_closes_mcp_session(self, clean_env): + """Async close must tear down the HTTP MCP session.""" client = MCPToolClient(base_url="http://localhost:8000", mode="production") - assert client.use_production_mode is True + client._production_session_id = "test-session-async" - mock_ws = MagicMock() - mock_ws.closed = False - mock_ws.close = AsyncMock() + mock_http_client = AsyncMock() + client._http_client = mock_http_client with patch.object( client, "_production_mcp_request", - side_effect=[ - {"result": {"session_id": "test-session"}}, - {"result": {}}, - ], - ) as mock_mcp_request: - with patch( - "openenv.core.env_client.ws_connect", - new_callable=AsyncMock, - return_value=mock_ws, - ): - sync_client = client.sync() - sync_client.connect() - assert client._production_session_id == "test-session" + new_callable=AsyncMock, + return_value={"result": {"status": "closed"}}, + ) as mock_mcp_req: + await client.close() - sync_client.close() + mock_mcp_req.assert_awaited_once_with( + "openenv/session/close", + {"session_id": "test-session-async"}, + ) + assert client._production_session_id is None + mock_http_client.aclose.assert_awaited_once() + assert client._http_client is None - assert client._production_session_id is None - assert mock_mcp_request.call_count == 2 - mock_mcp_request.assert_any_call( - "openenv/session/close", - {"session_id": "test-session"}, - ) + @pytest.mark.asyncio + async def test_websocket_disconnect_preserves_attached_http_session(self): + """Attaching `/ws` to an existing session must not destroy it on disconnect.""" + from fastapi import FastAPI + from openenv.core.env_server.http_server import HTTPEnvServer + from openenv.core.env_server.interfaces import Environment + from openenv.core.env_server.types import Action, Observation + from starlette.testclient import TestClient + + class MinimalAction(Action): + pass + + class MinimalObservation(Observation): + pass + + class MinimalEnvironment(Environment): + def reset(self): + return MinimalObservation() + + def step(self, action): + return MinimalObservation() + + @property + def state(self): + return {} + + server = HTTPEnvServer( + env=MinimalEnvironment, + action_cls=MinimalAction, + observation_cls=MinimalObservation, + ) + app = FastAPI() + server.register_routes(app) + + session_id, env_instance = await server._create_session() + assert session_id in server._sessions + + with TestClient(app) as test_client: + with test_client.websocket_connect(f"/ws?session_id={session_id}") as ws: + ws.send_json({"type": "close"}) + + assert session_id in server._sessions + assert server._sessions[session_id] is env_instance + + await server._destroy_session(session_id) + assert session_id not in server._sessions + + @pytest.mark.asyncio + async def test_websocket_disconnect_destroys_websocket_created_session(self): + """A WebSocket-created session must still be destroyed on disconnect.""" + from fastapi import FastAPI + from openenv.core.env_server.http_server import HTTPEnvServer + from openenv.core.env_server.interfaces import Environment + from openenv.core.env_server.types import Action, Observation + from starlette.testclient import TestClient + + class MinimalAction(Action): + pass + + class MinimalObservation(Observation): + pass + + class MinimalEnvironment(Environment): + def reset(self): + return MinimalObservation() + + def step(self, action): + return MinimalObservation() + + @property + def state(self): + return {} + + server = HTTPEnvServer( + env=MinimalEnvironment, + action_cls=MinimalAction, + observation_cls=MinimalObservation, + ) + app = FastAPI() + server.register_routes(app) + + with TestClient(app) as test_client: + with test_client.websocket_connect("/ws") as ws: + assert len(server._sessions) == 1 + ws.send_json({"type": "close"}) + + assert len(server._sessions) == 0 # ============================================================================