From d08b2499b37f4ee259ebdc84b81d110358049caf Mon Sep 17 00:00:00 2001
From: mugenkyou <sachinskyte@gmail.com>
Date: Wed, 16 Sep 2026 09:22:42 +0530
Subject: [PATCH 01/11] fix: route MCP production mode through JSON-RPC

---
 src/openenv/core/mcp_client.py    |  2 +-
 tests/core/test_mode_selection.py | 89 ++++++++++++++++++++-----------
 2 files changed, 60 insertions(+), 31 deletions(-)

diff --git a/src/openenv/core/mcp_client.py b/src/openenv/core/mcp_client.py
index 7634bb0be9..cc5a23ca5e 100644
--- a/src/openenv/core/mcp_client.py
+++ b/src/openenv/core/mcp_client.py
@@ -154,7 +154,7 @@ def __init__(
             mode=mode,
         )
         self._tools_cache: Optional[List[Tool]] = None
-        self.use_production_mode = False
+        self.use_production_mode = self._mode == "production"
         self._production_session_id: Optional[str] = None
         self._production_session_lock = asyncio.Lock()
         self._jsonrpc_request_id = 0
diff --git a/tests/core/test_mode_selection.py b/tests/core/test_mode_selection.py
index cbbf543d48..0eea0d6be0 100644
--- a/tests/core/test_mode_selection.py
+++ b/tests/core/test_mode_selection.py
@@ -193,39 +193,67 @@ async def test_simulation_mode_uses_gym_protocol(self, clean_env, mock_websocket
                     )
 
     @pytest.mark.asyncio
-    async def test_production_mode_uses_jsonrpc_protocol(
-        self, clean_env, mock_websocket
-    ):
-        """Test that production mode uses JSON-RPC format for tool calls."""
+    async def test_production_mode_uses_jsonrpc_protocol(self, clean_env):
+        """Test that production mode uses HTTP JSON-RPC format for tool listing."""
         client = MCPToolClient(base_url="http://localhost:8000", mode="production")
