From 91fce5dda8612f493ee670f8700e508479333c60 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 16 Sep 2026 10:17:19 +0000 Subject: [PATCH 1/8] fix(mcp): share one production session across HTTP MCP and WebSocket Land contributor #1169 lifecycle work on top of #1175: create the HTTP MCP session first, attach `/ws` with that session_id, and only destroy WebSocket-owned sessions on disconnect. Hardens `_production_mcp_url` against temporary query rewrites. Co-authored-by: mugenkyou --- src/openenv/core/env_server/http_server.py | 69 +++++-- src/openenv/core/mcp_client.py | 25 ++- tests/core/test_mode_selection.py | 201 +++++++++++++++------ 3 files changed, 218 insertions(+), 77 deletions(-) diff --git a/src/openenv/core/env_server/http_server.py b/src/openenv/core/env_server/http_server.py index 01954c5bd5..85f47a2519 100644 --- a/src/openenv/core/env_server/http_server.py +++ b/src/openenv/core/env_server/http_server.py @@ -134,6 +134,7 @@ async def _maybe_await(value: Any) -> Any: ConcurrencyConfigurationError, EnvironmentFactoryError, SessionCapacityError, + SessionNotFoundError, ) @@ -1487,33 +1488,56 @@ 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. + By default each WebSocket connection gets its own environment instance. + Pass `?session_id=` to attach to an existing HTTP MCP session so + Gym (`reset`/`step`/`state`) and tools share one server-side env. + Attached sockets do not destroy the session on disconnect; only + WebSocket-created sessions are cleaned up here. + + 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" - ) - - # Keep MCP session open for entire websocket lifetime - # (avoids reconnect overhead on every message) + requested_session_id = websocket.query_params.get("session_id") + if requested_session_id: + # Attach to an existing HTTP MCP (or other) session. + async with self._session_lock: + existing = self._sessions.get(requested_session_id, _MISSING) + if existing is _MISSING: + raise SessionNotFoundError(requested_session_id) + if existing is None: + raise RuntimeError( + f"Session {requested_session_id} is still initializing; " + "retry shortly" + ) + session_id = requested_session_id + session_env = existing + 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 WebSocket-owned sessions (attached + # sessions already hold MCP via `_session_stacks`). async with AsyncExitStack() as stack: - mcp_session_factory = getattr(session_env, "mcp_session", None) - if callable(mcp_session_factory): - mcp_session_cm = cast( - AsyncContextManager[Any], mcp_session_factory() - ) - await stack.enter_async_context(mcp_session_cm) + if owns_session: + mcp_session_factory = getattr(session_env, "mcp_session", None) + if callable(mcp_session_factory): + mcp_session_cm = cast( + AsyncContextManager[Any], mcp_session_factory() + ) + await stack.enter_async_context(mcp_session_cm) while True: # Receive message from client @@ -1682,13 +1706,20 @@ async def websocket_endpoint(websocket: WebSocket): } ) await websocket.send_text(error_resp.model_dump_json()) + except SessionNotFoundError as e: + error_resp = WSErrorResponse( + data={"message": str(e), "code": WSErrorCode.SESSION_ERROR} + ) + await websocket.send_text(error_resp.model_dump_json()) except Exception as e: error_resp = WSErrorResponse( data={"message": str(e), "code": WSErrorCode.SESSION_ERROR} ) await websocket.send_text(error_resp.model_dump_json()) finally: - if session_id: + # Only destroy sessions this WebSocket created. Attached HTTP MCP + # sessions are closed via `openenv/session/close` on the client. + 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..dd611be35f 100644 --- a/src/openenv/core/mcp_client.py +++ b/src/openenv/core/mcp_client.py @@ -168,6 +168,9 @@ def _next_request_id(self) -> int: def _production_mcp_url(self) -> str: """Build HTTP MCP endpoint URL from the client's websocket URL.""" url = self._ws_url.replace("ws://", "http://").replace("wss://", "https://") + # Strip query/fragment so a temporary `?session_id=` rewrite on `_ws_url` + # cannot poison session create/close posts to `/mcp`. + url = url.split("?", 1)[0].split("#", 1)[0] if url.endswith("/ws"): url = url[: -len("/ws")] return url.rstrip("/") + "/mcp" @@ -202,16 +205,24 @@ 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 a persistent HTTP + MCP session first, then open the WebSocket with that `session_id` so + `reset` / `step` / `state` and `list_tools` / `call_tool` share one + server-side environment (avoids dual-session capacity failures). """ 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: + # Always restore so failed connect / close still posts to `/mcp`. + 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..e52e7f9915 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.""" + """Test that connect() in production mode initializes the HTTP MCP session AND connects WebSocket using the same session ID.""" 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,17 +272,18 @@ 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: + # Explicit connect (e.g. from async with client:) await client.connect() + # Should create HTTP session and connect WS with session_id query param 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") + # Subsequent call_tool should reuse the same session result = await client.call_tool("echo", message="hello world") assert result == "hello world" assert mock_mcp_request.call_count == 2 @@ -304,61 +302,162 @@ async def test_production_mode_connect_failure_cleans_up_resources(self, clean_e 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): + """Test that production sync close() closes the MCP session and releases HTTP client.""" + client = MCPToolClient( + base_url="http://localhost:8000", mode="production" + ).sync() + client._async._production_session_id = "test-session-sync" - with patch( - "openenv.core.env_client.ws_connect", + mock_http_client = AsyncMock() + client._async._http_client = mock_http_client + + 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): + """Test that production sync context-manager exit closes the MCP session and releases HTTP client.""" + 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): + """Test that production async close() closes the MCP session and releases HTTP client.""" 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): + """Test that a WebSocket attaching to an existing session does NOT destroy that session when disconnected.""" + from fastapi import FastAPI + from openenv.core.env_server.http_server import HTTPEnvServer + from starlette.testclient import TestClient + + # Import sibling module fixtures (tests/ is not a package on PYTHONPATH=src:envs). + from test_production_mode_routes import ( + MinimalAction, + MinimalEnvironment, + MinimalObservation, + ) + + server = HTTPEnvServer( + env=MinimalEnvironment, + action_cls=MinimalAction, + observation_cls=MinimalObservation, + ) + app = FastAPI() + server.register_routes(app) + + # 1. Create session via HTTP / HTTPEnvServer + session_id, env_instance = await server._create_session() + assert session_id in server._sessions + + # 2. Attach WebSocket to existing session_id + with TestClient(app) as test_client: + with test_client.websocket_connect(f"/ws?session_id={session_id}") as ws: + ws.send_json({"type": "close"}) + + # 3. Session must remain alive because WebSocket didn't create it + assert session_id in server._sessions + assert server._sessions[session_id] is env_instance + + # 4. Clean up + 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): + """Test that a WebSocket creating its own session DOES destroy that session when disconnected.""" + from fastapi import FastAPI + from openenv.core.env_server.http_server import HTTPEnvServer + from starlette.testclient import TestClient + from test_production_mode_routes import ( + MinimalAction, + MinimalEnvironment, + MinimalObservation, + ) + + 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"}) + + # After WebSocket disconnects, the session created by WebSocket should be destroyed + assert len(server._sessions) == 0 # ============================================================================ From 8a13ccb28bac91fbc6c0c85efe168db98ee80e5e Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 16 Sep 2026 10:23:39 +0000 Subject: [PATCH 2/8] fix(mcp): adopt #1169 tip and convert ws base_url for /mcp MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Sync contributor shared-session client/server/tests from #1169 a8eec758, then fix Bugbot medium: `_production_mcp_url` must map ws://|wss:// → http(s) before posting to `/mcp`. Co-authored-by: mugenkyou --- src/openenv/core/env_server/http_server.py | 71 +++++++++----------- src/openenv/core/mcp_client.py | 26 ++++---- tests/core/test_mode_selection.py | 78 +++------------------- tests/core/test_production_mode_routes.py | 72 ++++++++++++++++++++ 4 files changed, 127 insertions(+), 120 deletions(-) diff --git a/src/openenv/core/env_server/http_server.py b/src/openenv/core/env_server/http_server.py index 85f47a2519..5f3b056447 100644 --- a/src/openenv/core/env_server/http_server.py +++ b/src/openenv/core/env_server/http_server.py @@ -134,7 +134,6 @@ async def _maybe_await(value: Any) -> Any: ConcurrencyConfigurationError, EnvironmentFactoryError, SessionCapacityError, - SessionNotFoundError, ) @@ -1488,15 +1487,11 @@ async def websocket_endpoint(websocket: WebSocket): """ WebSocket endpoint for persistent environment sessions. - By default each WebSocket connection gets its own environment instance. - Pass `?session_id=` to attach to an existing HTTP MCP session so - Gym (`reset`/`step`/`state`) and tools share one server-side env. - Attached sockets do not destroy the session on disconnect; only - WebSocket-created sessions are cleaned up here. - - The client sends WSResetMessage, WSStepMessage, WSStateMessage, or - WSCloseMessage; the server responds with WSObservationResponse, - WSStateResponse, or WSErrorResponse. + By default, each WebSocket connection gets its own environment instance. + A client can instead attach to an existing HTTP MCP session by passing its + session ID in the query string. The client sends WSResetMessage, + WSStepMessage, WSStateMessage, or WSCloseMessage; the server responds with + WSObservationResponse, WSStateResponse, or WSErrorResponse. """ await websocket.accept() @@ -1507,37 +1502,40 @@ async def websocket_endpoint(websocket: WebSocket): try: requested_session_id = websocket.query_params.get("session_id") if requested_session_id: - # Attach to an existing HTTP MCP (or other) session. async with self._session_lock: - existing = self._sessions.get(requested_session_id, _MISSING) - if existing is _MISSING: - raise SessionNotFoundError(requested_session_id) - if existing is None: + attached_env = self._sessions.get( + requested_session_id, _MISSING + ) + if attached_env is _MISSING: + raise RuntimeError( + f"Unknown session_id: {requested_session_id}" + ) + if attached_env is None: raise RuntimeError( - f"Session {requested_session_id} is still initializing; " - "retry shortly" + f"Session {requested_session_id} is still initializing" ) session_id = requested_session_id - session_env = existing + session_env = attached_env + self._update_session_activity(session_id) 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 WebSocket-owned sessions (attached - # sessions already hold MCP via `_session_stacks`). + 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) + async with AsyncExitStack() as stack: - if owns_session: - mcp_session_factory = getattr(session_env, "mcp_session", None) - if callable(mcp_session_factory): - mcp_session_cm = cast( - AsyncContextManager[Any], mcp_session_factory() - ) - await stack.enter_async_context(mcp_session_cm) + mcp_session_factory = getattr(session_env, "mcp_session", None) + if owns_session and callable(mcp_session_factory): + mcp_session_cm = cast( + AsyncContextManager[Any], mcp_session_factory() + ) + await stack.enter_async_context(mcp_session_cm) while True: # Receive message from client @@ -1706,20 +1704,13 @@ async def websocket_endpoint(websocket: WebSocket): } ) await websocket.send_text(error_resp.model_dump_json()) - except SessionNotFoundError as e: - error_resp = WSErrorResponse( - data={"message": str(e), "code": WSErrorCode.SESSION_ERROR} - ) - await websocket.send_text(error_resp.model_dump_json()) except Exception as e: error_resp = WSErrorResponse( data={"message": str(e), "code": WSErrorCode.SESSION_ERROR} ) await websocket.send_text(error_resp.model_dump_json()) finally: - # Only destroy sessions this WebSocket created. Attached HTTP MCP - # sessions are closed via `openenv/session/close` on the client. - if session_id and owns_session: + if owns_session and session_id: 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 dd611be35f..b97c50ec64 100644 --- a/src/openenv/core/mcp_client.py +++ b/src/openenv/core/mcp_client.py @@ -166,14 +166,16 @@ def _next_request_id(self) -> int: return self._jsonrpc_request_id def _production_mcp_url(self) -> str: - """Build HTTP MCP endpoint URL from the client's websocket URL.""" - url = self._ws_url.replace("ws://", "http://").replace("wss://", "https://") - # Strip query/fragment so a temporary `?session_id=` rewrite on `_ws_url` - # cannot poison session create/close posts to `/mcp`. - url = url.split("?", 1)[0].split("#", 1)[0] - if url.endswith("/ws"): - url = url[: -len("/ws")] - return url.rstrip("/") + "/mcp" + """Build the HTTP MCP endpoint URL from the stable base URL.""" + if self._base_url is None: + raise RuntimeError("MCP client is not connected to a server.") + # `_base_url` may be `ws://` / `wss://` (documented for EnvClient); httpx needs HTTP. + url = ( + self._base_url.replace("ws://", "http://") + .replace("wss://", "https://") + .rstrip("/") + ) + return url + "/mcp" async def _get_http_client(self) -> Any: """Return a shared httpx.AsyncClient, creating one lazily.""" @@ -205,10 +207,9 @@ async def _connect_async(self) -> EnvClient: """ Establish connection to the server. - In production mode (`use_production_mode=True`), create a persistent HTTP - MCP session first, then open the WebSocket with that `session_id` so - `reset` / `step` / `state` and `list_tools` / `call_tool` share one - server-side environment (avoids dual-session capacity failures). + In production mode (use_production_mode=True), creates an HTTP MCP session + and connects the WebSocket using that session ID so that WebSocket (reset/step/state) + and HTTP MCP (list_tools/call_tool) share the exact same server-side environment session. """ if getattr(self, "use_production_mode", False): try: @@ -221,7 +222,6 @@ async def _connect_async(self) -> EnvClient: try: await super()._connect_async() finally: - # Always restore so failed connect / close still posts to `/mcp`. self._ws_url = original_ws_url except Exception: await self.close() diff --git a/tests/core/test_mode_selection.py b/tests/core/test_mode_selection.py index e52e7f9915..1dffa0a9c9 100644 --- a/tests/core/test_mode_selection.py +++ b/tests/core/test_mode_selection.py @@ -255,6 +255,17 @@ async def test_production_mode_call_tool_uses_jsonrpc_protocol(self, clean_env): }, ) + def test_production_mcp_url_converts_ws_base_url(self, clean_env): + """ws:// / wss:// base URLs must become http(s) for production /mcp posts.""" + ws_client = MCPToolClient(base_url="ws://localhost:8000", mode="production") + assert ws_client._production_mcp_url() == "http://localhost:8000/mcp" + + wss_client = MCPToolClient(base_url="wss://example.com", mode="production") + assert wss_client._production_mcp_url() == "https://example.com/mcp" + + http_client = MCPToolClient(base_url="http://localhost:8000", mode="production") + assert http_client._production_mcp_url() == "http://localhost:8000/mcp" + @pytest.mark.asyncio async def test_production_mode_connect_creates_single_session_with_websocket( self, clean_env @@ -392,73 +403,6 @@ async def test_production_mode_async_close_closes_mcp_session(self, clean_env): mock_http_client.aclose.assert_awaited_once() assert client._http_client is None - @pytest.mark.asyncio - async def test_websocket_disconnect_preserves_attached_http_session(self): - """Test that a WebSocket attaching to an existing session does NOT destroy that session when disconnected.""" - from fastapi import FastAPI - from openenv.core.env_server.http_server import HTTPEnvServer - from starlette.testclient import TestClient - - # Import sibling module fixtures (tests/ is not a package on PYTHONPATH=src:envs). - from test_production_mode_routes import ( - MinimalAction, - MinimalEnvironment, - MinimalObservation, - ) - - server = HTTPEnvServer( - env=MinimalEnvironment, - action_cls=MinimalAction, - observation_cls=MinimalObservation, - ) - app = FastAPI() - server.register_routes(app) - - # 1. Create session via HTTP / HTTPEnvServer - session_id, env_instance = await server._create_session() - assert session_id in server._sessions - - # 2. Attach WebSocket to existing session_id - with TestClient(app) as test_client: - with test_client.websocket_connect(f"/ws?session_id={session_id}") as ws: - ws.send_json({"type": "close"}) - - # 3. Session must remain alive because WebSocket didn't create it - assert session_id in server._sessions - assert server._sessions[session_id] is env_instance - - # 4. Clean up - 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): - """Test that a WebSocket creating its own session DOES destroy that session when disconnected.""" - from fastapi import FastAPI - from openenv.core.env_server.http_server import HTTPEnvServer - from starlette.testclient import TestClient - from test_production_mode_routes import ( - MinimalAction, - MinimalEnvironment, - MinimalObservation, - ) - - 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"}) - - # After WebSocket disconnects, the session created by WebSocket should be destroyed - assert len(server._sessions) == 0 - # ============================================================================ # Mode Immutability Tests diff --git a/tests/core/test_production_mode_routes.py b/tests/core/test_production_mode_routes.py index b78df5680a..93e88ef264 100644 --- a/tests/core/test_production_mode_routes.py +++ b/tests/core/test_production_mode_routes.py @@ -786,6 +786,78 @@ def test_session_create_from_websocket_is_idempotent(self, app): response2 = json.loads(response_text2) assert response2["data"]["result"]["session_id"] == ws_session_id + def test_websocket_can_attach_to_http_session_without_destroying_it(self, app): + """An attached WebSocket shares and preserves an HTTP-created session.""" + client = TestClient(app) + create_response = client.post( + "/mcp", + json={ + "jsonrpc": "2.0", + "method": "openenv/session/create", + "params": {}, + "id": 1, + }, + ) + session_id = create_response.json()["result"]["session_id"] + + with client.websocket_connect(f"/ws?session_id={session_id}") as websocket: + websocket.send_json({"type": "state"}) + state_response = websocket.receive_json() + assert state_response["type"] == "state" + websocket.send_json({"type": "close"}) + + tools_response = client.post( + "/mcp", + json={ + "jsonrpc": "2.0", + "method": "tools/list", + "params": {"session_id": session_id}, + "id": 2, + }, + ) + assert "result" in tools_response.json() + + close_response = client.post( + "/mcp", + json={ + "jsonrpc": "2.0", + "method": "openenv/session/close", + "params": {"session_id": session_id}, + "id": 3, + }, + ) + assert close_response.json()["result"]["closed"] is True + + def test_websocket_still_destroys_its_own_session(self, app): + """A WebSocket-created session is destroyed when the socket closes.""" + client = TestClient(app) + + with client.websocket_connect("/ws") as websocket: + websocket.send_json( + { + "type": "mcp", + "data": { + "jsonrpc": "2.0", + "method": "openenv/session/create", + "params": {}, + "id": 1, + }, + } + ) + session_id = websocket.receive_json()["data"]["result"]["session_id"] + websocket.send_json({"type": "close"}) + + tools_response = client.post( + "/mcp", + json={ + "jsonrpc": "2.0", + "method": "tools/list", + "params": {"session_id": session_id}, + "id": 2, + }, + ) + assert tools_response.json()["error"]["code"] == -32602 + def test_session_close_missing_session_id_param(self, app): """Test openenv/session/close without session_id returns INVALID_PARAMS.""" from starlette.testclient import TestClient From cedd6d9dd3964346dccabc3895ba4b8127180dd0 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 16 Sep 2026 10:25:25 +0000 Subject: [PATCH 3/8] =?UTF-8?q?fix(mcp):=20sync=20#1169=20ownership=20guar?= =?UTF-8?q?ds=20+=20ws=E2=86=92http=20MCP=20URL?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull latest contributor tip (77a9915f session ownership / attach locks) and keep the Bugbot medium fix converting ws://|wss:// base URLs for /mcp. Co-authored-by: mugenkyou Co-authored-by: burtenshaw --- src/openenv/core/env_server/http_server.py | 66 ++++++++++++++++------ src/openenv/core/mcp_client.py | 33 ++++++----- tests/core/test_mode_selection.py | 7 ++- tests/core/test_production_mode_routes.py | 53 ++++++++++++++++- 4 files changed, 127 insertions(+), 32 deletions(-) diff --git a/src/openenv/core/env_server/http_server.py b/src/openenv/core/env_server/http_server.py index 5f3b056447..a3e242bc2d 100644 --- a/src/openenv/core/env_server/http_server.py +++ b/src/openenv/core/env_server/http_server.py @@ -248,6 +248,7 @@ def __init__( self._session_executors: Dict[str, ThreadPoolExecutor] = {} self._session_stacks: Dict[str, AsyncExitStack] = {} self._session_info: Dict[str, SessionInfo] = {} + self._session_websocket_attachments: set[str] = set() self._session_lock = asyncio.Lock() # Create thread pool for running sync code in async context @@ -451,6 +452,7 @@ async def _destroy_session(self, session_id: str) -> None: executor = self._session_executors.pop(session_id, None) stack = self._session_stacks.pop(session_id, None) self._session_info.pop(session_id, None) + self._session_websocket_attachments.discard(session_id) await self._cleanup_session_resources(env, executor, stack) @@ -522,7 +524,10 @@ async def _reap_idle_sessions(self) -> None: stale_ids: list[str] = [] async with self._session_lock: for sid, info in self._session_info.items(): - if now - info.last_activity_at > timeout: + if ( + sid not in self._session_websocket_attachments + and now - info.last_activity_at > timeout + ): stale_ids.append(sid) for sid in stale_ids: # Re-check under lock: activity may have arrived since @@ -532,7 +537,11 @@ async def _reap_idle_sessions(self) -> None: now = time.time() async with self._session_lock: info = self._session_info.get(sid) - if info is None or (now - info.last_activity_at) <= timeout: + if ( + info is None + or sid in self._session_websocket_attachments + or (now - info.last_activity_at) <= timeout + ): continue await self._destroy_session(sid) except asyncio.CancelledError: @@ -804,15 +813,29 @@ async def mcp_handler( ) async with self._session_lock: - env = self._sessions.pop(target_session_id, _MISSING) - if env is not _MISSING: + if target_session_id in self._session_websocket_attachments: + env = _MISSING + attached = True + executor = None + stack = None + else: + attached = False + env = self._sessions.pop(target_session_id, _MISSING) + if not attached and env is not _MISSING: executor = self._session_executors.pop(target_session_id, None) stack = self._session_stacks.pop(target_session_id, None) self._session_info.pop(target_session_id, None) - else: + elif not attached: executor = None stack = None + if attached: + return JsonRpcResponse.error_response( + JsonRpcErrorCode.INVALID_REQUEST, + f"Session {target_session_id} has an active WebSocket", + request_id=request_id, + ) + if env is _MISSING: return JsonRpcResponse.error_response( JsonRpcErrorCode.INVALID_PARAMS, @@ -1498,6 +1521,7 @@ async def websocket_endpoint(websocket: WebSocket): session_id = None session_env = None owns_session = False + attached_session = False try: requested_session_id = websocket.query_params.get("session_id") @@ -1506,16 +1530,23 @@ async def websocket_endpoint(websocket: WebSocket): attached_env = self._sessions.get( requested_session_id, _MISSING ) - if attached_env is _MISSING: - raise RuntimeError( - f"Unknown session_id: {requested_session_id}" - ) - if attached_env is None: - raise RuntimeError( - f"Session {requested_session_id} is still initializing" - ) - session_id = requested_session_id - session_env = attached_env + if attached_env is _MISSING: + raise RuntimeError( + f"Unknown session_id: {requested_session_id}" + ) + if attached_env is None: + raise RuntimeError( + f"Session {requested_session_id} is still initializing" + ) + if requested_session_id in self._session_websocket_attachments: + raise RuntimeError( + f"Session {requested_session_id} already has " + "an attached WebSocket" + ) + self._session_websocket_attachments.add(requested_session_id) + session_id = requested_session_id + session_env = attached_env + attached_session = True self._update_session_activity(session_id) else: session_id, session_env = await self._create_session() @@ -1710,7 +1741,10 @@ async def websocket_endpoint(websocket: WebSocket): ) await websocket.send_text(error_resp.model_dump_json()) finally: - if owns_session and session_id: + if attached_session and session_id: + async with self._session_lock: + self._session_websocket_attachments.discard(session_id) + elif owns_session and session_id: 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 b97c50ec64..3e383616c5 100644 --- a/src/openenv/core/mcp_client.py +++ b/src/openenv/core/mcp_client.py @@ -56,6 +56,7 @@ import asyncio from typing import Any, Dict, List, Optional +from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit from pydantic import ConfigDict @@ -156,6 +157,7 @@ def __init__( self._tools_cache: Optional[List[Tool]] = None self.use_production_mode = self._mode == "production" self._production_session_id: Optional[str] = None + self._production_connect_lock = asyncio.Lock() self._production_session_lock = asyncio.Lock() self._jsonrpc_request_id = 0 self._http_client: Optional[Any] = None # lazily-created httpx.AsyncClient @@ -212,20 +214,25 @@ async def _connect_async(self) -> EnvClient: and HTTP MCP (list_tools/call_tool) share the exact same server-side environment session. """ if getattr(self, "use_production_mode", False): - try: - 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}" + async with self._production_connect_lock: try: - await super()._connect_async() - finally: - self._ws_url = original_ws_url - except Exception: - await self.close() - raise + self._start_provider_if_needed() + session_id = await self._ensure_production_session() + original_ws_url = self._ws_url + if original_ws_url is None: + raise RuntimeError("MCP client has no WebSocket URL.") + + parts = urlsplit(original_ws_url) + query = dict(parse_qsl(parts.query, keep_blank_values=True)) + query["session_id"] = session_id + self._ws_url = urlunsplit(parts._replace(query=urlencode(query))) + try: + await super()._connect_async() + finally: + self._ws_url = original_ws_url + except Exception: + await self.close() + raise return self return await super()._connect_async() diff --git a/tests/core/test_mode_selection.py b/tests/core/test_mode_selection.py index 1dffa0a9c9..6e9f1cb941 100644 --- a/tests/core/test_mode_selection.py +++ b/tests/core/test_mode_selection.py @@ -273,6 +273,8 @@ async def test_production_mode_connect_creates_single_session_with_websocket( """Test that connect() in production mode initializes the HTTP MCP session AND connects WebSocket using the same session ID.""" client = MCPToolClient(base_url="http://localhost:8000", mode="production") assert client.use_production_mode is True + client._ws_url = f"{client._ws_url}?some_session_id=keep" + original_ws_url = client._ws_url with patch.object( client, @@ -290,7 +292,10 @@ async def test_production_mode_connect_creates_single_session_with_websocket( # Should create HTTP session and connect WS with session_id query param mock_ws_connect.assert_called_once() - assert "session_id=test-session" in mock_ws_connect.call_args[0][0] + connected_url = mock_ws_connect.call_args[0][0] + assert "session_id=test-session" in connected_url + assert "some_session_id=keep" in connected_url + assert client._ws_url == original_ws_url assert client._production_session_id == "test-session" mock_mcp_request.assert_called_once_with("openenv/session/create") diff --git a/tests/core/test_production_mode_routes.py b/tests/core/test_production_mode_routes.py index 93e88ef264..e856ed9a6c 100644 --- a/tests/core/test_production_mode_routes.py +++ b/tests/core/test_production_mode_routes.py @@ -804,6 +804,19 @@ def test_websocket_can_attach_to_http_session_without_destroying_it(self, app): websocket.send_json({"type": "state"}) state_response = websocket.receive_json() assert state_response["type"] == "state" + + active_close_response = client.post( + "/mcp", + json={ + "jsonrpc": "2.0", + "method": "openenv/session/close", + "params": {"session_id": session_id}, + "id": 2, + }, + ) + assert ( + "active WebSocket" in active_close_response.json()["error"]["message"] + ) websocket.send_json({"type": "close"}) tools_response = client.post( @@ -812,7 +825,7 @@ def test_websocket_can_attach_to_http_session_without_destroying_it(self, app): "jsonrpc": "2.0", "method": "tools/list", "params": {"session_id": session_id}, - "id": 2, + "id": 3, }, ) assert "result" in tools_response.json() @@ -823,7 +836,43 @@ def test_websocket_can_attach_to_http_session_without_destroying_it(self, app): "jsonrpc": "2.0", "method": "openenv/session/close", "params": {"session_id": session_id}, - "id": 3, + "id": 4, + }, + ) + assert close_response.json()["result"]["closed"] is True + + def test_http_session_allows_only_one_attached_websocket(self, app): + """A second WebSocket cannot concurrently mutate the same session.""" + client = TestClient(app) + create_response = client.post( + "/mcp", + json={ + "jsonrpc": "2.0", + "method": "openenv/session/create", + "params": {}, + "id": 1, + }, + ) + session_id = create_response.json()["result"]["session_id"] + + with client.websocket_connect(f"/ws?session_id={session_id}") as first_socket: + with client.websocket_connect( + f"/ws?session_id={session_id}" + ) as second_socket: + error_response = second_socket.receive_json() + assert ( + "already has an attached WebSocket" + in (error_response["data"]["message"]) + ) + first_socket.send_json({"type": "close"}) + + close_response = client.post( + "/mcp", + json={ + "jsonrpc": "2.0", + "method": "openenv/session/close", + "params": {"session_id": session_id}, + "id": 2, }, ) assert close_response.json()["result"]["closed"] is True From 1742eb56a35668ae1a7302c4ef190730d6d489ee Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 16 Sep 2026 10:32:16 +0000 Subject: [PATCH 4/8] fix(mcp): detach WebSocket before HTTP session close Bugbot high on #1169: openenv/session/close was rejected while /ws was attached, so production close cleared the local id and leaked the server session. Disconnect first; make HTTP session/close authoritative if a socket is still attached. Co-authored-by: mugenkyou --- src/openenv/core/env_server/http_server.py | 24 +++------- src/openenv/core/mcp_client.py | 36 ++++++++++----- tests/core/test_mode_selection.py | 53 ++++++++++++++++----- tests/core/test_production_mode_routes.py | 54 ++++++++++++++++------ 4 files changed, 115 insertions(+), 52 deletions(-) diff --git a/src/openenv/core/env_server/http_server.py b/src/openenv/core/env_server/http_server.py index a3e242bc2d..425eb27b05 100644 --- a/src/openenv/core/env_server/http_server.py +++ b/src/openenv/core/env_server/http_server.py @@ -812,30 +812,20 @@ async def mcp_handler( request_id=request_id, ) + # HTTP `openenv/session/close` is authoritative: drop any + # WebSocket attachment marker and destroy the session. Clients + # should detach first; this still recovers if they race. async with self._session_lock: - if target_session_id in self._session_websocket_attachments: - env = _MISSING - attached = True - executor = None - stack = None - else: - attached = False - env = self._sessions.pop(target_session_id, _MISSING) - if not attached and env is not _MISSING: + self._session_websocket_attachments.discard(target_session_id) + env = self._sessions.pop(target_session_id, _MISSING) + if env is not _MISSING: executor = self._session_executors.pop(target_session_id, None) stack = self._session_stacks.pop(target_session_id, None) self._session_info.pop(target_session_id, None) - elif not attached: + else: executor = None stack = None - if attached: - return JsonRpcResponse.error_response( - JsonRpcErrorCode.INVALID_REQUEST, - f"Session {target_session_id} has an active WebSocket", - request_id=request_id, - ) - if env is _MISSING: return JsonRpcResponse.error_response( JsonRpcErrorCode.INVALID_PARAMS, diff --git a/src/openenv/core/mcp_client.py b/src/openenv/core/mcp_client.py index 3e383616c5..de5a4ec712 100644 --- a/src/openenv/core/mcp_client.py +++ b/src/openenv/core/mcp_client.py @@ -171,13 +171,16 @@ def _production_mcp_url(self) -> str: """Build the HTTP MCP endpoint URL from the stable base URL.""" if self._base_url is None: raise RuntimeError("MCP client is not connected to a server.") - # `_base_url` may be `ws://` / `wss://` (documented for EnvClient); httpx needs HTTP. - url = ( - self._base_url.replace("ws://", "http://") - .replace("wss://", "https://") - .rstrip("/") + parts = urlsplit(self._base_url) + scheme = {"ws": "http", "wss": "https"}.get(parts.scheme, parts.scheme) + return urlunsplit( + parts._replace( + scheme=scheme, + path=parts.path.rstrip("/") + "/mcp", + query="", + fragment="", + ) ) - return url + "/mcp" async def _get_http_client(self) -> Any: """Return a shared httpx.AsyncClient, creating one lazily.""" @@ -382,19 +385,30 @@ async def _close_async(self) -> None: """ Close client resources. - In production MCP mode, this also closes the server-side persistent - MCP session (best effort) before closing websocket/provider resources. + In production MCP mode, detach the WebSocket first so + `openenv/session/close` is allowed by the server attachment guard, + then close the HTTP MCP session (best effort). Override `_close_async` rather than `close` so sync teardown (`SyncEnvClient.close`, sync `__exit__`, and `_dispatch`) still cleans up the HTTP MCP session. """ - if self._production_session_id is not None: + session_id = self._production_session_id + if session_id is not None: + # Detach `/ws` before HTTP session close — the server rejects + # `openenv/session/close` while a WebSocket is still attached. + try: + await self._disconnect_async() + except Exception: + pass try: - await self._production_mcp_request( + data = await self._production_mcp_request( "openenv/session/close", - {"session_id": self._production_session_id}, + {"session_id": session_id}, ) + # JSON-RPC errors are HTTP 200; treat them as failed cleanup. + if isinstance(data, dict) and data.get("error"): + pass except Exception: # Best effort cleanup - do not mask normal close behavior pass diff --git a/tests/core/test_mode_selection.py b/tests/core/test_mode_selection.py index 6e9f1cb941..648892d16c 100644 --- a/tests/core/test_mode_selection.py +++ b/tests/core/test_mode_selection.py @@ -166,6 +166,20 @@ def test_invalid_env_var_raises_error(self): class TestModeBehavior: """Test that different modes result in different client behavior.""" + @pytest.mark.parametrize( + ("base_url", "expected_url"), + [ + ("http://localhost:8000", "http://localhost:8000/mcp"), + ("https://example.com/env", "https://example.com/env/mcp"), + ("ws://localhost:8000", "http://localhost:8000/mcp"), + ("wss://example.com/env", "https://example.com/env/mcp"), + ], + ) + def test_production_mcp_url_uses_http_scheme(self, base_url, expected_url): + """HTTP MCP requests normalize WebSocket base URL schemes.""" + client = MCPToolClient(base_url=base_url, mode="production") + assert client._production_mcp_url() == expected_url + @pytest.mark.asyncio async def test_simulation_mode_uses_gym_protocol(self, clean_env, mock_websocket): """Test that simulation mode uses Gym-style WebSocket messages.""" @@ -255,17 +269,6 @@ async def test_production_mode_call_tool_uses_jsonrpc_protocol(self, clean_env): }, ) - def test_production_mcp_url_converts_ws_base_url(self, clean_env): - """ws:// / wss:// base URLs must become http(s) for production /mcp posts.""" - ws_client = MCPToolClient(base_url="ws://localhost:8000", mode="production") - assert ws_client._production_mcp_url() == "http://localhost:8000/mcp" - - wss_client = MCPToolClient(base_url="wss://example.com", mode="production") - assert wss_client._production_mcp_url() == "https://example.com/mcp" - - http_client = MCPToolClient(base_url="http://localhost:8000", mode="production") - assert http_client._production_mcp_url() == "http://localhost:8000/mcp" - @pytest.mark.asyncio async def test_production_mode_connect_creates_single_session_with_websocket( self, clean_env @@ -329,6 +332,34 @@ async def test_production_mode_connect_failure_cleans_up_resources(self, clean_e mock_close.assert_called_once() + @pytest.mark.asyncio + async def test_production_close_disconnects_websocket_before_session_close( + self, clean_env + ): + """Close must detach /ws before openenv/session/close (attachment guard).""" + client = MCPToolClient(base_url="http://localhost:8000", mode="production") + client._production_session_id = "sess-attach" + client._ws = MagicMock() + client._ws_loop = None + + calls = [] + + async def fake_disconnect(): + calls.append("disconnect") + client._ws = None + + async def fake_mcp(method, params=None): + calls.append((method, params)) + return {"result": {"closed": True}} + + with patch.object(client, "_disconnect_async", side_effect=fake_disconnect): + with patch.object(client, "_production_mcp_request", side_effect=fake_mcp): + await client.close() + + assert calls[0] == "disconnect" + assert calls[1] == ("openenv/session/close", {"session_id": "sess-attach"}) + assert client._production_session_id is None + def test_production_mode_sync_close_closes_mcp_session(self, clean_env): """Test that production sync close() closes the MCP session and releases HTTP client.""" client = MCPToolClient( diff --git a/tests/core/test_production_mode_routes.py b/tests/core/test_production_mode_routes.py index e856ed9a6c..7d9c36b8d3 100644 --- a/tests/core/test_production_mode_routes.py +++ b/tests/core/test_production_mode_routes.py @@ -787,7 +787,7 @@ def test_session_create_from_websocket_is_idempotent(self, app): assert response2["data"]["result"]["session_id"] == ws_session_id def test_websocket_can_attach_to_http_session_without_destroying_it(self, app): - """An attached WebSocket shares and preserves an HTTP-created session.""" + """An attached WebSocket shares an HTTP-created session until HTTP close.""" client = TestClient(app) create_response = client.post( "/mcp", @@ -805,18 +805,7 @@ def test_websocket_can_attach_to_http_session_without_destroying_it(self, app): state_response = websocket.receive_json() assert state_response["type"] == "state" - active_close_response = client.post( - "/mcp", - json={ - "jsonrpc": "2.0", - "method": "openenv/session/close", - "params": {"session_id": session_id}, - "id": 2, - }, - ) - assert ( - "active WebSocket" in active_close_response.json()["error"]["message"] - ) + # Detach without destroying the HTTP-owned session. websocket.send_json({"type": "close"}) tools_response = client.post( @@ -841,6 +830,45 @@ def test_websocket_can_attach_to_http_session_without_destroying_it(self, app): ) assert close_response.json()["result"]["closed"] is True + def test_http_session_close_destroys_session_even_while_websocket_attached( + self, app + ): + """HTTP openenv/session/close is authoritative over an attached WebSocket.""" + client = TestClient(app) + create_response = client.post( + "/mcp", + json={ + "jsonrpc": "2.0", + "method": "openenv/session/create", + "params": {}, + "id": 1, + }, + ) + session_id = create_response.json()["result"]["session_id"] + + with client.websocket_connect(f"/ws?session_id={session_id}"): + close_response = client.post( + "/mcp", + json={ + "jsonrpc": "2.0", + "method": "openenv/session/close", + "params": {"session_id": session_id}, + "id": 2, + }, + ) + assert close_response.json()["result"]["closed"] is True + + tools_response = client.post( + "/mcp", + json={ + "jsonrpc": "2.0", + "method": "tools/list", + "params": {"session_id": session_id}, + "id": 3, + }, + ) + assert "error" in tools_response.json() + def test_http_session_allows_only_one_attached_websocket(self, app): """A second WebSocket cannot concurrently mutate the same session.""" client = TestClient(app) From 822945d4d2fa16f77d353ebbbcce6e3ec378175c Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 16 Sep 2026 10:33:22 +0000 Subject: [PATCH 5/8] fix(mcp): adopt #1169 cancel-safe detach-before-close Sync production `_close_async` nesting from fork tip 41cb4e58 (disconnect before session/close; keep provider teardown in finally). Server-side authoritative HTTP close from prior commit remains for attach races. Co-authored-by: mugenkyou --- src/openenv/core/mcp_client.py | 63 +++++++++++++++++----------------- 1 file changed, 31 insertions(+), 32 deletions(-) diff --git a/src/openenv/core/mcp_client.py b/src/openenv/core/mcp_client.py index de5a4ec712..dab8cede53 100644 --- a/src/openenv/core/mcp_client.py +++ b/src/openenv/core/mcp_client.py @@ -385,45 +385,44 @@ async def _close_async(self) -> None: """ Close client resources. - In production MCP mode, detach the WebSocket first so - `openenv/session/close` is allowed by the server attachment guard, - then close the HTTP MCP session (best effort). + In production MCP mode, this also closes the server-side persistent + MCP session (best effort) after detaching the WebSocket and before + closing HTTP/provider resources. Override `_close_async` rather than `close` so sync teardown (`SyncEnvClient.close`, sync `__exit__`, and `_dispatch`) still cleans up the HTTP MCP session. """ - session_id = self._production_session_id - if session_id is not None: - # Detach `/ws` before HTTP session close — the server rejects - # `openenv/session/close` while a WebSocket is still attached. + try: + # The WebSocket shares the HTTP-created session. Detach it first so + # the server's ownership guard permits the explicit session close. + await self._disconnect_async() + finally: try: - await self._disconnect_async() - except Exception: - pass - try: - data = await self._production_mcp_request( - "openenv/session/close", - {"session_id": session_id}, - ) - # JSON-RPC errors are HTTP 200; treat them as failed cleanup. - if isinstance(data, dict) and data.get("error"): - pass - except Exception: - # Best effort cleanup - do not mask normal close behavior - pass - finally: - self._production_session_id = None - - if self._http_client is not None: - try: - await self._http_client.aclose() - except Exception: - pass + if self._production_session_id is not None: + try: + await self._production_mcp_request( + "openenv/session/close", + {"session_id": self._production_session_id}, + ) + except Exception: + # Best effort cleanup - do not mask normal close behavior + pass + finally: + self._production_session_id = None finally: - self._http_client = None - - await super()._close_async() + try: + if self._http_client is not None: + try: + await self._http_client.aclose() + except Exception: + pass + finally: + self._http_client = None + finally: + # This is intentionally inside the outer finally so + # cancellation cannot skip provider teardown. + await super()._close_async() class MCPToolClient(MCPClientBase): From 0fa2399d3a44d707c43729481ac3db6e6a12004a Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 16 Sep 2026 10:35:37 +0000 Subject: [PATCH 6/8] chore(mcp): document best-effort httpx close in production teardown Co-authored-by: benjamin.burtenshaw --- src/openenv/core/mcp_client.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/openenv/core/mcp_client.py b/src/openenv/core/mcp_client.py index dab8cede53..86e30fc6d0 100644 --- a/src/openenv/core/mcp_client.py +++ b/src/openenv/core/mcp_client.py @@ -416,6 +416,7 @@ async def _close_async(self) -> None: try: await self._http_client.aclose() except Exception: + # Best effort — do not mask provider/websocket teardown pass finally: self._http_client = None From 8a71e411d8fec153c3e7f5f86d20a3b7bbad33f0 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 16 Sep 2026 10:41:05 +0000 Subject: [PATCH 7/8] fix(mcp): adopt #1169 deferred close until WebSocket detaches Sync on-repo #1182 with contributor tip 715d5af6: HTTP session/close while attached marks closing and destroys on detach (pending_closes), instead of force-destroying under an active WebSocket. Add disconnect- before-close unit coverage. Co-authored-by: benjamin.burtenshaw --- src/openenv/core/env_server/http_server.py | 38 ++++++++++--- tests/core/test_mode_selection.py | 24 ++++++++ tests/core/test_production_mode_routes.py | 64 ++++++++-------------- 3 files changed, 77 insertions(+), 49 deletions(-) diff --git a/src/openenv/core/env_server/http_server.py b/src/openenv/core/env_server/http_server.py index 425eb27b05..4ddd96a794 100644 --- a/src/openenv/core/env_server/http_server.py +++ b/src/openenv/core/env_server/http_server.py @@ -249,6 +249,7 @@ def __init__( self._session_stacks: Dict[str, AsyncExitStack] = {} self._session_info: Dict[str, SessionInfo] = {} self._session_websocket_attachments: set[str] = set() + self._session_pending_closes: set[str] = set() self._session_lock = asyncio.Lock() # Create thread pool for running sync code in async context @@ -453,6 +454,7 @@ async def _destroy_session(self, session_id: str) -> None: stack = self._session_stacks.pop(session_id, None) self._session_info.pop(session_id, None) self._session_websocket_attachments.discard(session_id) + self._session_pending_closes.discard(session_id) await self._cleanup_session_resources(env, executor, stack) @@ -812,20 +814,34 @@ async def mcp_handler( request_id=request_id, ) - # HTTP `openenv/session/close` is authoritative: drop any - # WebSocket attachment marker and destroy the session. Clients - # should detach first; this still recovers if they race. async with self._session_lock: - self._session_websocket_attachments.discard(target_session_id) - env = self._sessions.pop(target_session_id, _MISSING) - if env is not _MISSING: + if target_session_id in self._session_websocket_attachments: + env = _MISSING + attached = True + self._session_pending_closes.add(target_session_id) + executor = None + stack = None + else: + attached = False + env = self._sessions.pop(target_session_id, _MISSING) + if not attached and env is not _MISSING: executor = self._session_executors.pop(target_session_id, None) stack = self._session_stacks.pop(target_session_id, None) self._session_info.pop(target_session_id, None) - else: + elif not attached: executor = None stack = None + if attached: + return JsonRpcResponse.success( + result={ + "session_id": target_session_id, + "closed": False, + "closing": True, + }, + request_id=request_id, + ) + if env is _MISSING: return JsonRpcResponse.error_response( JsonRpcErrorCode.INVALID_PARAMS, @@ -1732,8 +1748,12 @@ async def websocket_endpoint(websocket: WebSocket): await websocket.send_text(error_resp.model_dump_json()) finally: if attached_session and session_id: - async with self._session_lock: - self._session_websocket_attachments.discard(session_id) + # Do not await the lock before releasing ownership: a task + # cancelled while waiting would leave this session + # permanently exempt from idle reaping. + self._session_websocket_attachments.discard(session_id) + if session_id in self._session_pending_closes: + await asyncio.shield(self._destroy_session(session_id)) elif owns_session and session_id: await self._destroy_session(session_id) try: diff --git a/tests/core/test_mode_selection.py b/tests/core/test_mode_selection.py index 648892d16c..fdf4fc48f0 100644 --- a/tests/core/test_mode_selection.py +++ b/tests/core/test_mode_selection.py @@ -439,6 +439,30 @@ async def test_production_mode_async_close_closes_mcp_session(self, clean_env): mock_http_client.aclose.assert_awaited_once() assert client._http_client is None + @pytest.mark.asyncio + async def test_production_close_detaches_websocket_before_session_close( + self, clean_env + ): + """Shared WebSocket ownership is released before HTTP session teardown.""" + client = MCPToolClient(base_url="http://localhost:8000", mode="production") + client._production_session_id = "test-session" + teardown_events = [] + + async def disconnect(): + teardown_events.append("websocket") + + async def request(method, params=None): + teardown_events.append("session") + return {"result": {"closed": True}} + + with ( + patch.object(client, "_disconnect_async", side_effect=disconnect), + patch.object(client, "_production_mcp_request", side_effect=request), + ): + await client.close() + + assert teardown_events[:2] == ["websocket", "session"] + # ============================================================================ # Mode Immutability Tests diff --git a/tests/core/test_production_mode_routes.py b/tests/core/test_production_mode_routes.py index 7d9c36b8d3..737defd188 100644 --- a/tests/core/test_production_mode_routes.py +++ b/tests/core/test_production_mode_routes.py @@ -787,7 +787,7 @@ def test_session_create_from_websocket_is_idempotent(self, app): assert response2["data"]["result"]["session_id"] == ws_session_id def test_websocket_can_attach_to_http_session_without_destroying_it(self, app): - """An attached WebSocket shares an HTTP-created session until HTTP close.""" + """An attached WebSocket shares and preserves an HTTP-created session.""" client = TestClient(app) create_response = client.post( "/mcp", @@ -805,7 +805,21 @@ def test_websocket_can_attach_to_http_session_without_destroying_it(self, app): state_response = websocket.receive_json() assert state_response["type"] == "state" - # Detach without destroying the HTTP-owned session. + active_close_response = client.post( + "/mcp", + json={ + "jsonrpc": "2.0", + "method": "openenv/session/close", + "params": {"session_id": session_id}, + "id": 2, + }, + ) + close_result = active_close_response.json()["result"] + assert close_result == { + "session_id": session_id, + "closed": False, + "closing": True, + } websocket.send_json({"type": "close"}) tools_response = client.post( @@ -817,57 +831,27 @@ def test_websocket_can_attach_to_http_session_without_destroying_it(self, app): "id": 3, }, ) - assert "result" in tools_response.json() - - close_response = client.post( - "/mcp", - json={ - "jsonrpc": "2.0", - "method": "openenv/session/close", - "params": {"session_id": session_id}, - "id": 4, - }, - ) - assert close_response.json()["result"]["closed"] is True + assert tools_response.json()["error"]["code"] == -32602 - def test_http_session_close_destroys_session_even_while_websocket_attached( - self, app - ): - """HTTP openenv/session/close is authoritative over an attached WebSocket.""" - client = TestClient(app) - create_response = client.post( + replacement_response = client.post( "/mcp", json={ "jsonrpc": "2.0", "method": "openenv/session/create", "params": {}, - "id": 1, + "id": 4, }, ) - session_id = create_response.json()["result"]["session_id"] - - with client.websocket_connect(f"/ws?session_id={session_id}"): - close_response = client.post( - "/mcp", - json={ - "jsonrpc": "2.0", - "method": "openenv/session/close", - "params": {"session_id": session_id}, - "id": 2, - }, - ) - assert close_response.json()["result"]["closed"] is True - - tools_response = client.post( + replacement_id = replacement_response.json()["result"]["session_id"] + client.post( "/mcp", json={ "jsonrpc": "2.0", - "method": "tools/list", - "params": {"session_id": session_id}, - "id": 3, + "method": "openenv/session/close", + "params": {"session_id": replacement_id}, + "id": 5, }, ) - assert "error" in tools_response.json() def test_http_session_allows_only_one_attached_websocket(self, app): """A second WebSocket cannot concurrently mutate the same session.""" From d7c9a7e0fef4a150fb3eea10a93ef3a24070d888 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 16 Sep 2026 10:45:41 +0000 Subject: [PATCH 8/8] fix(mcp): cancel-safe production connect and refuse attach while closing Catch BaseException on production connect so CancelledError after HTTP session create still tears down capacity. Reject WebSocket attach when session is in _session_pending_closes to close the reattach/destroy race. Co-authored-by: benjamin.burtenshaw --- src/openenv/core/env_server/http_server.py | 4 +++ src/openenv/core/mcp_client.py | 9 ++++-- tests/core/test_mode_selection.py | 36 ++++++++++++++++++++++ tests/core/test_production_mode_routes.py | 34 ++++++++++++++++++++ 4 files changed, 81 insertions(+), 2 deletions(-) diff --git a/src/openenv/core/env_server/http_server.py b/src/openenv/core/env_server/http_server.py index 4ddd96a794..14d8d4fb3a 100644 --- a/src/openenv/core/env_server/http_server.py +++ b/src/openenv/core/env_server/http_server.py @@ -1544,6 +1544,10 @@ async def websocket_endpoint(websocket: WebSocket): raise RuntimeError( f"Session {requested_session_id} is still initializing" ) + if requested_session_id in self._session_pending_closes: + raise RuntimeError( + f"Session {requested_session_id} is closing" + ) if requested_session_id in self._session_websocket_attachments: raise RuntimeError( f"Session {requested_session_id} already has " diff --git a/src/openenv/core/mcp_client.py b/src/openenv/core/mcp_client.py index 86e30fc6d0..f9d73ab262 100644 --- a/src/openenv/core/mcp_client.py +++ b/src/openenv/core/mcp_client.py @@ -233,8 +233,13 @@ async def _connect_async(self) -> EnvClient: await super()._connect_async() finally: self._ws_url = original_ws_url - except Exception: - await self.close() + except BaseException: + # CancelledError is a BaseException: cleanup must still run + # after session create so capacity is not leaked. + try: + await asyncio.shield(self.close()) + except Exception: + pass raise return self diff --git a/tests/core/test_mode_selection.py b/tests/core/test_mode_selection.py index fdf4fc48f0..fa0a8854a6 100644 --- a/tests/core/test_mode_selection.py +++ b/tests/core/test_mode_selection.py @@ -21,6 +21,7 @@ - Environment: Code mode with mode-aware tool registration """ +import asyncio import os from unittest.mock import AsyncMock, MagicMock, patch @@ -463,6 +464,41 @@ async def request(method, params=None): assert teardown_events[:2] == ["websocket", "session"] + @pytest.mark.asyncio + async def test_production_connect_cancellation_closes_allocated_session( + self, clean_env + ): + """Cancelled WebSocket connect still releases the HTTP MCP session.""" + client = MCPToolClient(base_url="http://localhost:8000", mode="production") + client._ws_url = "ws://localhost:8000/ws" + close_calls = [] + + async def ensure_session(): + client._production_session_id = "cancelled-session" + return "cancelled-session" + + async def connect(*_args, **_kwargs): + raise asyncio.CancelledError() + + async def close(*_args, **_kwargs): + close_calls.append(client._production_session_id) + + with ( + patch.object(client, "_start_provider_if_needed"), + patch.object( + client, "_ensure_production_session", side_effect=ensure_session + ), + patch( + "openenv.core.env_client.EnvClient._connect_async", + side_effect=connect, + ), + patch.object(client, "close", new=close), + ): + with pytest.raises(asyncio.CancelledError): + await client._connect_async() + + assert close_calls == ["cancelled-session"] + # ============================================================================ # Mode Immutability Tests diff --git a/tests/core/test_production_mode_routes.py b/tests/core/test_production_mode_routes.py index 737defd188..a528806e74 100644 --- a/tests/core/test_production_mode_routes.py +++ b/tests/core/test_production_mode_routes.py @@ -889,6 +889,40 @@ def test_http_session_allows_only_one_attached_websocket(self, app): ) assert close_response.json()["result"]["closed"] is True + def test_websocket_cannot_attach_while_session_close_is_pending(self, app): + """Pending closes refuse reattach so delayed destroy cannot race a new WS.""" + client = TestClient(app) + create_response = client.post( + "/mcp", + json={ + "jsonrpc": "2.0", + "method": "openenv/session/create", + "params": {}, + "id": 1, + }, + ) + session_id = create_response.json()["result"]["session_id"] + + with client.websocket_connect(f"/ws?session_id={session_id}") as websocket: + close_response = client.post( + "/mcp", + json={ + "jsonrpc": "2.0", + "method": "openenv/session/close", + "params": {"session_id": session_id}, + "id": 2, + }, + ) + assert close_response.json()["result"]["closing"] is True + + with client.websocket_connect( + f"/ws?session_id={session_id}" + ) as second_socket: + error_response = second_socket.receive_json() + assert "is closing" in error_response["data"]["message"] + + websocket.send_json({"type": "close"}) + def test_websocket_still_destroys_its_own_session(self, app): """A WebSocket-created session is destroyed when the socket closes.""" client = TestClient(app)