fix: route MCP production mode through JSON-RPC

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,
@@ -1487,18 +1510,48 @@ 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.
+            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()
 
             session_id = None
             session_env = None
+            owns_session = False
+            attached_session = False
 
             try:
-                # Create session with dedicated environment
-                session_id, session_env = await self._create_session()
+                requested_session_id = websocket.query_params.get("session_id")
+                if requested_session_id:
+                    async with self._session_lock:
+                        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"
+                            )
+                        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()
+                    owns_session = True
+
                 if session_env is None:
                     raise RuntimeError(
                         "Session environment not initialized for websocket"
@@ -1509,7 +1562,7 @@ async def websocket_endpoint(websocket: WebSocket):
 
                 async with AsyncExitStack() as stack:
                     mcp_session_factory = getattr(session_env, "mcp_session", None)
-                    if callable(mcp_session_factory):
+                    if owns_session and callable(mcp_session_factory):
                         mcp_session_cm = cast(
                             AsyncContextManager[Any], mcp_session_factory()
                         )
@@ -1688,7 +1741,10 @@ async def websocket_endpoint(websocket: WebSocket):
                 )
                 await websocket.send_text(error_resp.model_dump_json())
             finally:
-                if 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()

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
@@ -166,11 +168,10 @@ 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://")
-        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.")
+        return self._base_url.rstrip("/") + "/mcp"
 
     async def _get_http_client(self) -> Any:
         """Return a shared httpx.AsyncClient, creating one lazily."""
@@ -202,19 +203,30 @@ 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), 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:
-                await super()._connect_async()
-                await self._ensure_production_session()
-            except Exception:
-                await self.close()
-                raise
+            async with self._production_connect_lock:
+                try:
+                    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()