fix: route MCP production mode through JSON-RPC

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
@@ -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()
+                await self._ensure_production_session()
+            except Exception:
+                await self.close()
+                raise
+            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:

tests/core/test_mode_selection.py

@@ -193,39 +193,125 @@ 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")
+        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": {},
+                            }
+                        ]
+                    }
+                },
+            ],
+        ) 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"}
+                )
 
-        with patch.object(client, "_send") as mock_send:
-            with patch.object(
-                client,
-                "_receive",
-                return_value={
-                    "type": "response",
-                    "data": {
-                        "observation": {"tools": []},
-                        "reward": None,
-                        "done": False,
+    @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",
                     },
-                },
-            ):
-                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'"
+    @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",
+                    },
+                )
+
+    @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()
 
-                    # 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"
+                mock_close.assert_called_once()
 
 
 # ============================================================================
@@ -283,6 +369,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."""