-
-        with patch.object(client, "_send") as mock_send:
-            with patch.object(
-                client,
-                "_receive",
-                return_value={
-                    "type": "response",
-                    "data": {
-                        "observation": {"tools": []},
-                        "reward": None,
-                        "done": False,
-                    },
+        assert client.use_production_mode is True
+
+        with patch.object(
+            client,
+            "_production_mcp_request",
+            side_effect=[
+                {"result": {"session_id": "test-session"}},
+                {
+                    "result": {
+                        "tools": [
+                            {
+                                "name": "echo",
+                                "description": "Echo message",
+                                "inputSchema": {},
+                            }
+                        ]
+                    }
                 },
-            ):
-                with patch.object(client, "_ws", mock_websocket):
-                    await client.list_tools()
-
-                    # Should send step message with list_tools action
-                    call_args = mock_send.call_args_list
-                    step_call = [
-                        call for call in call_args if call[0][0].get("type") == "step"
-                    ]
-                    assert len(step_call) > 0, "Should send message with type='step'"
+            ],
+        ) as mock_mcp_request:
+            with patch.object(client, "step") as mock_step:
+                tools = await client.list_tools()
+
+                mock_step.assert_not_called()
+                assert len(tools) == 1
+                assert tools[0].name == "echo"
+                assert mock_mcp_request.call_count == 2
+                mock_mcp_request.assert_called_with(
+                    "tools/list", {"session_id": "test-session"}
+                )
 
-                    # Check that the action payload is list_tools
-                    step_message = step_call[0][0][0]
-                    assert "data" in step_message
-                    assert step_message["data"].get("type") == "list_tools"
+    @pytest.mark.asyncio
+    async def test_production_mode_call_tool_uses_jsonrpc_protocol(self, clean_env):
+        """Test that call_tool in production mode uses HTTP JSON-RPC transport."""
+        client = MCPToolClient(base_url="http://localhost:8000", mode="production")
+        assert client.use_production_mode is True
+
+        with patch.object(
+            client,
+            "_production_mcp_request",
+            side_effect=[
+                {"result": {"session_id": "test-session"}},
+                {"result": {"data": "hello world"}},
+            ],
+        ) as mock_mcp_request:
+            with patch.object(client, "step") as mock_step:
+                result = await client.call_tool("echo", message="hello world")
+
+                mock_step.assert_not_called()
+                assert result == "hello world"
+                mock_mcp_request.assert_called_with(
+                    "tools/call",
+                    {
+                        "name": "echo",
+                        "arguments": {"message": "hello world"},
+                        "session_id": "test-session",
+                    },
+                )
 
 
 # ============================================================================
@@ -283,6 +311,7 @@ def test_mcp_client_defaults_to_production_mode(self, clean_env):
 
         # MCPToolClient should default to production mode
         assert client._mode == "production"
+        assert client.use_production_mode is True
 
     def test_mcp_client_cannot_use_simulation_mode(self, clean_env):
         """Test that MCPToolClient raises error if simulation mode is requested."""

From f996142159fc6f955e45834a2f6892b1c471b62a Mon Sep 17 00:00:00 2001
From: mugenkyou <sachinskyte@gmail.com>
Date: Wed, 16 Sep 2026 09:57:34 +0530
Subject: [PATCH 02/11] fix: avoid duplicate MCP sessions in production mode

---
 src/openenv/core/mcp_client.py    | 18 ++++++++++++++
 tests/core/test_mode_selection.py | 39 +++++++++++++++++++++++++++++++
 2 files changed, 57 insertions(+)

diff --git a/src/openenv/core/mcp_client.py b/src/openenv/core/mcp_client.py
index cc5a23ca5e..d9815fa658 100644
--- a/src/openenv/core/mcp_client.py
+++ b/src/openenv/core/mcp_client.py
@@ -198,6 +198,24 @@ async def _production_mcp_request(
         response.raise_for_status()
         return response.json()
 
+    async def _connect_async(self) -> EnvClient:
+        """
+        Establish connection to the server.
+
+        In production mode (use_production_mode=True), creates and caches a
+        persistent HTTP MCP session instead of establishing a WebSocket connection.
+        """
+        if getattr(self, "use_production_mode", False):
+            try:
+                self._start_provider_if_needed()
+            except Exception:
+                await self.close()
+                raise
+            await self._ensure_production_session()
+            return self
+
+        return await super()._connect_async()
+
     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/core/test_mode_selection.py b/tests/core/test_mode_selection.py
index 0eea0d6be0..de9498081d 100644
--- a/tests/core/test_mode_selection.py
+++ b/tests/core/test_mode_selection.py
@@ -255,6 +255,45 @@ async def test_production_mode_call_tool_uses_jsonrpc_protocol(self, clean_env):
                     },
                 )
 
+    @pytest.mark.asyncio
+    async def test_production_mode_connect_creates_http_session_without_websocket(
+        self, clean_env
+    ):
+        """Test that connect() in production mode initializes the HTTP MCP session without creating a WebSocket."""
+        client = MCPToolClient(base_url="http://localhost:8000", mode="production")
+        assert client.use_production_mode is True
+
+        with patch.object(
+            client,
+            "_production_mcp_request",
+            side_effect=[
+                {"result": {"session_id": "test-session"}},
+                {"result": {"data": "hello world"}},
+            ],
+        ) as mock_mcp_request:
+            with patch("openenv.core.env_client.ws_connect") as mock_ws_connect:
+                # Explicit connect (e.g. from async with client:)
+                await client.connect()
+
+                # Should create HTTP session and not connect WS
+                mock_ws_connect.assert_not_called()
+                assert client._ws is None
+                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
+                mock_mcp_request.assert_called_with(
+                    "tools/call",
+                    {
+                        "name": "echo",
+                        "arguments": {"message": "hello world"},
+                        "session_id": "test-session",
+                    },
+                )
+
 
 # ============================================================================
 # Mode Immutability Tests

From d1653233c7439bfaa9880ecd3acb785f75ea00f3 Mon Sep 17 00:00:00 2001
From: mugenkyou <sachinskyte@gmail.com>
Date: Wed, 16 Sep 2026 10:11:21 +0530
Subject: [PATCH 03/11] fix: clean up production MCP connection failures

