fix(mcp): share one production session across HTTP MCP and WebSocket (from #1169) #1182
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -248,6 +248,8 @@ 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_pending_closes: set[str] = set() | ||
| self._session_lock = asyncio.Lock() | ||
|
|
||
| # Create thread pool for running sync code in async context | ||
|
|
@@ -451,6 +453,8 @@ 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) | ||
| self._session_pending_closes.discard(session_id) | ||
|
|
||
| await self._cleanup_session_resources(env, executor, stack) | ||
|
|
||
|
|
@@ -522,7 +526,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 ( | ||
| 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 +539,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 sid in self._session_websocket_attachments | ||
| or (now - info.last_activity_at) <= timeout | ||
| ): | ||
| continue | ||
| await self._destroy_session(sid) | ||
| except asyncio.CancelledError: | ||
|
|
@@ -804,15 +815,33 @@ async def mcp_handler( | |
| ) | ||
|
|
||
| async with self._session_lock: | ||
| 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) | ||
| 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, | ||
|
|
@@ -1487,18 +1516,52 @@ async def websocket_endpoint(websocket: WebSocket): | |
| """ | ||
| WebSocket endpoint for persistent environment sessions. | ||
|
|
||
| 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: | ||
| 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_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 " | ||
| "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 +1572,7 @@ async def websocket_endpoint(websocket: WebSocket): | |
|
|
||
| async with AsyncExitStack() as stack: | ||
| mcp_session_factory = getattr(session_env, "mcp_session", None) | ||
| if owns_session and callable(mcp_session_factory): | ||
|
Contributor
Author
There was a problem hiding this comment. The Follow-up (pre-existing, non-blocking): for the |
||
| mcp_session_cm = cast( | ||
| AsyncContextManager[Any], mcp_session_factory() | ||
| ) | ||
|
|
@@ -1688,7 +1751,14 @@ async def websocket_endpoint(websocket: WebSocket): | |
| ) | ||
| await websocket.send_text(error_resp.model_dump_json()) | ||
| finally: | ||
| if attached_session and session_id: | ||
|
Contributor
Author
There was a problem hiding this comment. Non-blocking robustness note. On an abnormal WebSocket drop (e.g. a crashed production client), this branch detaches ( This matches the pre-existing persistent-HTTP- |
||
| # 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: | ||
| await websocket.close() | ||
|
|
||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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,19 @@ def _next_request_id(self) -> int: | |
| return self._jsonrpc_request_id | ||
|
|
||
| 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.") | ||
| 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="", | ||
| ) | ||
| ) | ||
|
Contributor
Author
There was a problem hiding this comment. Schemeless base URLs break MCP requestsLow Severity
  Reviewed by Cursor Bugbot for commit d7c9a7e. Configure here. |
||
|
|
||
| async def _get_http_client(self) -> Any: | ||
| """Return a shared httpx.AsyncClient, creating one lazily.""" | ||
|
|
@@ -202,19 +212,35 @@ async def _connect_async(self) -> EnvClient: | |
| """ | ||
| Establish connection to the server. | ||
|
|
||
| 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): | ||
| 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 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 | ||
|
|
||
| return await super()._connect_async() | ||
|
|
@@ -365,33 +391,44 @@ async def _close_async(self) -> None: | |
| Close client resources. | ||
|
|
||
| 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. | ||
| """ | ||
| 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: | ||
| 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: | ||
| try: | ||
| if self._http_client is not None: | ||
| try: | ||
| await self._http_client.aclose() | ||
| except Exception: | ||
|
github-code-quality[bot] marked this conversation as resolved.
Fixed
|
||
| # Best effort — do not mask provider/websocket teardown | ||
| 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): | ||
|
|
||
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Owned WebSockets skip exclusive attach tracking
Medium Severity
Plain
/wsconnections create a session but never join_session_websocket_attachments. The new attach guard only rejects IDs already in that set, so a second socket withsession_idcan share a WebSocket-owned environment. When the owner disconnects,_destroy_sessionstill runs and tears that environment down under the attached client.Additional Locations (1)
src/openenv/core/env_server/http_server.py#L1550-L1556Reviewed by Cursor Bugbot for commit d7c9a7e. Configure here.