---
 src/openenv/core/mcp_client.py    |  2 +-
 tests/core/test_mode_selection.py | 19 +++++++++++++++++++
 2 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/src/openenv/core/mcp_client.py b/src/openenv/core/mcp_client.py
index d9815fa658..bb1c48d196 100644
--- a/src/openenv/core/mcp_client.py
+++ b/src/openenv/core/mcp_client.py
@@ -208,10 +208,10 @@ async def _connect_async(self) -> EnvClient:
         if getattr(self, "use_production_mode", False):
             try:
                 self._start_provider_if_needed()
+                await self._ensure_production_session()
             except Exception:
                 await self.close()
                 raise
-            await self._ensure_production_session()
             return self
 
         return await super()._connect_async()
diff --git a/tests/core/test_mode_selection.py b/tests/core/test_mode_selection.py
index de9498081d..613dd296a3 100644
--- a/tests/core/test_mode_selection.py
+++ b/tests/core/test_mode_selection.py
@@ -294,6 +294,25 @@ async def test_production_mode_connect_creates_http_session_without_websocket(
                     },
                 )
 
+    @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."""
+        client = MCPToolClient(base_url="http://localhost:8000", mode="production")
+        assert client.use_production_mode is True
+
+        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()
+
 
 # ============================================================================
 # Mode Immutability Tests

From 1244762a6210728d7dea5c63e09437b69c5c4118 Mon Sep 17 00:00:00 2001
From: mugenkyou <sachinskyte@gmail.com>
Date: Wed, 16 Sep 2026 13:53:32 +0530
Subject: [PATCH 04/11] fix: clean up production MCP sessions on close

---
 src/openenv/core/mcp_client.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/openenv/core/mcp_client.py b/src/openenv/core/mcp_client.py
index bb1c48d196..0c857582fe 100644
--- a/src/openenv/core/mcp_client.py
+++ b/src/openenv/core/mcp_client.py
@@ -357,7 +357,7 @@ def _parse_state(self, payload: Dict[str, Any]) -> State:
             step_count=payload.get("step_count", 0),
         )
 
-    async def close(self) -> None:
+    async def _close_async(self) -> None:
         """
         Close client resources.
 
@@ -384,7 +384,7 @@ async def close(self) -> None:
             finally:
                 self._http_client = None
 
-        await super().close()
+        await super()._close_async()
 
 
 class MCPToolClient(MCPClientBase):

From 434ac3da775f67397715cce51b16b7a08cbae1fb Mon Sep 17 00:00:00 2001
From: mugenkyou <sachinskyte@gmail.com>
Date: Wed, 16 Sep 2026 15:37:03 +0530
Subject: [PATCH 05/11] fix: handle production MCP session lifecycle

---
 src/openenv/core/mcp_client.py    |  15 ++-
 tests/core/test_mode_selection.py | 164 ++++++++++++++++++++++++++++--
 2 files changed, 166 insertions(+), 13 deletions(-)

diff --git a/src/openenv/core/mcp_client.py b/src/openenv/core/mcp_client.py
index 0c857582fe..4d42dfff6d 100644
--- a/src/openenv/core/mcp_client.py
+++ b/src/openenv/core/mcp_client.py
@@ -202,13 +202,22 @@ async def _connect_async(self) -> EnvClient:
         """
         Establish connection to the server.
 
-        In production mode (use_production_mode=True), creates and caches a
-        persistent HTTP MCP session instead of establishing a WebSocket connection.
+        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:
                 self._start_provider_if_needed()
-                await self._ensure_production_session()
+                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 613dd296a3..fb8a3286f9 100644
--- a/tests/core/test_mode_selection.py
+++ b/tests/core/test_mode_selection.py
@@ -22,7 +22,7 @@
 """
 
 import os
-from unittest.mock import MagicMock, patch
+from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
 from fastmcp import FastMCP
@@ -256,10 +256,10 @@ async def test_production_mode_call_tool_uses_jsonrpc_protocol(self, clean_env):
                 )
 
     @pytest.mark.asyncio
-    async def test_production_mode_connect_creates_http_session_without_websocket(
+    async def test_production_mode_connect_creates_single_session_with_websocket(
         self, clean_env
     ):
-        """Test that connect() in production mode initializes the HTTP MCP session without creating a 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
 
@@ -271,13 +271,15 @@ async def test_production_mode_connect_creates_http_session_without_websocket(
                 {"result": {"data": "hello world"}},
             ],
         ) as mock_mcp_request:
-            with patch("openenv.core.env_client.ws_connect") as mock_ws_connect:
+            with patch(
+                "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 not connect WS
-                mock_ws_connect.assert_not_called()
-                assert client._ws is None
+                # 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]
                 assert client._production_session_id == "test-session"
                 mock_mcp_request.assert_called_once_with("openenv/session/create")
 
@@ -295,9 +297,7 @@ async def test_production_mode_connect_creates_http_session_without_websocket(
                 )
 
     @pytest.mark.asyncio
-    async def test_production_mode_connect_failure_cleans_up_resources(
-        self, clean_env
-    ):
+    async def test_production_mode_connect_failure_cleans_up_resources(self, clean_env):
         """Test that failure during production mode connect() triggers client.close() cleanup."""
         client = MCPToolClient(base_url="http://localhost:8000", mode="production")
         assert client.use_production_mode is True
@@ -313,6 +313,150 @@ async def test_production_mode_connect_failure_cleans_up_resources(
 
                 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"
+
+        mock_http_client = AsyncMock()
+        client._async._http_client = mock_http_client
+
+        with patch.object(
+            client._async,
+            "_production_mcp_request",
+            new_callable=AsyncMock,
+            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
+
+        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_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")
+        client._production_session_id = "test-session-async"
+
+        mock_http_client = AsyncMock()
+        client._http_client = mock_http_client
+
+        with patch.object(
+            client,
+            "_production_mcp_request",
+            new_callable=AsyncMock,
+            return_value={"result": {"status": "closed"}},
+        ) as mock_mcp_req:
+            await 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
+
+    @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
+        from tests.core.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 tests.core.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

From a8eec7583f99109955f696d4e3d4fbe9c93dc887 Mon Sep 17 00:00:00 2001
From: Cursor Agent <cursoragent@cursor.com>
Date: Wed, 16 Sep 2026 10:17:38 +0000
Subject: [PATCH 06/11] test: cover shared MCP WebSocket sessions

Co-authored-by: burtenshaw <burtenshaw@users.noreply.github.com>
---
 tests/core/test_mode_selection.py         | 65 --------------------
 tests/core/test_production_mode_routes.py | 72 +++++++++++++++++++++++
 2 files changed, 72 insertions(+), 65 deletions(-)

diff --git a/tests/core/test_mode_selection.py b/tests/core/test_mode_selection.py
index fb8a3286f9..88005aa077 100644
--- a/tests/core/test_mode_selection.py
+++ b/tests/core/test_mode_selection.py
@@ -392,71 +392,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
-        from tests.core.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 tests.core.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 77a9915f03ee6e67b74bfce8c74fe43abdb3e889 Mon Sep 17 00:00:00 2001
From: Cursor Agent <cursoragent@cursor.com>
Date: Wed, 16 Sep 2026 10:23:59 +0000
Subject: [PATCH 07/11] fix: guard shared MCP session ownership

Co-authored-by: burtenshaw <burtenshaw@users.noreply.github.com>
---
 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 43f7812595..0f22427dd4 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
@@ -206,20 +208,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 88005aa077..b168993c20 100644
--- a/tests/core/test_mode_selection.py
+++ b/tests/core/test_mode_selection.py
@@ -262,6 +262,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,
@@ -279,7 +281,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 17006d520232895677c2b4a262f080a864fb1aa7 Mon Sep 17 00:00:00 2001
From: Cursor Agent <cursoragent@cursor.com>
Date: Wed, 16 Sep 2026 10:25:22 +0000
Subject: [PATCH 08/11] fix: normalize WebSocket MCP URL schemes

Co-authored-by: burtenshaw <burtenshaw@users.noreply.github.com>
---
 src/openenv/core/mcp_client.py    | 11 ++++++++++-
 tests/core/test_mode_selection.py | 14 ++++++++++++++
 2 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/src/openenv/core/mcp_client.py b/src/openenv/core/mcp_client.py
index 0f22427dd4..89cccf8e5e 100644
--- a/src/openenv/core/mcp_client.py
+++ b/src/openenv/core/mcp_client.py
@@ -171,7 +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.")
-        return self._base_url.rstrip("/") + "/mcp"
+        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="",
+            )
+        )
 
     async def _get_http_client(self) -> Any:
         """Return a shared httpx.AsyncClient, creating one lazily."""
diff --git a/tests/core/test_mode_selection.py b/tests/core/test_mode_selection.py
index b168993c20..7ed1a86a0a 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."""

From 41cb4e5858b9c46fb16dfd83aacb24283852b8b4 Mon Sep 17 00:00:00 2001
From: Cursor Agent <cursoragent@cursor.com>
Date: Wed, 16 Sep 2026 10:31:38 +0000
Subject: [PATCH 09/11] fix: detach WebSocket before closing MCP session

Co-authored-by: burtenshaw <burtenshaw@users.noreply.github.com>
---
 src/openenv/core/mcp_client.py            | 50 ++++++++++++++---------
 tests/core/test_mode_selection.py         | 24 +++++++++++
 tests/core/test_production_mode_routes.py | 20 +++++++++
 3 files changed, 74 insertions(+), 20 deletions(-)

diff --git a/src/openenv/core/mcp_client.py b/src/openenv/core/mcp_client.py
index 89cccf8e5e..dab8cede53 100644
--- a/src/openenv/core/mcp_client.py
+++ b/src/openenv/core/mcp_client.py
@@ -386,33 +386,43 @@ 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.
+        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.
         """
-        if self._production_session_id is not None:
+        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._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
-
-        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):
diff --git a/tests/core/test_mode_selection.py b/tests/core/test_mode_selection.py
index 7ed1a86a0a..e2b8b44e4f 100644
--- a/tests/core/test_mode_selection.py
+++ b/tests/core/test_mode_selection.py
@@ -411,6 +411,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 e856ed9a6c..2784a4ab0e 100644
--- a/tests/core/test_production_mode_routes.py
+++ b/tests/core/test_production_mode_routes.py
@@ -841,6 +841,26 @@ def test_websocket_can_attach_to_http_session_without_destroying_it(self, app):
         )
         assert close_response.json()["result"]["closed"] is True
 
+        replacement_response = client.post(
+            "/mcp",
+            json={
+                "jsonrpc": "2.0",
+                "method": "openenv/session/create",
+                "params": {},
+                "id": 5,
+            },
+        )
+        replacement_id = replacement_response.json()["result"]["session_id"]
+        client.post(
+            "/mcp",
+            json={
+                "jsonrpc": "2.0",
+                "method": "openenv/session/close",
+                "params": {"session_id": replacement_id},
+                "id": 6,
+            },
+        )
+
     def test_http_session_allows_only_one_attached_websocket(self, app):
         """A second WebSocket cannot concurrently mutate the same session."""
         client = TestClient(app)

From 715d5af6f1a5af08bcd8c95dac1500b36be8ee1f Mon Sep 17 00:00:00 2001
From: Cursor Agent <cursoragent@cursor.com>
Date: Wed, 16 Sep 2026 10:37:32 +0000
Subject: [PATCH 10/11] fix: defer MCP session close until WebSocket detaches

Co-authored-by: burtenshaw <burtenshaw@users.noreply.github.com>
---
 src/openenv/core/env_server/http_server.py | 20 ++++++++++++-----
 tests/core/test_production_mode_routes.py  | 26 ++++++++--------------
 2 files changed, 24 insertions(+), 22 deletions(-)

diff --git a/src/openenv/core/env_server/http_server.py b/src/openenv/core/env_server/http_server.py
index a3e242bc2d..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)
 
@@ -816,6 +818,7 @@ async def mcp_handler(
                     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:
@@ -830,9 +833,12 @@ async def mcp_handler(
                         stack = None
 
                 if attached:
-                    return JsonRpcResponse.error_response(
-                        JsonRpcErrorCode.INVALID_REQUEST,
-                        f"Session {target_session_id} has an active WebSocket",
+                    return JsonRpcResponse.success(
+                        result={
+                            "session_id": target_session_id,
+                            "closed": False,
+                            "closing": True,
+                        },
                         request_id=request_id,
                     )
 
@@ -1742,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_production_mode_routes.py b/tests/core/test_production_mode_routes.py
index 2784a4ab0e..737defd188 100644
--- a/tests/core/test_production_mode_routes.py
+++ b/tests/core/test_production_mode_routes.py
@@ -814,9 +814,12 @@ def test_websocket_can_attach_to_http_session_without_destroying_it(self, app):
                     "id": 2,
                 },
             )
-            assert (
-                "active WebSocket" in active_close_response.json()["error"]["message"]
-            )
+            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(
@@ -828,18 +831,7 @@ 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
 
         replacement_response = client.post(
             "/mcp",
@@ -847,7 +839,7 @@ def test_websocket_can_attach_to_http_session_without_destroying_it(self, app):
                 "jsonrpc": "2.0",
                 "method": "openenv/session/create",
                 "params": {},
-                "id": 5,
+                "id": 4,
             },
         )
         replacement_id = replacement_response.json()["result"]["session_id"]
@@ -857,7 +849,7 @@ def test_websocket_can_attach_to_http_session_without_destroying_it(self, app):
                 "jsonrpc": "2.0",
                 "method": "openenv/session/close",
                 "params": {"session_id": replacement_id},
-                "id": 6,
+                "id": 5,
             },
         )
 

From 00173d4796e13c8540c453208c919199d8409b21 Mon Sep 17 00:00:00 2001
From: Cursor Agent <cursoragent@cursor.com>
Date: Wed, 16 Sep 2026 10:43:35 +0000
Subject: [PATCH 11/11] fix: clean up cancelled MCP connections

Co-authored-by: burtenshaw <burtenshaw@users.noreply.github.com>
---
 src/openenv/core/mcp_client.py    |  5 ++++-
 tests/core/test_mode_selection.py | 36 +++++++++++++++++++++++++++++++
 2 files changed, 40 insertions(+), 1 deletion(-)

diff --git a/src/openenv/core/mcp_client.py b/src/openenv/core/mcp_client.py
index dab8cede53..e172f32eef 100644
--- a/src/openenv/core/mcp_client.py
+++ b/src/openenv/core/mcp_client.py
@@ -233,7 +233,10 @@ async def _connect_async(self) -> EnvClient:
                         await super()._connect_async()
                     finally:
                         self._ws_url = original_ws_url
-                except Exception:
+                except BaseException:
+                    # Cancellation after the HTTP session is allocated must
+                    # release that session and any started provider before the
+                    # cancellation propagates.
                     await self.close()
                     raise
             return self
diff --git a/tests/core/test_mode_selection.py b/tests/core/test_mode_selection.py
index e2b8b44e4f..ad9465f171 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
 
@@ -332,6 +333,41 @@ async def test_production_mode_connect_failure_cleans_up_resources(self, clean_e
 
                 mock_close.assert_called_once()
 
+    @pytest.mark.asyncio
+    async def test_cancelled_production_connect_closes_allocated_session(
+        self, clean_env
+    ):
+        """Cancellation during WebSocket connect releases the HTTP session."""
+        client = MCPToolClient(base_url="http://localhost:8000", mode="production")
+        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": {"session_id": "test-session", "closed": True}},
+                ],
+            ) as mock_mcp_request,
+            patch(
+                "openenv.core.env_client.ws_connect",
+                new_callable=AsyncMock,
+                side_effect=asyncio.CancelledError,
+            ),
+        ):
+            with pytest.raises(asyncio.CancelledError):
+                await client.connect()
+
+        mock_mcp_request.assert_any_call(
+            "openenv/session/close",
+            {"session_id": "test-session"},
+        )
+        assert client._production_session_id is None
+        mock_http_client.aclose.assert_awaited_once()
+        assert client._http_client 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(