Upload folder using huggingface_hub

Dockerfile

@@ -17,7 +17,7 @@ RUN apt-get update && apt-get install -y \
 COPY envs/coding_env/ ./envs/coding_env/
 
 # Install openenv-core first from PyPI, then coding_env
-RUN pip install --no-cache-dir "openenv-core[core] @ git+https://github.com/meta-pytorch/OpenEnv.git@main" && \
+RUN pip install --no-cache-dir "openenv-core[core] @ git+https://github.com/meta-pytorch/OpenEnv.git@v0.2.3" && \
     pip install --no-cache-dir ./envs/coding_env/
 
 # Environment variables

README.md

@@ -8,7 +8,7 @@ pinned: false
 app_port: 8000
 base_path: /web
 tags:
-  - openenv-0.2.2
+  - openenv-0.2.3
   - openenv
 ---
 
@@ -17,7 +17,7 @@ tags:
 This Space is built from OpenEnv environment `coding_env`.
 
 - Space URL: `https://huggingface.co/spaces/openenv/coding_env`
-- OpenEnv pinned ref: `0.2.2`
+- OpenEnv pinned ref: `0.2.3`
 - Hub tag: `openenv`
 
 ### Connecting from Code

envs/coding_env/pyproject.toml

@@ -8,7 +8,7 @@ version = "0.1.0"
 description = "Coding Environment for OpenEnv"
 requires-python = ">=3.10"
 dependencies = [
-    "openenv-core[core] @ git+https://github.com/meta-pytorch/OpenEnv.git@main",
+    "openenv-core[core] @ git+https://github.com/meta-pytorch/OpenEnv.git@v0.2.3",
     "fastapi>=0.115.0",
     "pydantic>=2.0.0",
     "uvicorn[standard]>=0.24.0",

envs/coding_env/server/Dockerfile

@@ -17,7 +17,7 @@ RUN apt-get update && apt-get install -y \
 COPY envs/coding_env/ ./envs/coding_env/
 
 # Install openenv-core first from PyPI, then coding_env
-RUN pip install --no-cache-dir "openenv-core[core]>=0.2.1" && \
+RUN pip install --no-cache-dir "openenv-core[core]>=0.2.2" && \
     pip install --no-cache-dir ./envs/coding_env/
 
 # Environment variables

pyproject.toml

@@ -8,7 +8,7 @@ version = "0.1.0"
 description = "Coding Environment for OpenEnv"
 requires-python = ">=3.10"
 dependencies = [
-    "openenv-core[core] @ git+https://github.com/meta-pytorch/OpenEnv.git@main",
+    "openenv-core[core] @ git+https://github.com/meta-pytorch/OpenEnv.git@v0.2.3",
     "fastapi>=0.115.0",
     "pydantic>=2.0.0",
     "uvicorn[standard]>=0.24.0",

server/Dockerfile

@@ -17,7 +17,7 @@ RUN apt-get update && apt-get install -y \
 COPY envs/coding_env/ ./envs/coding_env/
 
 # Install openenv-core first from PyPI, then coding_env
-RUN pip install --no-cache-dir "openenv-core[core]>=0.2.1" && \
+RUN pip install --no-cache-dir "openenv-core[core]>=0.2.2" && \
     pip install --no-cache-dir ./envs/coding_env/
 
 # Environment variables

src/core/env_server/http_server.py

@@ -16,11 +16,15 @@ from __future__ import annotations
 import asyncio
 import inspect
 import json
+import logging
 import os
 import time
 import uuid
 from concurrent.futures import ThreadPoolExecutor
-from typing import Any, Callable, Dict, Optional, Type
+from contextlib import AsyncExitStack
+from typing import Any, AsyncContextManager, Callable, cast, Dict, Optional, Type
+
+_MISSING = object()
 
 from fastapi import (
     Body,
@@ -204,8 +208,9 @@ class HTTPEnvServer:
         self.observation_cls = observation_cls
 
         # Session management for WebSocket connections
-        self._sessions: Dict[str, Environment] = {}
+        self._sessions: Dict[str, Optional[Environment]] = {}
         self._session_executors: Dict[str, ThreadPoolExecutor] = {}
+        self._session_stacks: Dict[str, AsyncExitStack] = {}
         self._session_info: Dict[str, SessionInfo] = {}
         self._session_lock = asyncio.Lock()
 
@@ -213,6 +218,14 @@ class HTTPEnvServer:
         # This is needed for environments using sync libraries (e.g., Playwright)
         self._executor = ThreadPoolExecutor(max_workers=32)
 
+        # Idle session reaper configuration.
+        # Timeout is taken from ConcurrencyConfig.session_timeout;
+        # None means no timeout (default — reaper is a no-op).
+        self._session_idle_timeout_s: Optional[float] = (
+            self._concurrency_config.session_timeout
+        )
+        self._reaper_task: Optional[asyncio.Task[None]] = None
+
     def _validate_concurrency_safety(self) -> None:
         """
         Validate that the environment supports the configured concurrency level.
@@ -321,12 +334,37 @@ class HTTPEnvServer:
             )
             raise EnvironmentFactoryError(factory_name) from e
 
+        # Hold the MCP session open for the lifetime of this session,
+        # matching the WebSocket path's AsyncExitStack pattern.  This
+        # prevents per-request MCP transport teardown/reconnection and
+        # preserves FastMCP session state (ctx.set_state / ctx.get_state)
+        # across HTTP calls within the same OpenEnv session.
+        stack = AsyncExitStack()
+        try:
+            mcp_session_factory = getattr(env, "mcp_session", None)
+            if callable(mcp_session_factory):
+                mcp_session_cm = cast(AsyncContextManager[Any], mcp_session_factory())
+                await stack.enter_async_context(mcp_session_cm)
+        except Exception:
+            # MCP transport failed to start — clean up the reserved slot,
+            # the env, and the executor so they don't leak permanently
+            # against _max_concurrent_envs.
+            await stack.aclose()  # best-effort
+            async with self._session_lock:
+                self._sessions.pop(session_id, None)
+                self._session_executors.pop(session_id, None)
+                self._session_info.pop(session_id, None)
+            await self._cleanup_session_resources(env, executor)
+            raise
+
         async with self._session_lock:
             self._sessions[session_id] = env
+            self._session_stacks[session_id] = stack
+            now = time.time()
             self._session_info[session_id] = SessionInfo(
                 session_id=session_id,
                 created_at=current_time,
-                last_activity_at=current_time,
+                last_activity_at=now,
                 step_count=0,
                 environment_type=type(env).__name__,
             )
@@ -343,8 +381,27 @@ class HTTPEnvServer:
         async with self._session_lock:
             env = self._sessions.pop(session_id, None)
             executor = self._session_executors.pop(session_id, None)
+            stack = self._session_stacks.pop(session_id, None)
             self._session_info.pop(session_id, None)
 
+        await self._cleanup_session_resources(env, executor, stack)
+
+    async def _cleanup_session_resources(
+        self,
+        env: Optional[Environment],
+        executor: Optional[ThreadPoolExecutor],
+        stack: Optional[AsyncExitStack] = None,
+    ) -> None:
+        """Close an environment and shut down its executor (best-effort)."""
+        # Close the MCP session stack first — this gracefully exits the
+        # mcp_session() context (and the underlying FastMCP Client session)
+        # before we tear down the environment references.
+        if stack is not None:
+            try:
+                await stack.aclose()
+            except Exception:
+                pass  # Best effort cleanup
+
         # Run close() in the same executor where the env was created
         # This is required for thread-sensitive libraries like Playwright/greenlet
         if env is not None:
@@ -383,6 +440,51 @@ class HTTPEnvServer:
             if increment_step:
                 self._session_info[session_id].step_count += 1
 
+    async def _reap_idle_sessions(self) -> None:
+        """Background task that periodically destroys sessions idle beyond the timeout."""
+        timeout = self._session_idle_timeout_s
+        if timeout is None:
+            return  # no timeout configured — noop
+        interval = max(timeout / 4, 5.0)  # check frequently enough
+        while True:
+            try:
+                await asyncio.sleep(interval)
+                now = time.time()
+                stale_ids: list[str] = []
+                async with self._session_lock:
+                    for sid, info in self._session_info.items():
+                        if now - info.last_activity_at > timeout:
+                            stale_ids.append(sid)
+                for sid in stale_ids:
+                    # Re-check under lock: activity may have arrived since
+                    # the snapshot was taken, making this session active again.
+                    # Refresh `now` so slow _destroy_session calls don't cause
+                    # subsequent entries to be validated against a stale clock.
+                    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:
+                            continue
+                    await self._destroy_session(sid)
+            except asyncio.CancelledError:
+                break
+            except Exception as exc:
+                logging.getLogger(__name__).warning(
+                    "Idle-session reaper encountered an error (will retry): %s",
+                    exc,
+                )
+
+    def _start_reaper(self) -> None:
+        """Start the idle-session reaper if a timeout is configured."""
+        if self._session_idle_timeout_s is not None and self._reaper_task is None:
+            self._reaper_task = asyncio.create_task(self._reap_idle_sessions())
+
+    def _stop_reaper(self) -> None:
+        """Cancel the reaper background task."""
+        if self._reaper_task is not None:
+            self._reaper_task.cancel()
+            self._reaper_task = None
+
     def get_session_info(self, session_id: str) -> Optional[SessionInfo]:
         """
         Get information about a specific session.
@@ -458,6 +560,20 @@ class HTTPEnvServer:
                     f"Invalid mode: '{mode}'. Must be one of: {valid_modes}"
                 )
 
+        # Wire up idle-session reaper lifecycle via app events
+        server_ref = self
+
+        async def _start_session_reaper() -> None:
+            server_ref._start_reaper()
+
+        async def _stop_session_reaper() -> None:
+            server_ref._stop_reaper()
+
+        if not getattr(app.router, "_openenv_reaper_registered", False):
+            app.router.on_startup.append(_start_session_reaper)
+            app.router.on_shutdown.append(_stop_session_reaper)
+            app.router._openenv_reaper_registered = True  # type: ignore[attr-defined]
+
         # Helper function to handle reset endpoint
         async def reset_handler(
             request: ResetRequest = Body(default_factory=ResetRequest),
@@ -526,53 +642,214 @@ class HTTPEnvServer:
 
         # Helper function to handle MCP endpoint
         async def mcp_handler(
-            request: JsonRpcRequest, session_env: Optional[Environment] = None
+            request: JsonRpcRequest,
+            session_env: Optional[Environment] = None,
+            session_id: Optional[str] = None,
         ) -> JsonRpcResponse:
             """
             Handle MCP JSON-RPC requests.
 
-            Supports tools/list and tools/call methods in JSON-RPC 2.0 format.
+            Supports tools/list and tools/call methods in JSON-RPC 2.0 format,
+            plus OpenEnv session lifecycle methods for HTTP MCP:
+            - openenv/session/create
+            - openenv/session/close
             """
             method = request.method
             request_id = request.id
+            params = request.params
+            if not isinstance(params, dict):
+                return JsonRpcResponse.error_response(
+                    JsonRpcErrorCode.INVALID_PARAMS,
+                    "Params must be an object",
+                    request_id=request_id,
+                )
+
+            # OpenEnv extension methods for explicit MCP session management.
+            # This enables persistent MCP lifecycles over HTTP /mcp, matching WebSocket semantics.
+            if method == "openenv/session/create":
+                if session_env is not None and session_id is not None:
+                    return JsonRpcResponse.success(
+                        result={"session_id": session_id},
+                        request_id=request_id,
+                    )
+                try:
+                    created_session_id, _ = await self._create_session()
+                except SessionCapacityError as e:
+                    return JsonRpcResponse.error_response(
+                        JsonRpcErrorCode.SERVER_ERROR,
+                        str(e),
+                        request_id=request_id,
+                        data={
+                            "active_sessions": e.active_sessions,
+                            "max_sessions": e.max_sessions,
+                        },
+                    )
+                except EnvironmentFactoryError as e:
+                    return JsonRpcResponse.error_response(
+                        JsonRpcErrorCode.SERVER_ERROR,
+                        str(e),
+                        request_id=request_id,
+                        data={"factory_name": e.factory_name},
+                    )
+                return JsonRpcResponse.success(
+                    result={"session_id": created_session_id},
+                    request_id=request_id,
+                )
+
+            if method == "openenv/session/close":
+                target_session_id = params.get("session_id")
+                if not target_session_id:
+                    return JsonRpcResponse.error_response(
+                        JsonRpcErrorCode.INVALID_PARAMS,
+                        "Invalid params - 'session_id' is required",
+                        request_id=request_id,
+                    )
+
+                if session_id is not None and target_session_id == session_id:
+                    return JsonRpcResponse.error_response(
+                        JsonRpcErrorCode.INVALID_REQUEST,
+                        "Cannot close active WebSocket-managed session via MCP method",
+                        request_id=request_id,
+                    )
+
+                async with self._session_lock:
+                    env = self._sessions.pop(target_session_id, _MISSING)
+                    if env is not _MISSING:
+                        executor = self._session_executors.pop(target_session_id, None)
+                        stack = self._session_stacks.pop(target_session_id, None)
+                        self._session_info.pop(target_session_id, None)
+                    else:
+                        executor = None
+                        stack = None
+
+                if env is _MISSING:
+                    return JsonRpcResponse.error_response(
+                        JsonRpcErrorCode.INVALID_PARAMS,
+                        f"Unknown session_id: {target_session_id}",
+                        request_id=request_id,
+                    )
+
+                if env is None:
+                    # Session slot reserved but env factory still running;
+                    # re-insert the placeholder AND the executor so
+                    # _create_session can finish and the executor remains
+                    # tracked for eventual shutdown.
+                    async with self._session_lock:
+                        self._sessions[target_session_id] = None
+                        if executor is not None:
+                            self._session_executors[target_session_id] = executor
+                    return JsonRpcResponse.error_response(
+                        JsonRpcErrorCode.INVALID_REQUEST,
+                        f"Session {target_session_id} is still initializing; retry shortly",
+                        request_id=request_id,
+                    )
+
+                # env/executor/stack cleanup outside the lock
+                await self._cleanup_session_resources(env, executor, stack)
+                return JsonRpcResponse.success(
+                    result={"session_id": target_session_id, "closed": True},
+                    request_id=request_id,
+                )
+
+            requested_session_id = params.get("session_id")
+            managed_session_id = session_id
 
             # Use provided session environment or create temporary one
             if session_env is not None:
                 _env = session_env
                 should_close = False
+            elif requested_session_id:
+                async with self._session_lock:
+                    _env = self._sessions.get(requested_session_id, _MISSING)
+
+                if _env is _MISSING:
+                    return JsonRpcResponse.error_response(
+                        JsonRpcErrorCode.INVALID_PARAMS,
+                        f"Unknown session_id: {requested_session_id}",
+                        request_id=request_id,
+                    )
+
+                if _env is None:
+                    return JsonRpcResponse.error_response(
+                        JsonRpcErrorCode.INVALID_REQUEST,
+                        f"Session {requested_session_id} is still initializing; retry shortly",
+                        request_id=request_id,
+                    )
+
+                should_close = False
+                managed_session_id = requested_session_id
             else:
                 _env = self._env_factory()
                 should_close = True
             try:
+                mcp_client = getattr(_env, "mcp_client", None)
+                mcp_server = getattr(_env, "mcp_server", None)
+                mcp_session_factory = getattr(_env, "mcp_session", None)
+
                 if method == McpMethod.TOOLS_LIST:
                     # Check if environment is MCP-enabled
-                    if not hasattr(_env, "mcp_client"):
+                    if mcp_client is None and mcp_server is None:
                         return JsonRpcResponse.error_response(
                             JsonRpcErrorCode.INTERNAL_ERROR,
                             "Environment does not support MCP",
                             request_id=request_id,
                         )
 
-                    # Use async context manager for MCP client
-                    async with _env.mcp_client:
-                        tools = await _env.mcp_client.list_tools()
+                    if mcp_client:
+                        if managed_session_id and mcp_client.is_connected():
+                            # Session-managed with live transport — call
+                            # directly, no redundant re-entry.
+                            tools = await mcp_client.list_tools()
+                        elif callable(mcp_session_factory):
+                            # Stateless request, or session-managed but the
+                            # background transport was lost: (re-)open.
+                            mcp_session_cm = cast(
+                                AsyncContextManager[Any], mcp_session_factory()
+                            )
+                            async with mcp_session_cm:
+                                tools = await mcp_client.list_tools()
+                        else:
+                            async with mcp_client:
+                                tools = await mcp_client.list_tools()
+
+                        return JsonRpcResponse.success(
+                            result={
+                                "tools": [
+                                    t.model_dump()
+                                    if hasattr(t, "model_dump")
+                                    else dict(t)
+                                    for t in tools
+                                ]
+                            },
+                            request_id=request_id,
+                        )
 
-                    return JsonRpcResponse.success(
-                        result={
-                            "tools": [
-                                t.model_dump() if hasattr(t, "model_dump") else dict(t)
-                                for t in tools
-                            ]
-                        },
+                    if mcp_server:
+                        tools = []
+                        for _tool_name, tool in get_server_tools(mcp_server).items():
+                            tools.append(
+                                {
+                                    "name": tool.name,
+                                    "description": tool.description or "",
+                                    "inputSchema": tool.parameters or {},
+                                }
+                            )
+                        return JsonRpcResponse.success(
+                            result={"tools": tools},
+                            request_id=request_id,
+                        )
+
+                    return JsonRpcResponse.error_response(
+                        JsonRpcErrorCode.INTERNAL_ERROR,
+                        "MCP server not available",
                         request_id=request_id,
                     )
 
                 elif method == McpMethod.TOOLS_CALL:
-                    params = request.params
                     tool_name = params.get("name")
                     arguments = params.get("arguments", {})
 
-                    if not hasattr(_env, "mcp_client"):
+                    if mcp_client is None and mcp_server is None:
                         return JsonRpcResponse.error_response(
                             JsonRpcErrorCode.INTERNAL_ERROR,
                             "Environment does not support MCP",
@@ -581,15 +858,51 @@ class HTTPEnvServer:
 
                     if not tool_name:
                         return JsonRpcResponse.error_response(
-                            JsonRpcErrorCode.INVALID_REQUEST,
+                            JsonRpcErrorCode.INVALID_PARAMS,
                             "Missing 'name' in params",
                             request_id=request_id,
                         )
 
-                    # Use async context manager for MCP client
-                    async with _env.mcp_client:
-                        result = await _env.mcp_client.call_tool(
-                            name=tool_name, arguments=arguments
+                    if mcp_client:
+                        if managed_session_id and mcp_client.is_connected():
+                            # Session-managed with live transport.
+                            result = await mcp_client.call_tool(
+                                name=tool_name, arguments=arguments
+                            )
+                        elif callable(mcp_session_factory):
+                            # Stateless request, or session-managed but the
+                            # background transport was lost: (re-)open.
+                            mcp_session_cm = cast(
+                                AsyncContextManager[Any], mcp_session_factory()
+                            )
+                            async with mcp_session_cm:
+                                result = await mcp_client.call_tool(
+                                    name=tool_name, arguments=arguments
+                                )
+                        else:
+                            async with mcp_client:
+                                result = await mcp_client.call_tool(
+                                    name=tool_name, arguments=arguments
+                                )
+                    elif mcp_server:
+                        server_tools = get_server_tools(mcp_server)
+                        if tool_name in server_tools:
+                            tool = server_tools[tool_name]
+                            if inspect.iscoroutinefunction(tool.fn):
+                                result = await tool.fn(**arguments)
+                            else:
+                                result = tool.fn(**arguments)
+                        else:
+                            return JsonRpcResponse.error_response(
+                                JsonRpcErrorCode.INVALID_PARAMS,
+                                f"Tool not found: {tool_name}",
+                                request_id=request_id,
+                            )
+                    else:
+                        return JsonRpcResponse.error_response(
+                            JsonRpcErrorCode.INTERNAL_ERROR,
+                            "MCP server not available",
+                            request_id=request_id,
                         )
 
                     # Ensure result is JSON serializable
@@ -614,6 +927,11 @@ class HTTPEnvServer:
                     request_id=request_id,
                 )
             finally:
+                if managed_session_id:
+                    self._update_session_activity(
+                        managed_session_id,
+                        increment_step=(method == McpMethod.TOOLS_CALL),
+                    )
                 if should_close:
                     _env.close()
 
@@ -637,42 +955,59 @@ class HTTPEnvServer:
             try:
                 # Create session with dedicated environment
                 session_id, session_env = await self._create_session()
+                if session_env is None:
+                    raise RuntimeError(
+                        "Session environment not initialized for MCP websocket"
+                    )
 
-                while True:
-                    # Receive message from client
-                    raw_message = await websocket.receive_text()
-
-                    try:
-                        jsonrpc_dict = json.loads(raw_message)
-                        jsonrpc_request = JsonRpcRequest(**jsonrpc_dict)
-                    except json.JSONDecodeError as e:
-                        error_resp = JsonRpcResponse.error_response(
-                            JsonRpcErrorCode.PARSE_ERROR,
-                            f"Parse error: {e}",
-                        )
-                        await websocket.send_text(error_resp.model_dump_json())
-                        continue
-                    except ValidationError as e:
-                        error_resp = JsonRpcResponse.error_response(
-                            JsonRpcErrorCode.INVALID_REQUEST,
-                            f"Invalid request: {e}",
-                        )
-                        await websocket.send_text(error_resp.model_dump_json())
-                        continue
+                # If environment has an mcp_session context manager, hold it open
+                # for the lifetime of the websocket connection
 
-                    try:
-                        # Call mcp_handler with session environment
-                        response = await mcp_handler(
-                            jsonrpc_request, session_env=session_env
+                async with AsyncExitStack() as stack:
+                    mcp_session_factory = getattr(session_env, "mcp_session", None)
+                    if callable(mcp_session_factory):
+                        mcp_session_cm = cast(
+                            AsyncContextManager[Any], mcp_session_factory()
                         )
-                        await websocket.send_text(response.model_dump_json())
-                    except Exception as e:
-                        error_resp = JsonRpcResponse.error_response(
-                            JsonRpcErrorCode.INTERNAL_ERROR,
-                            str(e),
-                            request_id=jsonrpc_request.id,
-                        )
-                        await websocket.send_text(error_resp.model_dump_json())
+                        await stack.enter_async_context(mcp_session_cm)
+
+                    while True:
+                        # Receive message from client
+                        raw_message = await websocket.receive_text()
+
+                        try:
+                            jsonrpc_dict = json.loads(raw_message)
+                            jsonrpc_request = JsonRpcRequest(**jsonrpc_dict)
+                        except json.JSONDecodeError as e:
+                            error_resp = JsonRpcResponse.error_response(
+                                JsonRpcErrorCode.PARSE_ERROR,
+                                f"Parse error: {e}",
+                            )
+                            await websocket.send_text(error_resp.model_dump_json())
+                            continue
+                        except ValidationError as e:
+                            error_resp = JsonRpcResponse.error_response(
+                                JsonRpcErrorCode.INVALID_REQUEST,
+                                f"Invalid request: {e}",
+                            )
+                            await websocket.send_text(error_resp.model_dump_json())
+                            continue
+
+                        try:
+                            # Call mcp_handler with session environment
+                            response = await mcp_handler(
+                                jsonrpc_request,
+                                session_env=session_env,
+                                session_id=session_id,
+                            )
+                            await websocket.send_text(response.model_dump_json())
+                        except Exception as e:
+                            error_resp = JsonRpcResponse.error_response(
+                                JsonRpcErrorCode.INTERNAL_ERROR,
+                                str(e),
+                                request_id=jsonrpc_request.id,
+                            )
+                            await websocket.send_text(error_resp.model_dump_json())
 
             except WebSocketDisconnect:
                 pass
@@ -931,120 +1266,8 @@ all schema information needed to interact with the environment.
                     JsonRpcErrorCode.PARSE_ERROR
                 ).model_dump()
 
-            method = request.method
-            params = request.params
-            request_id = request.id
-
-            # Create a temporary environment for MCP access
-            _env = self._env_factory()
-
-            try:
-                # Check if environment supports MCP
-                if not hasattr(_env, "mcp_client") and not hasattr(_env, "mcp_server"):
-                    return JsonRpcResponse.error_response(
-                        JsonRpcErrorCode.INTERNAL_ERROR,
-                        "Environment does not support MCP",
-                        request_id=request_id,
-                    ).model_dump()
-
-                if method == McpMethod.TOOLS_LIST:
-                    # List tools from MCP server
-                    if hasattr(_env, "mcp_client") and _env.mcp_client:
-                        async with _env.mcp_client:
-                            tools = await _env.mcp_client.list_tools()
-                        return JsonRpcResponse.success(
-                            result={
-                                "tools": [
-                                    t.model_dump()
-                                    if hasattr(t, "model_dump")
-                                    else dict(t)
-                                    for t in tools
-                                ]
-                            },
-                            request_id=request_id,
-                        ).model_dump()
-                    elif hasattr(_env, "mcp_server") and _env.mcp_server:
-                        # Use server directly
-                        tools = []
-                        for tool_name, tool in get_server_tools(
-                            _env.mcp_server
-                        ).items():
-                            tool_dict = {
-                                "name": tool.name,
-                                "description": tool.description or "",
-                                "inputSchema": tool.parameters or {},
-                            }
-                            tools.append(tool_dict)
-                        return JsonRpcResponse.success(
-                            result={"tools": tools},
-                            request_id=request_id,
-                        ).model_dump()
-                    else:
-                        return JsonRpcResponse.error_response(
-                            JsonRpcErrorCode.INTERNAL_ERROR,
-                            "MCP server not available",
-                            request_id=request_id,
-                        ).model_dump()
-
-                elif method == McpMethod.TOOLS_CALL:
-                    tool_name = params.get("name")
-                    arguments = params.get("arguments", {})
-
-                    if not tool_name:
-                        return JsonRpcResponse.error_response(
-                            JsonRpcErrorCode.INVALID_PARAMS,
-                            "Invalid params - 'name' is required",
-                            request_id=request_id,
-                        ).model_dump()
-
-                    # Call tool via MCP
-                    if hasattr(_env, "mcp_client") and _env.mcp_client:
-                        async with _env.mcp_client:
-                            result = await _env.mcp_client.call_tool(
-                                name=tool_name, arguments=arguments
-                            )
-                    elif hasattr(_env, "mcp_server") and _env.mcp_server:
-                        # Call tool directly on FastMCP server
-                        server_tools = get_server_tools(_env.mcp_server)
-                        if tool_name in server_tools:
-                            tool = server_tools[tool_name]
-                            result = tool.fn(**arguments)
-                        else:
-                            return JsonRpcResponse.error_response(
-                                JsonRpcErrorCode.INVALID_PARAMS,
-                                f"Tool not found: {tool_name}",
-                                request_id=request_id,
-                            ).model_dump()
-                    else:
-                        return JsonRpcResponse.error_response(
-                            JsonRpcErrorCode.INTERNAL_ERROR,
-                            "MCP server not available",
-                            request_id=request_id,
-                        ).model_dump()
-
-                    # Make result JSON serializable
-                    serializable_result = _make_json_serializable(result)
-
-                    return JsonRpcResponse.success(
-                        result=serializable_result,
-                        request_id=request_id,
-                    ).model_dump()
-
-                else:
-                    return JsonRpcResponse.error_response(
-                        JsonRpcErrorCode.METHOD_NOT_FOUND,
-                        f"Method not found: {method}",
-                        request_id=request_id,
-                    ).model_dump()
-
-            except Exception as e:
-                return JsonRpcResponse.error_response(
-                    JsonRpcErrorCode.INTERNAL_ERROR,
-                    str(e),
-                    request_id=request_id,
-                ).model_dump()
-            finally:
-                _env.close()
+            response = await mcp_handler(request)
+            return response.model_dump()
 
         # Register WebSocket endpoint for persistent sessions
         @app.websocket("/ws")
@@ -1066,135 +1289,167 @@ all schema information needed to interact with the environment.
             try:
                 # Create session with dedicated environment
                 session_id, session_env = await self._create_session()
+                if session_env is None:
+                    raise RuntimeError(
+                        "Session environment not initialized for websocket"
+                    )
 
-                while True:
-                    # Receive message from client
-                    raw_message = await websocket.receive_text()
+                # Keep MCP session open for entire websocket lifetime
+                # (avoids reconnect overhead on every message)
 
-                    try:
-                        message_dict = json.loads(raw_message)
-                    except json.JSONDecodeError as e:
-                        error_resp = WSErrorResponse(
-                            data={
-                                "message": f"Invalid JSON: {e}",
-                                "code": WSErrorCode.INVALID_JSON,
-                            }
+                async with AsyncExitStack() as stack:
+                    mcp_session_factory = getattr(session_env, "mcp_session", None)
+                    if callable(mcp_session_factory):
+                        mcp_session_cm = cast(
+                            AsyncContextManager[Any], mcp_session_factory()
                         )
-                        await websocket.send_text(error_resp.model_dump_json())
-                        continue
-
-                    msg_type = message_dict.get("type", "")
+                        await stack.enter_async_context(mcp_session_cm)
+
+                    while True:
+                        # Receive message from client
+                        raw_message = await websocket.receive_text()
+
+                        try:
+                            message_dict = json.loads(raw_message)
+                        except json.JSONDecodeError as e:
+                            error_resp = WSErrorResponse(
+                                data={
+                                    "message": f"Invalid JSON: {e}",
+                                    "code": WSErrorCode.INVALID_JSON,
+                                }
+                            )
+                            await websocket.send_text(error_resp.model_dump_json())
+                            continue
 
-                    try:
-                        match msg_type:
-                            case "reset":
-                                msg = WSResetMessage(**message_dict)
+                        msg_type = message_dict.get("type", "")
 
-                                is_async = (
-                                    session_env.reset_async.__func__
-                                    is not Environment.reset_async
-                                )
+                        try:
+                            match msg_type:
+                                case "reset":
+                                    msg = WSResetMessage(**message_dict)
 
-                                if is_async:
-                                    sig = inspect.signature(session_env.reset_async)
-                                    valid_kwargs = self._get_valid_kwargs(sig, msg.data)
-                                    observation = await session_env.reset_async(
-                                        **valid_kwargs
+                                    is_async = (
+                                        session_env.reset_async.__func__
+                                        is not Environment.reset_async
                                     )
-                                else:
-                                    sig = inspect.signature(session_env.reset)
-                                    valid_kwargs = self._get_valid_kwargs(sig, msg.data)
-                                    observation = await self._run_in_session_executor(
-                                        session_id, session_env.reset, **valid_kwargs
-                                    )
-
-                                self._update_session_activity(session_id)
-
-                                response = WSObservationResponse(
-                                    data=serialize_observation(observation),
-                                )
 
-                            case "step":
-                                msg = WSStepMessage(**message_dict)
-                                action = deserialize_action(msg.data, self.action_cls)
+                                    if is_async:
+                                        sig = inspect.signature(session_env.reset_async)
+                                        valid_kwargs = self._get_valid_kwargs(
+                                            sig, msg.data
+                                        )
+                                        observation = await session_env.reset_async(
+                                            **valid_kwargs
+                                        )
+                                    else:
+                                        sig = inspect.signature(session_env.reset)
+                                        valid_kwargs = self._get_valid_kwargs(
+                                            sig, msg.data
+                                        )
+                                        observation = (
+                                            await self._run_in_session_executor(
+                                                session_id,
+                                                session_env.reset,
+                                                **valid_kwargs,
+                                            )
+                                        )
+
+                                    self._update_session_activity(session_id)
+
+                                    response = WSObservationResponse(
+                                        data=serialize_observation(observation),
+                                    )
 
-                                is_async = (
-                                    session_env.step_async.__func__
-                                    is not Environment.step_async
-                                )
+                                case "step":
+                                    msg = WSStepMessage(**message_dict)
+                                    action = deserialize_action(
+                                        msg.data, self.action_cls
+                                    )
 
-                                if is_async:
-                                    observation = await session_env.step_async(action)
-                                else:
-                                    observation = await self._run_in_session_executor(
-                                        session_id, session_env.step, action
+                                    is_async = (
+                                        session_env.step_async.__func__
+                                        is not Environment.step_async
                                     )
 
-                                self._update_session_activity(
-                                    session_id, increment_step=True
-                                )
+                                    if is_async:
+                                        observation = await session_env.step_async(
+                                            action
+                                        )
+                                    else:
+                                        observation = (
+                                            await self._run_in_session_executor(
+                                                session_id, session_env.step, action
+                                            )
+                                        )
+
+                                    self._update_session_activity(
+                                        session_id, increment_step=True
+                                    )
 
-                                response = WSObservationResponse(
-                                    data=serialize_observation(observation)
-                                )
+                                    response = WSObservationResponse(
+                                        data=serialize_observation(observation)
+                                    )
 
-                            case "state":
-                                msg = WSStateMessage(**message_dict)
-                                state = session_env.state
-                                if hasattr(state, "model_dump"):
-                                    state_data = state.model_dump()
-                                else:
-                                    state_data = dict(state) if state else {}
-
-                                response = WSStateResponse(data=state_data)
-
-                            case "close":
-                                msg = WSCloseMessage(**message_dict)
-                                break
-
-                            case "mcp":
-                                msg = WSMCPMessage(**message_dict)
-                                try:
-                                    rpc_request = JsonRpcRequest(**msg.data)
-                                except (ValidationError, Exception) as e:
-                                    rpc_response = JsonRpcResponse.error_response(
-                                        JsonRpcErrorCode.INVALID_REQUEST,
-                                        f"Invalid request: {e}",
+                                case "state":
+                                    msg = WSStateMessage(**message_dict)
+                                    state = session_env.state
+                                    if hasattr(state, "model_dump"):
+                                        state_data = state.model_dump()
+                                    else:
+                                        state_data = dict(state) if state else {}
+
+                                    response = WSStateResponse(data=state_data)
+
+                                case "close":
+                                    msg = WSCloseMessage(**message_dict)
+                                    break
+
+                                case "mcp":
+                                    msg = WSMCPMessage(**message_dict)
+                                    try:
+                                        rpc_request = JsonRpcRequest(**msg.data)
+                                    except (ValidationError, Exception) as e:
+                                        rpc_response = JsonRpcResponse.error_response(
+                                            JsonRpcErrorCode.INVALID_REQUEST,
+                                            f"Invalid request: {e}",
+                                        )
+                                    else:
+                                        rpc_response = await mcp_handler(
+                                            rpc_request,
+                                            session_env=session_env,
+                                            session_id=session_id,
+                                        )
+                                    response = WSMCPResponse(
+                                        data=rpc_response.model_dump()
                                     )
-                                else:
-                                    rpc_response = await mcp_handler(
-                                        rpc_request,
-                                        session_env=session_env,
+
+                                case _:
+                                    response = WSErrorResponse(
+                                        data={
+                                            "message": f"Unknown message type: {msg_type}",
+                                            "code": WSErrorCode.UNKNOWN_TYPE,
+                                        }
                                     )
-                                response = WSMCPResponse(data=rpc_response.model_dump())
-
-                            case _:
-                                response = WSErrorResponse(
-                                    data={
-                                        "message": f"Unknown message type: {msg_type}",
-                                        "code": WSErrorCode.UNKNOWN_TYPE,
-                                    }
-                                )
 
-                        await websocket.send_text(response.model_dump_json())
+                            await websocket.send_text(response.model_dump_json())
 
-                    except ValidationError as e:
-                        error_resp = WSErrorResponse(
-                            data={
-                                "message": "Invalid message",
-                                "code": WSErrorCode.VALIDATION_ERROR,
-                                "errors": e.errors(),
-                            }
-                        )
-                        await websocket.send_text(error_resp.model_dump_json())
-                    except Exception as e:
-                        error_resp = WSErrorResponse(
-                            data={
-                                "message": str(e),
-                                "code": WSErrorCode.EXECUTION_ERROR,
-                            }
-                        )
-                        await websocket.send_text(error_resp.model_dump_json())
+                        except ValidationError as e:
+                            error_resp = WSErrorResponse(
+                                data={
+                                    "message": "Invalid message",
+                                    "code": WSErrorCode.VALIDATION_ERROR,
+                                    "errors": e.errors(),
+                                }
+                            )
+                            await websocket.send_text(error_resp.model_dump_json())
+                        except Exception as e:
+                            error_resp = WSErrorResponse(
+                                data={
+                                    "message": str(e),
+                                    "code": WSErrorCode.EXECUTION_ERROR,
+                                }
+                            )
+                            await websocket.send_text(error_resp.model_dump_json())
 
             except WebSocketDisconnect:
                 pass
@@ -1276,7 +1531,7 @@ def create_app(
         from .web_interface import create_web_interface_app
 
         return create_web_interface_app(
-            env,
+            cast(Any, env),
             action_cls,
             observation_cls,
             env_name,

src/core/env_server/mcp_environment.py

@@ -56,6 +56,7 @@ import asyncio
 import inspect
 from abc import abstractmethod
 from collections import defaultdict
+from contextlib import asynccontextmanager
 from typing import Any, Callable, Dict, Optional
 
 from fastmcp import Client
@@ -164,6 +165,52 @@ class MCPEnvironment(Environment):
         # Track tool schemas for list_tools: {tool_name: {mode: schema}}
         self._mode_tool_schemas = defaultdict(dict)
 
+    def _require_mcp_client(self) -> Any:
+        """Return MCP client or raise if environment has been closed."""
+        if self.mcp_client is None:
+            raise RuntimeError("MCP client is not available; environment is closed")
+        return self.mcp_client
+
+    def _require_mcp_server(self) -> Any:
+        """Return MCP server or raise if environment has been closed."""
+        if self.mcp_server is None:
+            raise RuntimeError("MCP server is not available; environment is closed")
+        return self.mcp_server
+
+    @asynccontextmanager
+    async def mcp_session(self):
+        """
+        Context manager for MCP client sessions.
+
+        This wrapper serves two purposes:
+
+        1. **Null guard** — raises a clear error if ``close()`` has already
+           been called (``mcp_client`` is ``None``).
+
+        2. **AsyncExitStack adapter** — FastMCP's ``Client.__aenter__``
+           creates a background ``asyncio.Task`` for session management.
+           When entered directly via ``AsyncExitStack`` in the HTTP session
+           path (``_create_session``), this task can be cancelled by ASGI
+           harnesses (e.g. Starlette ``TestClient``) between requests,
+           corrupting session state.  Wrapping in an ``asynccontextmanager``
+           generator isolates the task lifecycle: the generator frame keeps
+           ``async with client:`` suspended at ``yield``, so cleanup only
+           runs when the stack explicitly closes the generator — not when
+           the event loop cancels orphaned tasks.
+
+        Delegates to FastMCP's ``Client`` context manager which is
+        reentrant: the first entry opens the transport and subsequent
+        (nested) entries simply increment an internal reference counter.
+        The transport is closed only when the outermost context exits.
+
+        No external lock is needed because ``Client._connect`` /
+        ``Client._disconnect`` already serialise connection state changes
+        through their own ``anyio.Lock``.
+        """
+        client = self._require_mcp_client()
+        async with client:
+            yield client
+
     @property
     def supports_code_mode(self) -> bool:
         """Check if this environment supports code mode (execute_code)."""
@@ -292,7 +339,8 @@ class MCPEnvironment(Environment):
 
             # If mode is None, register with FastMCP as usual
             if mode is None:
-                decorated_func = self.mcp_server.tool()(func)
+                mcp_server = self._require_mcp_server()
+                decorated_func = mcp_server.tool()(func)
                 self._mode_tools[tool_name][None] = func
                 return decorated_func
 
@@ -372,24 +420,49 @@ class MCPEnvironment(Environment):
             return self._step_impl(action, timeout_s=timeout_s, **kwargs)
 
     def _handle_list_tools(self) -> ListToolsObservation:
+        """Sync wrapper — delegates to the canonical async implementation."""
+        return run_async_safely(self._async_handle_list_tools())
+
+    async def _async_list_tools(self) -> list:
         """
-        Handle a ListToolsAction by querying the MCP server.
+        Async helper to list tools from the MCP client.
 
         Returns:
-            ListToolsObservation containing all available tools with their
-            names, descriptions, and input schemas, filtered by current mode.
+            List of tool objects from the MCP server.
         """
-        try:
-            # Get current mode
-            current_mode = getattr(self, "_mode", None)
+        async with self.mcp_session() as client:
+            return await client.list_tools()
 
-            # Start with tools from FastMCP server (mode=None tools)
-            tools_result = run_async_safely(self._async_list_tools())
+    def _handle_call_tool(
+        self,
+        action: CallToolAction,
+        timeout_s: Optional[float] = None,
+    ) -> CallToolObservation:
+        """Sync wrapper — delegates to the canonical async implementation."""
+        return run_async_safely(
+            self._async_handle_call_tool(action, timeout_s=timeout_s)
+        )
 
-            # Build list of Tool objects
-            tools = []
+    async def _async_call_tool(self, tool_name: str, arguments: dict) -> Any:
+        """
+        Async helper to call a tool on the MCP server.
 
-            # Add FastMCP tools that are not mode-specific
+        Args:
+            tool_name: Name of the tool to invoke.
+            arguments: Dictionary of arguments to pass to the tool.
+
+        Returns:
+            The result from the tool execution.
+        """
+        async with self.mcp_session() as client:
+            return await client.call_tool(tool_name, arguments)
+
+    async def _async_handle_list_tools(self) -> ListToolsObservation:
+        """Async version of _handle_list_tools — avoids run_async_safely."""
+        try:
+            current_mode = getattr(self, "_mode", None)
+            tools_result = await self._async_list_tools()
+            tools = []
             for tool in tools_result:
                 if tool.name not in self._mode_tool_schemas:
                     tools.append(
@@ -401,11 +474,8 @@ class MCPEnvironment(Environment):
                             else {},
                         )
                     )
-
-            # Add mode-specific tools available in current mode
             for tool_name, mode_schemas in self._mode_tool_schemas.items():
                 if None in mode_schemas:
-                    # Tool available in all modes
                     schema = mode_schemas[None]
                     tools.append(
                         Tool(
@@ -415,7 +485,6 @@ class MCPEnvironment(Environment):
                         )
                     )
                 elif current_mode in mode_schemas:
-                    # Tool available in current mode
                     schema = mode_schemas[current_mode]
                     tools.append(
                         Tool(
@@ -424,65 +493,30 @@ class MCPEnvironment(Environment):
                             input_schema=schema["input_schema"],
                         )
                     )
-
             return ListToolsObservation(tools=tools)
-
         except Exception as e:
-            # Return an observation with error in metadata
             return ListToolsObservation(
                 tools=[],
-                metadata={
-                    "error": str(e),
-                    "error_type": "list_tools_failed",
-                },
+                metadata={"error": str(e), "error_type": "list_tools_failed"},
             )
 
-    async def _async_list_tools(self) -> list:
-        """
-        Async helper to list tools from the MCP client.
-
-        Returns:
-            List of tool objects from the MCP server.
-        """
-        async with self.mcp_client:
-            return await self.mcp_client.list_tools()
-
-    def _handle_call_tool(
+    async def _async_handle_call_tool(
         self,
         action: CallToolAction,
         timeout_s: Optional[float] = None,
     ) -> CallToolObservation:
-        """
-        Handle a CallToolAction by invoking the specified tool.
-
-        Args:
-            action: The CallToolAction containing tool_name and arguments.
-            timeout_s: Timeout in seconds. Defaults to MCP_TOOL_CALL_TIMEOUT (30s).
-
-        Returns:
-            CallToolObservation with the tool's result or an error.
-        """
+        """Async version of _handle_call_tool — avoids run_async_safely."""
         timeout = timeout_s if timeout_s is not None else MCP_TOOL_CALL_TIMEOUT
-
-        # Check if this is a mode-specific tool
         tool_name = action.tool_name
         current_mode = getattr(self, "_mode", None)
 
         if tool_name in self._mode_tools:
             mode_info = self._mode_tools[tool_name]
-
-            # Check if tool is available in current mode
-            # Tool is available if:
-            # 1. It has a None mode (available in all modes), OR
-            # 2. It has an implementation for the current mode
             if None in mode_info:
-                # Use the mode-agnostic version
                 func = mode_info[None]
             elif current_mode in mode_info:
-                # Use the mode-specific version
                 func = mode_info[current_mode]
             else:
-                # Tool not available in current mode
                 return CallToolObservation(
                     tool_name=tool_name,
                     result=None,
@@ -491,16 +525,11 @@ class MCPEnvironment(Environment):
                         message=f"Tool '{tool_name}' not available in {current_mode} mode",
                     ),
                 )
-
-            # Call the mode-specific function directly
             try:
-                # Check if function is async and await if necessary
                 if inspect.iscoroutinefunction(func):
-                    result = run_async_safely(func(**action.arguments))
+                    result = await func(**action.arguments)
                 else:
                     result = func(**action.arguments)
-
-                # Wrap result in CallToolResult format to match FastMCP behavior
                 return CallToolObservation(
                     tool_name=tool_name,
                     result=CallToolResult(
@@ -521,22 +550,12 @@ class MCPEnvironment(Environment):
                     ),
                 )
 
-        # Not a mode-specific tool, use FastMCP
         try:
-            # Run the async call_tool with timeout
-            # Use run_async_safely to handle both sync and async contexts
-            result = run_async_safely(
-                asyncio.wait_for(
-                    self._async_call_tool(action.tool_name, action.arguments),
-                    timeout=timeout,
-                )
-            )
-
-            return CallToolObservation(
-                tool_name=action.tool_name,
-                result=result,
+            result = await asyncio.wait_for(
+                self._async_call_tool(action.tool_name, action.arguments),
+                timeout=timeout,
             )
-
+            return CallToolObservation(tool_name=action.tool_name, result=result)
         except asyncio.TimeoutError:
             return CallToolObservation(
                 tool_name=action.tool_name,
@@ -546,11 +565,8 @@ class MCPEnvironment(Environment):
                     message=f"Tool '{action.tool_name}' timed out after {timeout} seconds",
                 ),
             )
-
         except Exception as e:
             error_message = str(e)
-
-            # Determine error type based on the exception
             if (
                 "not found" in error_message.lower()
                 or "unknown tool" in error_message.lower()
@@ -563,29 +579,34 @@ class MCPEnvironment(Environment):
                 error_type = ToolErrorType.INVALID_ARGS
             else:
                 error_type = ToolErrorType.EXECUTION_ERROR
-
             return CallToolObservation(
                 tool_name=action.tool_name,
                 result=None,
-                error=ToolError(
-                    error_type=error_type,
-                    message=error_message,
-                ),
+                error=ToolError(error_type=error_type, message=error_message),
             )
 
-    async def _async_call_tool(self, tool_name: str, arguments: dict) -> Any:
+    async def step_async(
+        self,
+        action: Action,
+        timeout_s: Optional[float] = None,
+        **kwargs: Any,
+    ) -> Observation:
         """
-        Async helper to call a tool on the MCP server.
+        Async step that routes MCP actions without going through run_async_safely.
 
-        Args:
-            tool_name: Name of the tool to invoke.
-            arguments: Dictionary of arguments to pass to the tool.
-
-        Returns:
-            The result from the tool execution.
+        The WebSocket handler calls this directly on the outer event loop, where
+        the MCP session is already open, avoiding the thread/event-loop deadlock
+        that occurs when the sync step() path is used via run_in_executor.
         """
-        async with self.mcp_client:
-            return await self.mcp_client.call_tool(tool_name, arguments)
+        if isinstance(action, ListToolsAction):
+            return await self._async_handle_list_tools()
+        elif isinstance(action, CallToolAction):
+            return await self._async_handle_call_tool(action, timeout_s=timeout_s)
+        else:
+            loop = asyncio.get_event_loop()
+            return await loop.run_in_executor(
+                None, lambda: self._step_impl(action, timeout_s=timeout_s, **kwargs)
+            )
 
     @abstractmethod
     def _step_impl(

src/core/env_server/serialization.py

@@ -14,14 +14,28 @@ HTTP server and web interface implementations.
 
 from typing import Any, Dict, Type
 
+from .mcp_types import CallToolAction, ListToolsAction
 from .types import Action, Observation
 
+# MCP action types keyed by their "type" discriminator value.
+# These are checked before the environment's own action_cls so that
+# ListToolsAction / CallToolAction payloads are never rejected by an
+# unrelated Pydantic model.
+_MCP_ACTION_TYPES: Dict[str, Type[Action]] = {
+    "list_tools": ListToolsAction,
+    "call_tool": CallToolAction,
+}
+
 
 def deserialize_action(action_data: Dict[str, Any], action_cls: Type[Action]) -> Action:
     """
     Convert JSON dict to Action instance using Pydantic validation.
 
-    This is a basic deserialization that works for most environments.
+    MCP action types (``list_tools``, ``call_tool``) are recognised
+    automatically via the ``"type"`` discriminator field, regardless of
+    the environment's configured ``action_cls``.  All other payloads
+    fall through to ``action_cls.model_validate()``.
+
     For special cases (e.g., tensor fields, custom type conversions),
     use deserialize_action_with_preprocessing().
 
@@ -38,6 +52,17 @@ def deserialize_action(action_data: Dict[str, Any], action_cls: Type[Action]) ->
     Note:
         This uses Pydantic's model_validate() for automatic validation.
     """
+    # Route MCP action types before falling through to the env action_cls.
+    # Only intercept when action_cls is the generic Action base or itself an
+    # MCP type (i.e. the server hosts an MCP environment).  This avoids
+    # silently bypassing env-specific validation for non-MCP environments
+    # that happen to use "call_tool" / "list_tools" as a type discriminator.
+    action_type = action_data.get("type")
+    if action_type in _MCP_ACTION_TYPES:
+        mcp_cls = _MCP_ACTION_TYPES[action_type]
+        if action_cls is Action or action_cls in _MCP_ACTION_TYPES.values():
+            return mcp_cls.model_validate(action_data)
+
     return action_cls.model_validate(action_data)
 
 
@@ -62,6 +87,15 @@ def deserialize_action_with_preprocessing(
     Raises:
         ValidationError: If action_data is invalid for the action class
     """
+    # Route MCP action types before preprocessing (they don't need it).
+    # Same guard as deserialize_action: only intercept when action_cls is
+    # the generic Action base or itself an MCP type.
+    action_type = action_data.get("type")
+    if action_type in _MCP_ACTION_TYPES:
+        mcp_cls = _MCP_ACTION_TYPES[action_type]
+        if action_cls is Action or action_cls in _MCP_ACTION_TYPES.values():
+            return mcp_cls.model_validate(action_data)
+
     processed_data = {}
 
     for key, value in action_data.items():

src/core/env_server/web_interface.py

@@ -15,13 +15,15 @@ option (e.g. openenv push --enable-interface) or ENABLE_WEB_INTERFACE env var.
 from __future__ import annotations
 
 import asyncio
+import inspect
 import json
 from concurrent.futures import ThreadPoolExecutor
 from datetime import datetime
 from typing import Any, Callable, Dict, List, Optional, Type
 
 import gradio as gr
-from fastapi import FastAPI, WebSocket, WebSocketDisconnect
+from fastapi import Body, FastAPI, HTTPException, status, WebSocket, WebSocketDisconnect
+from fastapi.responses import RedirectResponse
 from pydantic import BaseModel, ConfigDict, Field
 
 from .gradio_theme import OPENENV_GRADIO_CSS, OPENENV_GRADIO_THEME
@@ -269,6 +271,28 @@ class WebInterfaceManager:
         # Thread pool for running sync code (e.g., Playwright sync API) in async context
         self._executor = ThreadPoolExecutor(max_workers=1)
 
+    @staticmethod
+    def _get_valid_kwargs(
+        sig: inspect.Signature,
+        kwargs: Dict[str, Any],
+        skip_params: Optional[set[str]] = None,
+    ) -> Dict[str, Any]:
+        """Filter kwargs to only those accepted by the target function."""
+        skip_params = skip_params or set()
+        valid_kwargs: Dict[str, Any] = {}
+        has_var_kwargs = any(
+            param.kind == inspect.Parameter.VAR_KEYWORD
+            for param in sig.parameters.values()
+        )
+
+        for key, value in kwargs.items():
+            if key in skip_params:
+                continue
+            if key in sig.parameters or has_var_kwargs:
+                valid_kwargs[key] = value
+
+        return valid_kwargs
+
     async def _run_sync_in_thread_pool(self, func, *args, **kwargs):
         """Run a synchronous function in the thread pool executor.
 
@@ -317,11 +341,24 @@ class WebInterfaceManager:
         for client in disconnected_clients:
             self.connected_clients.remove(client)
 
-    async def reset_environment(self) -> Dict[str, Any]:
+    async def reset_environment(
+        self, reset_kwargs: Optional[Dict[str, Any]] = None
+    ) -> Dict[str, Any]:
         """Reset the environment and update state."""
-        # Run sync reset in thread pool to avoid blocking event loop
-        # and to support environments using sync libraries (e.g., Playwright)
-        observation: Observation = await self._run_sync_in_thread_pool(self.env.reset)
+        reset_kwargs = reset_kwargs or {}
+
+        is_async = self.env.reset_async.__func__ is not Environment.reset_async
+        sig = inspect.signature(self.env.reset_async if is_async else self.env.reset)
+        valid_kwargs = self._get_valid_kwargs(sig, reset_kwargs)
+
+        if is_async:
+            observation = await self.env.reset_async(**valid_kwargs)
+        else:
+            # Run sync reset in thread pool to avoid blocking event loop
+            # and to support environments using sync libraries (e.g., Playwright)
+            observation = await self._run_sync_in_thread_pool(
+                self.env.reset, **valid_kwargs
+            )
         state: State = self.env.state
 
         # Serialize observation once using shared utility
@@ -428,6 +465,16 @@ def create_web_interface_app(
     web_manager = WebInterfaceManager(env, action_cls, observation_cls, metadata)
 
     # Web API routes first (so they take precedence over Gradio mount at /web)
+    @app.get("/", include_in_schema=False)
+    async def web_root():
+        """Redirect the app root to the Gradio interface."""
+        return RedirectResponse(url="/web/")
+
+    @app.get("/web", include_in_schema=False)
+    async def web_root_no_slash():
+        """Redirect /web to /web/ for mounted Gradio deployments behind proxies."""
+        return RedirectResponse(url="/web/")
+
     @app.get("/web/metadata")
     async def web_metadata():
         """Get environment metadata."""
@@ -449,9 +496,9 @@ def create_web_interface_app(
             await web_manager.disconnect_websocket(websocket)
 
     @app.post("/web/reset")
-    async def web_reset():
+    async def web_reset(request: Optional[Dict[str, Any]] = Body(default=None)):
         """Reset endpoint for web interface."""
-        return await web_manager.reset_environment()
+        return await web_manager.reset_environment(request)
 
     @app.post("/web/step")
     async def web_step(request: Dict[str, Any]):
@@ -475,7 +522,13 @@ def create_web_interface_app(
     @app.get("/web/state")
     async def web_state():
         """State endpoint for web interface."""
-        return web_manager.get_state()
+        try:
+            return web_manager.get_state()
+        except RuntimeError as exc:
+            raise HTTPException(
+                status_code=status.HTTP_409_CONFLICT,
+                detail=str(exc),
+            ) from exc
 
     action_fields = _extract_action_fields(action_cls)
     is_chat_env = _is_chat_env(action_cls)
@@ -505,7 +558,7 @@ def create_web_interface_app(
             )
         gradio_blocks = gr.TabbedInterface(
             [default_blocks, custom_blocks],
-            tab_names=["Playground", "Visualization"],
+            tab_names=["Playground", "Custom"],
             title=get_gradio_display_title(metadata),
         )
     else:

src/core/mcp_client.py

@@ -52,6 +52,7 @@ Example (sync wrapper):
     ...     result = env.call_tool("echo_message", message="Hello!")
 """
 
+import asyncio
 from typing import Any, Dict, List, Optional
 
 from .client_types import StepResult
@@ -118,6 +119,66 @@ class MCPClientBase(EnvClient[Any, Observation, State]):
         )
         self._tools_cache: Optional[List[Tool]] = None
         self.use_production_mode = False
+        self._production_session_id: Optional[str] = None
+        self._production_session_lock = asyncio.Lock()
+        self._jsonrpc_request_id = 0
+        self._http_client: Optional[Any] = None  # lazily-created httpx.AsyncClient
+
+    def _next_request_id(self) -> int:
+        """Generate a monotonically increasing JSON-RPC request id."""
+        self._jsonrpc_request_id += 1
+        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"
+
+    async def _get_http_client(self) -> Any:
+        """Return a shared httpx.AsyncClient, creating one lazily."""
+        if self._http_client is None:
+            import httpx
+
+            self._http_client = httpx.AsyncClient()
+        return self._http_client
+
+    async def _production_mcp_request(
+        self, method: str, params: Optional[Dict[str, Any]] = None
+    ) -> Dict[str, Any]:
+        """Send a JSON-RPC request to HTTP /mcp and return parsed JSON response."""
+        client = await self._get_http_client()
+        response = await client.post(
+            self._production_mcp_url(),
+            json={
+                "jsonrpc": "2.0",
+                "method": method,
+                "params": params or {},
+                "id": self._next_request_id(),
+            },
+            timeout=self._message_timeout,
+        )
+        response.raise_for_status()
+        return response.json()
+
+    async def _ensure_production_session(self) -> str:
+        """Create and cache a persistent HTTP MCP session id if needed."""
+        async with self._production_session_lock:
+            if self._production_session_id is not None:
+                return self._production_session_id
+
+            data = await self._production_mcp_request("openenv/session/create")
+            if "error" in data:
+                message = data.get("error", {}).get("message", "unknown error")
+                raise RuntimeError(f"Failed to create MCP session: {message}")
+
+            session_id = data.get("result", {}).get("session_id")
+            if not session_id:
+                raise RuntimeError("Failed to create MCP session: missing session_id")
+
+            self._production_session_id = session_id
+            return session_id
 
     async def list_tools(self, use_cache: bool = True) -> List[Tool]:
         """
@@ -138,26 +199,18 @@ class MCPClientBase(EnvClient[Any, Observation, State]):
         if use_cache and self._tools_cache is not None:
             return self._tools_cache
 
-        # Use production mode HTTP endpoint if enabled
-        if self.use_production_mode:
-            import requests
-
-            # Convert ws:// URL to http:// URL
-            url = self._ws_url.replace("ws://", "http://").replace("wss://", "https://")
-            # Remove /ws suffix if present and add /mcp
-            url = url.rstrip("/ws").rstrip("/") + "/mcp"
-
+        # Use production mode HTTP endpoint if enabled.
+        # Some tests instantiate with __new__ and skip __init__, so default missing flag to False.
+        if getattr(self, "use_production_mode", False):
             try:
-                response = requests.post(
-                    url,
-                    json={
-                        "jsonrpc": "2.0",
-                        "method": "tools/list",
-                        "params": {},
-                        "id": 1,
-                    },
+                session_id = await self._ensure_production_session()
+                data = await self._production_mcp_request(
+                    "tools/list",
+                    {"session_id": session_id},
                 )
-                data = response.json()
+                if "error" in data:
+                    message = data.get("error", {}).get("message", "unknown error")
+                    raise RuntimeError(f"list_tools failed: {message}")
                 if "result" in data and "tools" in data["result"]:
                     tools = [
                         Tool(
@@ -177,7 +230,12 @@ class MCPClientBase(EnvClient[Any, Observation, State]):
             return []
 
         result = await self.step(ListToolsAction())
-        self._tools_cache = result.observation.tools
+        if isinstance(result.observation, ListToolsObservation):
+            self._tools_cache = result.observation.tools
+            return self._tools_cache
+
+        # Unexpected observation type; keep API stable with an empty tool list.
+        self._tools_cache = []
         return self._tools_cache
 
     def _step_payload(self, action: Any) -> Dict[str, Any]:
@@ -251,6 +309,35 @@ class MCPClientBase(EnvClient[Any, Observation, State]):
             step_count=payload.get("step_count", 0),
         )
 
+    async def close(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.
+        """
+        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
+
+        if self._http_client is not None:
+            try:
+                await self._http_client.aclose()
+            except Exception:
+                pass
+            finally:
+                self._http_client = None
+
+        await super().close()
+
 
 class MCPToolClient(MCPClientBase):
     """
@@ -316,6 +403,26 @@ class MCPToolClient(MCPClientBase):
             >>> result = await env.call_tool("greet", name="Claude")
             >>> print(result)  # "Hello, Claude!"
         """
+        if getattr(self, "use_production_mode", False):
+            session_id = await self._ensure_production_session()
+            data = await self._production_mcp_request(
+                "tools/call",
+                {
+                    "name": name,
+                    "arguments": kwargs,
+                    "session_id": session_id,
+                },
+            )
+
+            if "error" in data:
+                message = data.get("error", {}).get("message", "unknown error")
+                raise RuntimeError(f"Tool '{name}' failed: {message}")
+
+            result = data.get("result")
+            if isinstance(result, dict) and "data" in result:
+                return result["data"]
+            return result
+
         action = CallToolAction(tool_name=name, arguments=kwargs)
         result = await self.step(action)
         obs = result.observation

src/core/openenv/__init__.py

@@ -14,10 +14,18 @@ __all__ = [
     "SyncEnvClient",
 ]
 
-try:
-    __version__ = metadata.version("openenv")  # type: ignore[arg-type]
-except metadata.PackageNotFoundError:  # pragma: no cover - local dev
-    __version__ = "0.0.0"
+
+def _load_package_version() -> str:
+    """Resolve the installed distribution version for the OpenEnv package."""
+    for distribution_name in ("openenv-core", "openenv"):
+        try:
+            return metadata.version(distribution_name)
+        except metadata.PackageNotFoundError:
+            continue
+    return "0.0.0"
+
+
+__version__ = _load_package_version()
 
 
 _LAZY_MODULES = {

src/core/openenv/cli/templates/openenv_env/pyproject.toml

@@ -17,7 +17,7 @@ dependencies = [
     # Core OpenEnv runtime (provides FastAPI server + HTTP client types)
     # install from github
     # "openenv-core[core] @ git+https://github.com/meta-pytorch/OpenEnv.git",
-    "openenv-core[core]>=0.2.1",
+    "openenv-core[core]>=0.2.2",
     # Environment-specific dependencies
     # Add all dependencies needed for your environment here
     # Examples:

src/core/openenv/core/env_server/http_server.py

@@ -16,11 +16,15 @@ from __future__ import annotations
 import asyncio
 import inspect
 import json
+import logging
 import os
 import time
 import uuid
 from concurrent.futures import ThreadPoolExecutor
-from typing import Any, Callable, Dict, Optional, Type
+from contextlib import AsyncExitStack
+from typing import Any, AsyncContextManager, Callable, cast, Dict, Optional, Type
+
+_MISSING = object()
 
 from fastapi import (
     Body,
@@ -204,8 +208,9 @@ class HTTPEnvServer:
         self.observation_cls = observation_cls
 
         # Session management for WebSocket connections
-        self._sessions: Dict[str, Environment] = {}
+        self._sessions: Dict[str, Optional[Environment]] = {}
         self._session_executors: Dict[str, ThreadPoolExecutor] = {}
+        self._session_stacks: Dict[str, AsyncExitStack] = {}
         self._session_info: Dict[str, SessionInfo] = {}
         self._session_lock = asyncio.Lock()
 
@@ -213,6 +218,14 @@ class HTTPEnvServer:
         # This is needed for environments using sync libraries (e.g., Playwright)
         self._executor = ThreadPoolExecutor(max_workers=32)
 
+        # Idle session reaper configuration.
+        # Timeout is taken from ConcurrencyConfig.session_timeout;
+        # None means no timeout (default — reaper is a no-op).
+        self._session_idle_timeout_s: Optional[float] = (
+            self._concurrency_config.session_timeout
+        )
+        self._reaper_task: Optional[asyncio.Task[None]] = None
+
     def _validate_concurrency_safety(self) -> None:
         """
         Validate that the environment supports the configured concurrency level.
@@ -321,12 +334,37 @@ class HTTPEnvServer:
             )
             raise EnvironmentFactoryError(factory_name) from e
 
+        # Hold the MCP session open for the lifetime of this session,
+        # matching the WebSocket path's AsyncExitStack pattern.  This
+        # prevents per-request MCP transport teardown/reconnection and
+        # preserves FastMCP session state (ctx.set_state / ctx.get_state)
+        # across HTTP calls within the same OpenEnv session.
+        stack = AsyncExitStack()
+        try:
+            mcp_session_factory = getattr(env, "mcp_session", None)
+            if callable(mcp_session_factory):
+                mcp_session_cm = cast(AsyncContextManager[Any], mcp_session_factory())
+                await stack.enter_async_context(mcp_session_cm)
+        except Exception:
+            # MCP transport failed to start — clean up the reserved slot,
+            # the env, and the executor so they don't leak permanently
+            # against _max_concurrent_envs.
+            await stack.aclose()  # best-effort
+            async with self._session_lock:
+                self._sessions.pop(session_id, None)
+                self._session_executors.pop(session_id, None)
+                self._session_info.pop(session_id, None)
+            await self._cleanup_session_resources(env, executor)
+            raise
+
         async with self._session_lock:
             self._sessions[session_id] = env
+            self._session_stacks[session_id] = stack
+            now = time.time()
             self._session_info[session_id] = SessionInfo(
                 session_id=session_id,
                 created_at=current_time,
-                last_activity_at=current_time,
+                last_activity_at=now,
                 step_count=0,
                 environment_type=type(env).__name__,
             )
@@ -343,8 +381,27 @@ class HTTPEnvServer:
         async with self._session_lock:
             env = self._sessions.pop(session_id, None)
             executor = self._session_executors.pop(session_id, None)
+            stack = self._session_stacks.pop(session_id, None)
             self._session_info.pop(session_id, None)
 
+        await self._cleanup_session_resources(env, executor, stack)
+
+    async def _cleanup_session_resources(
+        self,
+        env: Optional[Environment],
+        executor: Optional[ThreadPoolExecutor],
+        stack: Optional[AsyncExitStack] = None,
+    ) -> None:
+        """Close an environment and shut down its executor (best-effort)."""
+        # Close the MCP session stack first — this gracefully exits the
+        # mcp_session() context (and the underlying FastMCP Client session)
+        # before we tear down the environment references.
+        if stack is not None:
+            try:
+                await stack.aclose()
+            except Exception:
+                pass  # Best effort cleanup
+
         # Run close() in the same executor where the env was created
         # This is required for thread-sensitive libraries like Playwright/greenlet
         if env is not None:
@@ -383,6 +440,51 @@ class HTTPEnvServer:
             if increment_step:
                 self._session_info[session_id].step_count += 1
 
+    async def _reap_idle_sessions(self) -> None:
+        """Background task that periodically destroys sessions idle beyond the timeout."""
+        timeout = self._session_idle_timeout_s
+        if timeout is None:
+            return  # no timeout configured — noop
+        interval = max(timeout / 4, 5.0)  # check frequently enough
+        while True:
+            try:
+                await asyncio.sleep(interval)
+                now = time.time()
+                stale_ids: list[str] = []
+                async with self._session_lock:
+                    for sid, info in self._session_info.items():
+                        if now - info.last_activity_at > timeout:
+                            stale_ids.append(sid)
+                for sid in stale_ids:
+                    # Re-check under lock: activity may have arrived since
+                    # the snapshot was taken, making this session active again.
+                    # Refresh `now` so slow _destroy_session calls don't cause
+                    # subsequent entries to be validated against a stale clock.
+                    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:
+                            continue
+                    await self._destroy_session(sid)
+            except asyncio.CancelledError:
+                break
+            except Exception as exc:
+                logging.getLogger(__name__).warning(
+                    "Idle-session reaper encountered an error (will retry): %s",
+                    exc,
+                )
+
+    def _start_reaper(self) -> None:
+        """Start the idle-session reaper if a timeout is configured."""
+        if self._session_idle_timeout_s is not None and self._reaper_task is None:
+            self._reaper_task = asyncio.create_task(self._reap_idle_sessions())
+
+    def _stop_reaper(self) -> None:
+        """Cancel the reaper background task."""
+        if self._reaper_task is not None:
+            self._reaper_task.cancel()
+            self._reaper_task = None
+
     def get_session_info(self, session_id: str) -> Optional[SessionInfo]:
         """
         Get information about a specific session.
@@ -458,6 +560,20 @@ class HTTPEnvServer:
                     f"Invalid mode: '{mode}'. Must be one of: {valid_modes}"
                 )
 
+        # Wire up idle-session reaper lifecycle via app events
+        server_ref = self
+
+        async def _start_session_reaper() -> None:
+            server_ref._start_reaper()
+
+        async def _stop_session_reaper() -> None:
+            server_ref._stop_reaper()
+
+        if not getattr(app.router, "_openenv_reaper_registered", False):
+            app.router.on_startup.append(_start_session_reaper)
+            app.router.on_shutdown.append(_stop_session_reaper)
+            app.router._openenv_reaper_registered = True  # type: ignore[attr-defined]
+
         # Helper function to handle reset endpoint
         async def reset_handler(
             request: ResetRequest = Body(default_factory=ResetRequest),
@@ -526,53 +642,214 @@ class HTTPEnvServer:
 
         # Helper function to handle MCP endpoint
         async def mcp_handler(
-            request: JsonRpcRequest, session_env: Optional[Environment] = None
+            request: JsonRpcRequest,
+            session_env: Optional[Environment] = None,
+            session_id: Optional[str] = None,
         ) -> JsonRpcResponse:
             """
             Handle MCP JSON-RPC requests.
 
-            Supports tools/list and tools/call methods in JSON-RPC 2.0 format.
+            Supports tools/list and tools/call methods in JSON-RPC 2.0 format,
+            plus OpenEnv session lifecycle methods for HTTP MCP:
+            - openenv/session/create
+            - openenv/session/close
             """
             method = request.method
             request_id = request.id
+            params = request.params
+            if not isinstance(params, dict):
+                return JsonRpcResponse.error_response(
+                    JsonRpcErrorCode.INVALID_PARAMS,
+                    "Params must be an object",
+                    request_id=request_id,
+                )
+
+            # OpenEnv extension methods for explicit MCP session management.
+            # This enables persistent MCP lifecycles over HTTP /mcp, matching WebSocket semantics.
+            if method == "openenv/session/create":
+                if session_env is not None and session_id is not None:
+                    return JsonRpcResponse.success(
+                        result={"session_id": session_id},
+                        request_id=request_id,
+                    )
+                try:
+                    created_session_id, _ = await self._create_session()
+                except SessionCapacityError as e:
+                    return JsonRpcResponse.error_response(
+                        JsonRpcErrorCode.SERVER_ERROR,
+                        str(e),
+                        request_id=request_id,
+                        data={
+                            "active_sessions": e.active_sessions,
+                            "max_sessions": e.max_sessions,
+                        },
+                    )
+                except EnvironmentFactoryError as e:
+                    return JsonRpcResponse.error_response(
+                        JsonRpcErrorCode.SERVER_ERROR,
+                        str(e),
+                        request_id=request_id,
+                        data={"factory_name": e.factory_name},
+                    )
+                return JsonRpcResponse.success(
+                    result={"session_id": created_session_id},
+                    request_id=request_id,
+                )
+
+            if method == "openenv/session/close":
+                target_session_id = params.get("session_id")
+                if not target_session_id:
+                    return JsonRpcResponse.error_response(
+                        JsonRpcErrorCode.INVALID_PARAMS,
+                        "Invalid params - 'session_id' is required",
+                        request_id=request_id,
+                    )
+
+                if session_id is not None and target_session_id == session_id:
+                    return JsonRpcResponse.error_response(
+                        JsonRpcErrorCode.INVALID_REQUEST,
+                        "Cannot close active WebSocket-managed session via MCP method",
+                        request_id=request_id,
+                    )
+
+                async with self._session_lock:
+                    env = self._sessions.pop(target_session_id, _MISSING)
+                    if env is not _MISSING:
+                        executor = self._session_executors.pop(target_session_id, None)
+                        stack = self._session_stacks.pop(target_session_id, None)
+                        self._session_info.pop(target_session_id, None)
+                    else:
+                        executor = None
+                        stack = None
+
+                if env is _MISSING:
+                    return JsonRpcResponse.error_response(
+                        JsonRpcErrorCode.INVALID_PARAMS,
+                        f"Unknown session_id: {target_session_id}",
+                        request_id=request_id,
+                    )
+
+                if env is None:
+                    # Session slot reserved but env factory still running;
+                    # re-insert the placeholder AND the executor so
+                    # _create_session can finish and the executor remains
+                    # tracked for eventual shutdown.
+                    async with self._session_lock:
+                        self._sessions[target_session_id] = None
+                        if executor is not None:
+                            self._session_executors[target_session_id] = executor
+                    return JsonRpcResponse.error_response(
+                        JsonRpcErrorCode.INVALID_REQUEST,
+                        f"Session {target_session_id} is still initializing; retry shortly",
+                        request_id=request_id,
+                    )
+
+                # env/executor/stack cleanup outside the lock
+                await self._cleanup_session_resources(env, executor, stack)
+                return JsonRpcResponse.success(
+                    result={"session_id": target_session_id, "closed": True},
+                    request_id=request_id,
+                )
+
+            requested_session_id = params.get("session_id")
+            managed_session_id = session_id
 
             # Use provided session environment or create temporary one
             if session_env is not None:
                 _env = session_env
                 should_close = False
+            elif requested_session_id:
+                async with self._session_lock:
+                    _env = self._sessions.get(requested_session_id, _MISSING)
+
+                if _env is _MISSING:
+                    return JsonRpcResponse.error_response(
+                        JsonRpcErrorCode.INVALID_PARAMS,
+                        f"Unknown session_id: {requested_session_id}",
+                        request_id=request_id,
+                    )
+
+                if _env is None:
+                    return JsonRpcResponse.error_response(
+                        JsonRpcErrorCode.INVALID_REQUEST,
+                        f"Session {requested_session_id} is still initializing; retry shortly",
+                        request_id=request_id,
+                    )
+
+                should_close = False
+                managed_session_id = requested_session_id
             else:
                 _env = self._env_factory()
                 should_close = True
             try:
+                mcp_client = getattr(_env, "mcp_client", None)
+                mcp_server = getattr(_env, "mcp_server", None)
+                mcp_session_factory = getattr(_env, "mcp_session", None)
+
                 if method == McpMethod.TOOLS_LIST:
                     # Check if environment is MCP-enabled
-                    if not hasattr(_env, "mcp_client"):
+                    if mcp_client is None and mcp_server is None:
                         return JsonRpcResponse.error_response(
                             JsonRpcErrorCode.INTERNAL_ERROR,
                             "Environment does not support MCP",
                             request_id=request_id,
                         )
 
-                    # Use async context manager for MCP client
-                    async with _env.mcp_client:
-                        tools = await _env.mcp_client.list_tools()
+                    if mcp_client:
+                        if managed_session_id and mcp_client.is_connected():
+                            # Session-managed with live transport — call
+                            # directly, no redundant re-entry.
+                            tools = await mcp_client.list_tools()
+                        elif callable(mcp_session_factory):
+                            # Stateless request, or session-managed but the
+                            # background transport was lost: (re-)open.
+                            mcp_session_cm = cast(
+                                AsyncContextManager[Any], mcp_session_factory()
+                            )
+                            async with mcp_session_cm:
+                                tools = await mcp_client.list_tools()
+                        else:
+                            async with mcp_client:
+                                tools = await mcp_client.list_tools()
+
+                        return JsonRpcResponse.success(
+                            result={
+                                "tools": [
+                                    t.model_dump()
+                                    if hasattr(t, "model_dump")
+                                    else dict(t)
+                                    for t in tools
+                                ]
+                            },
+                            request_id=request_id,
+                        )
 
-                    return JsonRpcResponse.success(
-                        result={
-                            "tools": [
-                                t.model_dump() if hasattr(t, "model_dump") else dict(t)
-                                for t in tools
-                            ]
-                        },
+                    if mcp_server:
+                        tools = []
+                        for _tool_name, tool in get_server_tools(mcp_server).items():
+                            tools.append(
+                                {
+                                    "name": tool.name,
+                                    "description": tool.description or "",
+                                    "inputSchema": tool.parameters or {},
+                                }
+                            )
+                        return JsonRpcResponse.success(
+                            result={"tools": tools},
+                            request_id=request_id,
+                        )
+
+                    return JsonRpcResponse.error_response(
+                        JsonRpcErrorCode.INTERNAL_ERROR,
+                        "MCP server not available",
                         request_id=request_id,
                     )
 
                 elif method == McpMethod.TOOLS_CALL:
-                    params = request.params
                     tool_name = params.get("name")
                     arguments = params.get("arguments", {})
 
-                    if not hasattr(_env, "mcp_client"):
+                    if mcp_client is None and mcp_server is None:
                         return JsonRpcResponse.error_response(
                             JsonRpcErrorCode.INTERNAL_ERROR,
                             "Environment does not support MCP",
@@ -581,15 +858,51 @@ class HTTPEnvServer:
 
                     if not tool_name:
                         return JsonRpcResponse.error_response(
-                            JsonRpcErrorCode.INVALID_REQUEST,
+                            JsonRpcErrorCode.INVALID_PARAMS,
                             "Missing 'name' in params",
                             request_id=request_id,
                         )
 
-                    # Use async context manager for MCP client
-                    async with _env.mcp_client:
-                        result = await _env.mcp_client.call_tool(
-                            name=tool_name, arguments=arguments
+                    if mcp_client:
+                        if managed_session_id and mcp_client.is_connected():
+                            # Session-managed with live transport.
+                            result = await mcp_client.call_tool(
+                                name=tool_name, arguments=arguments
+                            )
+                        elif callable(mcp_session_factory):
+                            # Stateless request, or session-managed but the
+                            # background transport was lost: (re-)open.
+                            mcp_session_cm = cast(
+                                AsyncContextManager[Any], mcp_session_factory()
+                            )
+                            async with mcp_session_cm:
+                                result = await mcp_client.call_tool(
+                                    name=tool_name, arguments=arguments
+                                )
+                        else:
+                            async with mcp_client:
+                                result = await mcp_client.call_tool(
+                                    name=tool_name, arguments=arguments
+                                )
+                    elif mcp_server:
+                        server_tools = get_server_tools(mcp_server)
+                        if tool_name in server_tools:
+                            tool = server_tools[tool_name]
+                            if inspect.iscoroutinefunction(tool.fn):
+                                result = await tool.fn(**arguments)
+                            else:
+                                result = tool.fn(**arguments)
+                        else:
+                            return JsonRpcResponse.error_response(
+                                JsonRpcErrorCode.INVALID_PARAMS,
+                                f"Tool not found: {tool_name}",
+                                request_id=request_id,
+                            )
+                    else:
+                        return JsonRpcResponse.error_response(
+                            JsonRpcErrorCode.INTERNAL_ERROR,
+                            "MCP server not available",
+                            request_id=request_id,
                         )
 
                     # Ensure result is JSON serializable
@@ -614,6 +927,11 @@ class HTTPEnvServer:
                     request_id=request_id,
                 )
             finally:
+                if managed_session_id:
+                    self._update_session_activity(
+                        managed_session_id,
+                        increment_step=(method == McpMethod.TOOLS_CALL),
+                    )
                 if should_close:
                     _env.close()
 
@@ -637,42 +955,59 @@ class HTTPEnvServer:
             try:
                 # Create session with dedicated environment
                 session_id, session_env = await self._create_session()
+                if session_env is None:
+                    raise RuntimeError(
+                        "Session environment not initialized for MCP websocket"
+                    )
 
-                while True:
-                    # Receive message from client
-                    raw_message = await websocket.receive_text()
-
-                    try:
-                        jsonrpc_dict = json.loads(raw_message)
-                        jsonrpc_request = JsonRpcRequest(**jsonrpc_dict)
-                    except json.JSONDecodeError as e:
-                        error_resp = JsonRpcResponse.error_response(
-                            JsonRpcErrorCode.PARSE_ERROR,
-                            f"Parse error: {e}",
-                        )
-                        await websocket.send_text(error_resp.model_dump_json())
-                        continue
-                    except ValidationError as e:
-                        error_resp = JsonRpcResponse.error_response(
-                            JsonRpcErrorCode.INVALID_REQUEST,
-                            f"Invalid request: {e}",
-                        )
-                        await websocket.send_text(error_resp.model_dump_json())
-                        continue
+                # If environment has an mcp_session context manager, hold it open
+                # for the lifetime of the websocket connection
 
-                    try:
-                        # Call mcp_handler with session environment
-                        response = await mcp_handler(
-                            jsonrpc_request, session_env=session_env
+                async with AsyncExitStack() as stack:
+                    mcp_session_factory = getattr(session_env, "mcp_session", None)
+                    if callable(mcp_session_factory):
+                        mcp_session_cm = cast(
+                            AsyncContextManager[Any], mcp_session_factory()
                         )
-                        await websocket.send_text(response.model_dump_json())
-                    except Exception as e:
-                        error_resp = JsonRpcResponse.error_response(
-                            JsonRpcErrorCode.INTERNAL_ERROR,
-                            str(e),
-                            request_id=jsonrpc_request.id,
-                        )
-                        await websocket.send_text(error_resp.model_dump_json())
+                        await stack.enter_async_context(mcp_session_cm)
+
+                    while True:
+                        # Receive message from client
+                        raw_message = await websocket.receive_text()
+
+                        try:
+                            jsonrpc_dict = json.loads(raw_message)
+                            jsonrpc_request = JsonRpcRequest(**jsonrpc_dict)
+                        except json.JSONDecodeError as e:
+                            error_resp = JsonRpcResponse.error_response(
+                                JsonRpcErrorCode.PARSE_ERROR,
+                                f"Parse error: {e}",
+                            )
+                            await websocket.send_text(error_resp.model_dump_json())
+                            continue
+                        except ValidationError as e:
+                            error_resp = JsonRpcResponse.error_response(
+                                JsonRpcErrorCode.INVALID_REQUEST,
+                                f"Invalid request: {e}",
+                            )
+                            await websocket.send_text(error_resp.model_dump_json())
+                            continue
+
+                        try:
+                            # Call mcp_handler with session environment
+                            response = await mcp_handler(
+                                jsonrpc_request,
+                                session_env=session_env,
+                                session_id=session_id,
+                            )
+                            await websocket.send_text(response.model_dump_json())
+                        except Exception as e:
+                            error_resp = JsonRpcResponse.error_response(
+                                JsonRpcErrorCode.INTERNAL_ERROR,
+                                str(e),
+                                request_id=jsonrpc_request.id,
+                            )
+                            await websocket.send_text(error_resp.model_dump_json())
 
             except WebSocketDisconnect:
                 pass
@@ -931,120 +1266,8 @@ all schema information needed to interact with the environment.
                     JsonRpcErrorCode.PARSE_ERROR
                 ).model_dump()
 
-            method = request.method
-            params = request.params
-            request_id = request.id
-
-            # Create a temporary environment for MCP access
-            _env = self._env_factory()
-
-            try:
-                # Check if environment supports MCP
-                if not hasattr(_env, "mcp_client") and not hasattr(_env, "mcp_server"):
-                    return JsonRpcResponse.error_response(
-                        JsonRpcErrorCode.INTERNAL_ERROR,
-                        "Environment does not support MCP",
-                        request_id=request_id,
-                    ).model_dump()
-
-                if method == McpMethod.TOOLS_LIST:
-                    # List tools from MCP server
-                    if hasattr(_env, "mcp_client") and _env.mcp_client:
-                        async with _env.mcp_client:
-                            tools = await _env.mcp_client.list_tools()
-                        return JsonRpcResponse.success(
-                            result={
-                                "tools": [
-                                    t.model_dump()
-                                    if hasattr(t, "model_dump")
-                                    else dict(t)
-                                    for t in tools
-                                ]
-                            },
-                            request_id=request_id,
-                        ).model_dump()
-                    elif hasattr(_env, "mcp_server") and _env.mcp_server:
-                        # Use server directly
-                        tools = []
-                        for tool_name, tool in get_server_tools(
-                            _env.mcp_server
-                        ).items():
-                            tool_dict = {
-                                "name": tool.name,
-                                "description": tool.description or "",
-                                "inputSchema": tool.parameters or {},
-                            }
-                            tools.append(tool_dict)
-                        return JsonRpcResponse.success(
-                            result={"tools": tools},
-                            request_id=request_id,
-                        ).model_dump()
-                    else:
-                        return JsonRpcResponse.error_response(
-                            JsonRpcErrorCode.INTERNAL_ERROR,
-                            "MCP server not available",
-                            request_id=request_id,
-                        ).model_dump()
-
-                elif method == McpMethod.TOOLS_CALL:
-                    tool_name = params.get("name")
-                    arguments = params.get("arguments", {})
-
-                    if not tool_name:
-                        return JsonRpcResponse.error_response(
-                            JsonRpcErrorCode.INVALID_PARAMS,
-                            "Invalid params - 'name' is required",
-                            request_id=request_id,
-                        ).model_dump()
-
-                    # Call tool via MCP
-                    if hasattr(_env, "mcp_client") and _env.mcp_client:
-                        async with _env.mcp_client:
-                            result = await _env.mcp_client.call_tool(
-                                name=tool_name, arguments=arguments
-                            )
-                    elif hasattr(_env, "mcp_server") and _env.mcp_server:
-                        # Call tool directly on FastMCP server
-                        server_tools = get_server_tools(_env.mcp_server)
-                        if tool_name in server_tools:
-                            tool = server_tools[tool_name]
-                            result = tool.fn(**arguments)
-                        else:
-                            return JsonRpcResponse.error_response(
-                                JsonRpcErrorCode.INVALID_PARAMS,
-                                f"Tool not found: {tool_name}",
-                                request_id=request_id,
-                            ).model_dump()
-                    else:
-                        return JsonRpcResponse.error_response(
-                            JsonRpcErrorCode.INTERNAL_ERROR,
-                            "MCP server not available",
-                            request_id=request_id,
-                        ).model_dump()
-
-                    # Make result JSON serializable
-                    serializable_result = _make_json_serializable(result)
-
-                    return JsonRpcResponse.success(
-                        result=serializable_result,
-                        request_id=request_id,
-                    ).model_dump()
-
-                else:
-                    return JsonRpcResponse.error_response(
-                        JsonRpcErrorCode.METHOD_NOT_FOUND,
-                        f"Method not found: {method}",
-                        request_id=request_id,
-                    ).model_dump()
-
-            except Exception as e:
-                return JsonRpcResponse.error_response(
-                    JsonRpcErrorCode.INTERNAL_ERROR,
-                    str(e),
-                    request_id=request_id,
-                ).model_dump()
-            finally:
-                _env.close()
+            response = await mcp_handler(request)
+            return response.model_dump()
 
         # Register WebSocket endpoint for persistent sessions
         @app.websocket("/ws")
@@ -1066,135 +1289,167 @@ all schema information needed to interact with the environment.
             try:
                 # Create session with dedicated environment
                 session_id, session_env = await self._create_session()
+                if session_env is None:
+                    raise RuntimeError(
+                        "Session environment not initialized for websocket"
+                    )
 
-                while True:
-                    # Receive message from client
-                    raw_message = await websocket.receive_text()
+                # Keep MCP session open for entire websocket lifetime
+                # (avoids reconnect overhead on every message)
 
-                    try:
-                        message_dict = json.loads(raw_message)
-                    except json.JSONDecodeError as e:
-                        error_resp = WSErrorResponse(
-                            data={
-                                "message": f"Invalid JSON: {e}",
-                                "code": WSErrorCode.INVALID_JSON,
-                            }
+                async with AsyncExitStack() as stack:
+                    mcp_session_factory = getattr(session_env, "mcp_session", None)
+                    if callable(mcp_session_factory):
+                        mcp_session_cm = cast(
+                            AsyncContextManager[Any], mcp_session_factory()
                         )
-                        await websocket.send_text(error_resp.model_dump_json())
-                        continue
-
-                    msg_type = message_dict.get("type", "")
+                        await stack.enter_async_context(mcp_session_cm)
+
+                    while True:
+                        # Receive message from client
+                        raw_message = await websocket.receive_text()
+
+                        try:
+                            message_dict = json.loads(raw_message)
+                        except json.JSONDecodeError as e:
+                            error_resp = WSErrorResponse(
+                                data={
+                                    "message": f"Invalid JSON: {e}",
+                                    "code": WSErrorCode.INVALID_JSON,
+                                }
+                            )
+                            await websocket.send_text(error_resp.model_dump_json())
+                            continue
 
-                    try:
-                        match msg_type:
-                            case "reset":
-                                msg = WSResetMessage(**message_dict)
+                        msg_type = message_dict.get("type", "")
 
-                                is_async = (
-                                    session_env.reset_async.__func__
-                                    is not Environment.reset_async
-                                )
+                        try:
+                            match msg_type:
+                                case "reset":
+                                    msg = WSResetMessage(**message_dict)
 
-                                if is_async:
-                                    sig = inspect.signature(session_env.reset_async)
-                                    valid_kwargs = self._get_valid_kwargs(sig, msg.data)
-                                    observation = await session_env.reset_async(
-                                        **valid_kwargs
+                                    is_async = (
+                                        session_env.reset_async.__func__
+                                        is not Environment.reset_async
                                     )
-                                else:
-                                    sig = inspect.signature(session_env.reset)
-                                    valid_kwargs = self._get_valid_kwargs(sig, msg.data)
-                                    observation = await self._run_in_session_executor(
-                                        session_id, session_env.reset, **valid_kwargs
-                                    )
-
-                                self._update_session_activity(session_id)
-
-                                response = WSObservationResponse(
-                                    data=serialize_observation(observation),
-                                )
 
-                            case "step":
-                                msg = WSStepMessage(**message_dict)
-                                action = deserialize_action(msg.data, self.action_cls)
+                                    if is_async:
+                                        sig = inspect.signature(session_env.reset_async)
+                                        valid_kwargs = self._get_valid_kwargs(
+                                            sig, msg.data
+                                        )
+                                        observation = await session_env.reset_async(
+                                            **valid_kwargs
+                                        )
+                                    else:
+                                        sig = inspect.signature(session_env.reset)
+                                        valid_kwargs = self._get_valid_kwargs(
+                                            sig, msg.data
+                                        )
+                                        observation = (
+                                            await self._run_in_session_executor(
+                                                session_id,
+                                                session_env.reset,
+                                                **valid_kwargs,
+                                            )
+                                        )
+
+                                    self._update_session_activity(session_id)
+
+                                    response = WSObservationResponse(
+                                        data=serialize_observation(observation),
+                                    )
 
-                                is_async = (
-                                    session_env.step_async.__func__
-                                    is not Environment.step_async
-                                )
+                                case "step":
+                                    msg = WSStepMessage(**message_dict)
+                                    action = deserialize_action(
+                                        msg.data, self.action_cls
+                                    )
 
-                                if is_async:
-                                    observation = await session_env.step_async(action)
-                                else:
-                                    observation = await self._run_in_session_executor(
-                                        session_id, session_env.step, action
+                                    is_async = (
+                                        session_env.step_async.__func__
+                                        is not Environment.step_async
                                     )
 
-                                self._update_session_activity(
-                                    session_id, increment_step=True
-                                )
+                                    if is_async:
+                                        observation = await session_env.step_async(
+                                            action
+                                        )
+                                    else:
+                                        observation = (
+                                            await self._run_in_session_executor(
+                                                session_id, session_env.step, action
+                                            )
+                                        )
+
+                                    self._update_session_activity(
+                                        session_id, increment_step=True
+                                    )
 
-                                response = WSObservationResponse(
-                                    data=serialize_observation(observation)
-                                )
+                                    response = WSObservationResponse(
+                                        data=serialize_observation(observation)
+                                    )
 
-                            case "state":
-                                msg = WSStateMessage(**message_dict)
-                                state = session_env.state
-                                if hasattr(state, "model_dump"):
-                                    state_data = state.model_dump()
-                                else:
-                                    state_data = dict(state) if state else {}
-
-                                response = WSStateResponse(data=state_data)
-
-                            case "close":
-                                msg = WSCloseMessage(**message_dict)
-                                break
-
-                            case "mcp":
-                                msg = WSMCPMessage(**message_dict)
-                                try:
-                                    rpc_request = JsonRpcRequest(**msg.data)
-                                except (ValidationError, Exception) as e:
-                                    rpc_response = JsonRpcResponse.error_response(
-                                        JsonRpcErrorCode.INVALID_REQUEST,
-                                        f"Invalid request: {e}",
+                                case "state":
+                                    msg = WSStateMessage(**message_dict)
+                                    state = session_env.state
+                                    if hasattr(state, "model_dump"):
+                                        state_data = state.model_dump()
+                                    else:
+                                        state_data = dict(state) if state else {}
+
+                                    response = WSStateResponse(data=state_data)
+
+                                case "close":
+                                    msg = WSCloseMessage(**message_dict)
+                                    break
+
+                                case "mcp":
+                                    msg = WSMCPMessage(**message_dict)
+                                    try:
+                                        rpc_request = JsonRpcRequest(**msg.data)
+                                    except (ValidationError, Exception) as e:
+                                        rpc_response = JsonRpcResponse.error_response(
+                                            JsonRpcErrorCode.INVALID_REQUEST,
+                                            f"Invalid request: {e}",
+                                        )
+                                    else:
+                                        rpc_response = await mcp_handler(
+                                            rpc_request,
+                                            session_env=session_env,
+                                            session_id=session_id,
+                                        )
+                                    response = WSMCPResponse(
+                                        data=rpc_response.model_dump()
                                     )
-                                else:
-                                    rpc_response = await mcp_handler(
-                                        rpc_request,
-                                        session_env=session_env,
+
+                                case _:
+                                    response = WSErrorResponse(
+                                        data={
+                                            "message": f"Unknown message type: {msg_type}",
+                                            "code": WSErrorCode.UNKNOWN_TYPE,
+                                        }
                                     )
-                                response = WSMCPResponse(data=rpc_response.model_dump())
-
-                            case _:
-                                response = WSErrorResponse(
-                                    data={
-                                        "message": f"Unknown message type: {msg_type}",
-                                        "code": WSErrorCode.UNKNOWN_TYPE,
-                                    }
-                                )
 
-                        await websocket.send_text(response.model_dump_json())
+                            await websocket.send_text(response.model_dump_json())
 
-                    except ValidationError as e:
-                        error_resp = WSErrorResponse(
-                            data={
-                                "message": "Invalid message",
-                                "code": WSErrorCode.VALIDATION_ERROR,
-                                "errors": e.errors(),
-                            }
-                        )
-                        await websocket.send_text(error_resp.model_dump_json())
-                    except Exception as e:
-                        error_resp = WSErrorResponse(
-                            data={
-                                "message": str(e),
-                                "code": WSErrorCode.EXECUTION_ERROR,
-                            }
-                        )
-                        await websocket.send_text(error_resp.model_dump_json())
+                        except ValidationError as e:
+                            error_resp = WSErrorResponse(
+                                data={
+                                    "message": "Invalid message",
+                                    "code": WSErrorCode.VALIDATION_ERROR,
+                                    "errors": e.errors(),
+                                }
+                            )
+                            await websocket.send_text(error_resp.model_dump_json())
+                        except Exception as e:
+                            error_resp = WSErrorResponse(
+                                data={
+                                    "message": str(e),
+                                    "code": WSErrorCode.EXECUTION_ERROR,
+                                }
+                            )
+                            await websocket.send_text(error_resp.model_dump_json())
 
             except WebSocketDisconnect:
                 pass
@@ -1276,7 +1531,7 @@ def create_app(
         from .web_interface import create_web_interface_app
 
         return create_web_interface_app(
-            env,
+            cast(Any, env),
             action_cls,
             observation_cls,
             env_name,

src/core/openenv/core/env_server/mcp_environment.py

@@ -56,6 +56,7 @@ import asyncio
 import inspect
 from abc import abstractmethod
 from collections import defaultdict
+from contextlib import asynccontextmanager
 from typing import Any, Callable, Dict, Optional
 
 from fastmcp import Client
@@ -164,6 +165,52 @@ class MCPEnvironment(Environment):
         # Track tool schemas for list_tools: {tool_name: {mode: schema}}
         self._mode_tool_schemas = defaultdict(dict)
 
+    def _require_mcp_client(self) -> Any:
+        """Return MCP client or raise if environment has been closed."""
+        if self.mcp_client is None:
+            raise RuntimeError("MCP client is not available; environment is closed")
+        return self.mcp_client
+
+    def _require_mcp_server(self) -> Any:
+        """Return MCP server or raise if environment has been closed."""
+        if self.mcp_server is None:
+            raise RuntimeError("MCP server is not available; environment is closed")
+        return self.mcp_server
+
+    @asynccontextmanager
+    async def mcp_session(self):
+        """
+        Context manager for MCP client sessions.
+
+        This wrapper serves two purposes:
+
+        1. **Null guard** — raises a clear error if ``close()`` has already
+           been called (``mcp_client`` is ``None``).
+
+        2. **AsyncExitStack adapter** — FastMCP's ``Client.__aenter__``
+           creates a background ``asyncio.Task`` for session management.
+           When entered directly via ``AsyncExitStack`` in the HTTP session
+           path (``_create_session``), this task can be cancelled by ASGI
+           harnesses (e.g. Starlette ``TestClient``) between requests,
+           corrupting session state.  Wrapping in an ``asynccontextmanager``
+           generator isolates the task lifecycle: the generator frame keeps
+           ``async with client:`` suspended at ``yield``, so cleanup only
+           runs when the stack explicitly closes the generator — not when
+           the event loop cancels orphaned tasks.
+
+        Delegates to FastMCP's ``Client`` context manager which is
+        reentrant: the first entry opens the transport and subsequent
+        (nested) entries simply increment an internal reference counter.
+        The transport is closed only when the outermost context exits.
+
+        No external lock is needed because ``Client._connect`` /
+        ``Client._disconnect`` already serialise connection state changes
+        through their own ``anyio.Lock``.
+        """
+        client = self._require_mcp_client()
+        async with client:
+            yield client
+
     @property
     def supports_code_mode(self) -> bool:
         """Check if this environment supports code mode (execute_code)."""
@@ -292,7 +339,8 @@ class MCPEnvironment(Environment):
 
             # If mode is None, register with FastMCP as usual
             if mode is None:
-                decorated_func = self.mcp_server.tool()(func)
+                mcp_server = self._require_mcp_server()
+                decorated_func = mcp_server.tool()(func)
                 self._mode_tools[tool_name][None] = func
                 return decorated_func
 
@@ -372,24 +420,49 @@ class MCPEnvironment(Environment):
             return self._step_impl(action, timeout_s=timeout_s, **kwargs)
 
     def _handle_list_tools(self) -> ListToolsObservation:
+        """Sync wrapper — delegates to the canonical async implementation."""
+        return run_async_safely(self._async_handle_list_tools())
+
+    async def _async_list_tools(self) -> list:
         """
-        Handle a ListToolsAction by querying the MCP server.
+        Async helper to list tools from the MCP client.
 
         Returns:
-            ListToolsObservation containing all available tools with their
-            names, descriptions, and input schemas, filtered by current mode.
+            List of tool objects from the MCP server.
         """
-        try:
-            # Get current mode
-            current_mode = getattr(self, "_mode", None)
+        async with self.mcp_session() as client:
+            return await client.list_tools()
 
-            # Start with tools from FastMCP server (mode=None tools)
-            tools_result = run_async_safely(self._async_list_tools())
+    def _handle_call_tool(
+        self,
+        action: CallToolAction,
+        timeout_s: Optional[float] = None,
+    ) -> CallToolObservation:
+        """Sync wrapper — delegates to the canonical async implementation."""
+        return run_async_safely(
+            self._async_handle_call_tool(action, timeout_s=timeout_s)
+        )
 
-            # Build list of Tool objects
-            tools = []
+    async def _async_call_tool(self, tool_name: str, arguments: dict) -> Any:
+        """
+        Async helper to call a tool on the MCP server.
 
-            # Add FastMCP tools that are not mode-specific
+        Args:
+            tool_name: Name of the tool to invoke.
+            arguments: Dictionary of arguments to pass to the tool.
+
+        Returns:
+            The result from the tool execution.
+        """
+        async with self.mcp_session() as client:
+            return await client.call_tool(tool_name, arguments)
+
+    async def _async_handle_list_tools(self) -> ListToolsObservation:
+        """Async version of _handle_list_tools — avoids run_async_safely."""
+        try:
+            current_mode = getattr(self, "_mode", None)
+            tools_result = await self._async_list_tools()
+            tools = []
             for tool in tools_result:
                 if tool.name not in self._mode_tool_schemas:
                     tools.append(
@@ -401,11 +474,8 @@ class MCPEnvironment(Environment):
                             else {},
                         )
                     )
-
-            # Add mode-specific tools available in current mode
             for tool_name, mode_schemas in self._mode_tool_schemas.items():
                 if None in mode_schemas:
-                    # Tool available in all modes
                     schema = mode_schemas[None]
                     tools.append(
                         Tool(
@@ -415,7 +485,6 @@ class MCPEnvironment(Environment):
                         )
                     )
                 elif current_mode in mode_schemas:
-                    # Tool available in current mode
                     schema = mode_schemas[current_mode]
                     tools.append(
                         Tool(
@@ -424,65 +493,30 @@ class MCPEnvironment(Environment):
                             input_schema=schema["input_schema"],
                         )
                     )
-
             return ListToolsObservation(tools=tools)
-
         except Exception as e:
-            # Return an observation with error in metadata
             return ListToolsObservation(
                 tools=[],
-                metadata={
-                    "error": str(e),
-                    "error_type": "list_tools_failed",
-                },
+                metadata={"error": str(e), "error_type": "list_tools_failed"},
             )
 
-    async def _async_list_tools(self) -> list:
-        """
-        Async helper to list tools from the MCP client.
-
-        Returns:
-            List of tool objects from the MCP server.
-        """
-        async with self.mcp_client:
-            return await self.mcp_client.list_tools()
-
-    def _handle_call_tool(
+    async def _async_handle_call_tool(
         self,
         action: CallToolAction,
         timeout_s: Optional[float] = None,
     ) -> CallToolObservation:
-        """
-        Handle a CallToolAction by invoking the specified tool.
-
-        Args:
-            action: The CallToolAction containing tool_name and arguments.
-            timeout_s: Timeout in seconds. Defaults to MCP_TOOL_CALL_TIMEOUT (30s).
-
-        Returns:
-            CallToolObservation with the tool's result or an error.
-        """
+        """Async version of _handle_call_tool — avoids run_async_safely."""
         timeout = timeout_s if timeout_s is not None else MCP_TOOL_CALL_TIMEOUT
-
-        # Check if this is a mode-specific tool
         tool_name = action.tool_name
         current_mode = getattr(self, "_mode", None)
 
         if tool_name in self._mode_tools:
             mode_info = self._mode_tools[tool_name]
-
-            # Check if tool is available in current mode
-            # Tool is available if:
-            # 1. It has a None mode (available in all modes), OR
-            # 2. It has an implementation for the current mode
             if None in mode_info:
-                # Use the mode-agnostic version
                 func = mode_info[None]
             elif current_mode in mode_info:
-                # Use the mode-specific version
                 func = mode_info[current_mode]
             else:
-                # Tool not available in current mode
                 return CallToolObservation(
                     tool_name=tool_name,
                     result=None,
@@ -491,16 +525,11 @@ class MCPEnvironment(Environment):
                         message=f"Tool '{tool_name}' not available in {current_mode} mode",
                     ),
                 )
-
-            # Call the mode-specific function directly
             try:
-                # Check if function is async and await if necessary
                 if inspect.iscoroutinefunction(func):
-                    result = run_async_safely(func(**action.arguments))
+                    result = await func(**action.arguments)
                 else:
                     result = func(**action.arguments)
-
-                # Wrap result in CallToolResult format to match FastMCP behavior
                 return CallToolObservation(
                     tool_name=tool_name,
                     result=CallToolResult(
@@ -521,22 +550,12 @@ class MCPEnvironment(Environment):
                     ),
                 )
 
-        # Not a mode-specific tool, use FastMCP
         try:
-            # Run the async call_tool with timeout
-            # Use run_async_safely to handle both sync and async contexts
-            result = run_async_safely(
-                asyncio.wait_for(
-                    self._async_call_tool(action.tool_name, action.arguments),
-                    timeout=timeout,
-                )
-            )
-
-            return CallToolObservation(
-                tool_name=action.tool_name,
-                result=result,
+            result = await asyncio.wait_for(
+                self._async_call_tool(action.tool_name, action.arguments),
+                timeout=timeout,
             )
-
+            return CallToolObservation(tool_name=action.tool_name, result=result)
         except asyncio.TimeoutError:
             return CallToolObservation(
                 tool_name=action.tool_name,
@@ -546,11 +565,8 @@ class MCPEnvironment(Environment):
                     message=f"Tool '{action.tool_name}' timed out after {timeout} seconds",
                 ),
             )
-
         except Exception as e:
             error_message = str(e)
-
-            # Determine error type based on the exception
             if (
                 "not found" in error_message.lower()
                 or "unknown tool" in error_message.lower()
@@ -563,29 +579,34 @@ class MCPEnvironment(Environment):
                 error_type = ToolErrorType.INVALID_ARGS
             else:
                 error_type = ToolErrorType.EXECUTION_ERROR
-
             return CallToolObservation(
                 tool_name=action.tool_name,
                 result=None,
-                error=ToolError(
-                    error_type=error_type,
-                    message=error_message,
-                ),
+                error=ToolError(error_type=error_type, message=error_message),
             )
 
-    async def _async_call_tool(self, tool_name: str, arguments: dict) -> Any:
+    async def step_async(
+        self,
+        action: Action,
+        timeout_s: Optional[float] = None,
+        **kwargs: Any,
+    ) -> Observation:
         """
-        Async helper to call a tool on the MCP server.
+        Async step that routes MCP actions without going through run_async_safely.
 
-        Args:
-            tool_name: Name of the tool to invoke.
-            arguments: Dictionary of arguments to pass to the tool.
-
-        Returns:
-            The result from the tool execution.
+        The WebSocket handler calls this directly on the outer event loop, where
+        the MCP session is already open, avoiding the thread/event-loop deadlock
+        that occurs when the sync step() path is used via run_in_executor.
         """
-        async with self.mcp_client:
-            return await self.mcp_client.call_tool(tool_name, arguments)
+        if isinstance(action, ListToolsAction):
+            return await self._async_handle_list_tools()
+        elif isinstance(action, CallToolAction):
+            return await self._async_handle_call_tool(action, timeout_s=timeout_s)
+        else:
+            loop = asyncio.get_event_loop()
+            return await loop.run_in_executor(
+                None, lambda: self._step_impl(action, timeout_s=timeout_s, **kwargs)
+            )
 
     @abstractmethod
     def _step_impl(

src/core/openenv/core/env_server/serialization.py

@@ -14,14 +14,28 @@ HTTP server and web interface implementations.
 
 from typing import Any, Dict, Type
 
+from .mcp_types import CallToolAction, ListToolsAction
 from .types import Action, Observation
 
+# MCP action types keyed by their "type" discriminator value.
+# These are checked before the environment's own action_cls so that
+# ListToolsAction / CallToolAction payloads are never rejected by an
+# unrelated Pydantic model.
+_MCP_ACTION_TYPES: Dict[str, Type[Action]] = {
+    "list_tools": ListToolsAction,
+    "call_tool": CallToolAction,
+}
+
 
 def deserialize_action(action_data: Dict[str, Any], action_cls: Type[Action]) -> Action:
     """
     Convert JSON dict to Action instance using Pydantic validation.
 
-    This is a basic deserialization that works for most environments.
+    MCP action types (``list_tools``, ``call_tool``) are recognised
+    automatically via the ``"type"`` discriminator field, regardless of
+    the environment's configured ``action_cls``.  All other payloads
+    fall through to ``action_cls.model_validate()``.
+
     For special cases (e.g., tensor fields, custom type conversions),
     use deserialize_action_with_preprocessing().
 
@@ -38,6 +52,17 @@ def deserialize_action(action_data: Dict[str, Any], action_cls: Type[Action]) ->
     Note:
         This uses Pydantic's model_validate() for automatic validation.
     """
+    # Route MCP action types before falling through to the env action_cls.
+    # Only intercept when action_cls is the generic Action base or itself an
+    # MCP type (i.e. the server hosts an MCP environment).  This avoids
+    # silently bypassing env-specific validation for non-MCP environments
+    # that happen to use "call_tool" / "list_tools" as a type discriminator.
+    action_type = action_data.get("type")
+    if action_type in _MCP_ACTION_TYPES:
+        mcp_cls = _MCP_ACTION_TYPES[action_type]
+        if action_cls is Action or action_cls in _MCP_ACTION_TYPES.values():
+            return mcp_cls.model_validate(action_data)
+
     return action_cls.model_validate(action_data)
 
 
@@ -62,6 +87,15 @@ def deserialize_action_with_preprocessing(
     Raises:
         ValidationError: If action_data is invalid for the action class
     """
+    # Route MCP action types before preprocessing (they don't need it).
+    # Same guard as deserialize_action: only intercept when action_cls is
+    # the generic Action base or itself an MCP type.
+    action_type = action_data.get("type")
+    if action_type in _MCP_ACTION_TYPES:
+        mcp_cls = _MCP_ACTION_TYPES[action_type]
+        if action_cls is Action or action_cls in _MCP_ACTION_TYPES.values():
+            return mcp_cls.model_validate(action_data)
+
     processed_data = {}
 
     for key, value in action_data.items():

src/core/openenv/core/env_server/web_interface.py

@@ -15,13 +15,15 @@ option (e.g. openenv push --enable-interface) or ENABLE_WEB_INTERFACE env var.
 from __future__ import annotations
 
 import asyncio
+import inspect
 import json
 from concurrent.futures import ThreadPoolExecutor
 from datetime import datetime
 from typing import Any, Callable, Dict, List, Optional, Type
 
 import gradio as gr
-from fastapi import FastAPI, WebSocket, WebSocketDisconnect
+from fastapi import Body, FastAPI, HTTPException, status, WebSocket, WebSocketDisconnect
+from fastapi.responses import RedirectResponse
 from pydantic import BaseModel, ConfigDict, Field
 
 from .gradio_theme import OPENENV_GRADIO_CSS, OPENENV_GRADIO_THEME
@@ -269,6 +271,28 @@ class WebInterfaceManager:
         # Thread pool for running sync code (e.g., Playwright sync API) in async context
         self._executor = ThreadPoolExecutor(max_workers=1)
 
+    @staticmethod
+    def _get_valid_kwargs(
+        sig: inspect.Signature,
+        kwargs: Dict[str, Any],
+        skip_params: Optional[set[str]] = None,
+    ) -> Dict[str, Any]:
+        """Filter kwargs to only those accepted by the target function."""
+        skip_params = skip_params or set()
+        valid_kwargs: Dict[str, Any] = {}
+        has_var_kwargs = any(
+            param.kind == inspect.Parameter.VAR_KEYWORD
+            for param in sig.parameters.values()
+        )
+
+        for key, value in kwargs.items():
+            if key in skip_params:
+                continue
+            if key in sig.parameters or has_var_kwargs:
+                valid_kwargs[key] = value
+
+        return valid_kwargs
+
     async def _run_sync_in_thread_pool(self, func, *args, **kwargs):
         """Run a synchronous function in the thread pool executor.
 
@@ -317,11 +341,24 @@ class WebInterfaceManager:
         for client in disconnected_clients:
             self.connected_clients.remove(client)
 
-    async def reset_environment(self) -> Dict[str, Any]:
+    async def reset_environment(
+        self, reset_kwargs: Optional[Dict[str, Any]] = None
+    ) -> Dict[str, Any]:
         """Reset the environment and update state."""
-        # Run sync reset in thread pool to avoid blocking event loop
-        # and to support environments using sync libraries (e.g., Playwright)
-        observation: Observation = await self._run_sync_in_thread_pool(self.env.reset)
+        reset_kwargs = reset_kwargs or {}
+
+        is_async = self.env.reset_async.__func__ is not Environment.reset_async
+        sig = inspect.signature(self.env.reset_async if is_async else self.env.reset)
+        valid_kwargs = self._get_valid_kwargs(sig, reset_kwargs)
+
+        if is_async:
+            observation = await self.env.reset_async(**valid_kwargs)
+        else:
+            # Run sync reset in thread pool to avoid blocking event loop
+            # and to support environments using sync libraries (e.g., Playwright)
+            observation = await self._run_sync_in_thread_pool(
+                self.env.reset, **valid_kwargs
+            )
         state: State = self.env.state
 
         # Serialize observation once using shared utility
@@ -428,6 +465,16 @@ def create_web_interface_app(
     web_manager = WebInterfaceManager(env, action_cls, observation_cls, metadata)
 
     # Web API routes first (so they take precedence over Gradio mount at /web)
+    @app.get("/", include_in_schema=False)
+    async def web_root():
+        """Redirect the app root to the Gradio interface."""
+        return RedirectResponse(url="/web/")
+
+    @app.get("/web", include_in_schema=False)
+    async def web_root_no_slash():
+        """Redirect /web to /web/ for mounted Gradio deployments behind proxies."""
+        return RedirectResponse(url="/web/")
+
     @app.get("/web/metadata")
     async def web_metadata():
         """Get environment metadata."""
@@ -449,9 +496,9 @@ def create_web_interface_app(
             await web_manager.disconnect_websocket(websocket)
 
     @app.post("/web/reset")
-    async def web_reset():
+    async def web_reset(request: Optional[Dict[str, Any]] = Body(default=None)):
         """Reset endpoint for web interface."""
-        return await web_manager.reset_environment()
+        return await web_manager.reset_environment(request)
 
     @app.post("/web/step")
     async def web_step(request: Dict[str, Any]):
@@ -475,7 +522,13 @@ def create_web_interface_app(
     @app.get("/web/state")
     async def web_state():
         """State endpoint for web interface."""
-        return web_manager.get_state()
+        try:
+            return web_manager.get_state()
+        except RuntimeError as exc:
+            raise HTTPException(
+                status_code=status.HTTP_409_CONFLICT,
+                detail=str(exc),
+            ) from exc
 
     action_fields = _extract_action_fields(action_cls)
     is_chat_env = _is_chat_env(action_cls)
@@ -505,7 +558,7 @@ def create_web_interface_app(
             )
         gradio_blocks = gr.TabbedInterface(
             [default_blocks, custom_blocks],
-            tab_names=["Playground", "Visualization"],
+            tab_names=["Playground", "Custom"],
             title=get_gradio_display_title(metadata),
         )
     else:

src/core/openenv/core/mcp_client.py

@@ -52,6 +52,7 @@ Example (sync wrapper):
     ...     result = env.call_tool("echo_message", message="Hello!")
 """
 
+import asyncio
 from typing import Any, Dict, List, Optional
 
 from .client_types import StepResult
@@ -118,6 +119,66 @@ class MCPClientBase(EnvClient[Any, Observation, State]):
         )
         self._tools_cache: Optional[List[Tool]] = None
         self.use_production_mode = False
+        self._production_session_id: Optional[str] = None
+        self._production_session_lock = asyncio.Lock()
+        self._jsonrpc_request_id = 0
+        self._http_client: Optional[Any] = None  # lazily-created httpx.AsyncClient
+
+    def _next_request_id(self) -> int:
+        """Generate a monotonically increasing JSON-RPC request id."""
+        self._jsonrpc_request_id += 1
+        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"
+
+    async def _get_http_client(self) -> Any:
+        """Return a shared httpx.AsyncClient, creating one lazily."""
+        if self._http_client is None:
+            import httpx
+
+            self._http_client = httpx.AsyncClient()
+        return self._http_client
+
+    async def _production_mcp_request(
+        self, method: str, params: Optional[Dict[str, Any]] = None
+    ) -> Dict[str, Any]:
+        """Send a JSON-RPC request to HTTP /mcp and return parsed JSON response."""
+        client = await self._get_http_client()
+        response = await client.post(
+            self._production_mcp_url(),
+            json={
+                "jsonrpc": "2.0",
+                "method": method,
+                "params": params or {},
+                "id": self._next_request_id(),
+            },
+            timeout=self._message_timeout,
+        )
+        response.raise_for_status()
+        return response.json()
+
+    async def _ensure_production_session(self) -> str:
+        """Create and cache a persistent HTTP MCP session id if needed."""
+        async with self._production_session_lock:
+            if self._production_session_id is not None:
+                return self._production_session_id
+
+            data = await self._production_mcp_request("openenv/session/create")
+            if "error" in data:
+                message = data.get("error", {}).get("message", "unknown error")
+                raise RuntimeError(f"Failed to create MCP session: {message}")
+
+            session_id = data.get("result", {}).get("session_id")
+            if not session_id:
+                raise RuntimeError("Failed to create MCP session: missing session_id")
+
+            self._production_session_id = session_id
+            return session_id
 
     async def list_tools(self, use_cache: bool = True) -> List[Tool]:
         """
@@ -138,26 +199,18 @@ class MCPClientBase(EnvClient[Any, Observation, State]):
         if use_cache and self._tools_cache is not None:
             return self._tools_cache
 
-        # Use production mode HTTP endpoint if enabled
-        if self.use_production_mode:
-            import requests
-
-            # Convert ws:// URL to http:// URL
-            url = self._ws_url.replace("ws://", "http://").replace("wss://", "https://")
-            # Remove /ws suffix if present and add /mcp
-            url = url.rstrip("/ws").rstrip("/") + "/mcp"
-
+        # Use production mode HTTP endpoint if enabled.
+        # Some tests instantiate with __new__ and skip __init__, so default missing flag to False.
+        if getattr(self, "use_production_mode", False):
             try:
-                response = requests.post(
-                    url,
-                    json={
-                        "jsonrpc": "2.0",
-                        "method": "tools/list",
-                        "params": {},
-                        "id": 1,
-                    },
+                session_id = await self._ensure_production_session()
+                data = await self._production_mcp_request(
+                    "tools/list",
+                    {"session_id": session_id},
                 )
-                data = response.json()
+                if "error" in data:
+                    message = data.get("error", {}).get("message", "unknown error")
+                    raise RuntimeError(f"list_tools failed: {message}")
                 if "result" in data and "tools" in data["result"]:
                     tools = [
                         Tool(
@@ -177,7 +230,12 @@ class MCPClientBase(EnvClient[Any, Observation, State]):
             return []
 
         result = await self.step(ListToolsAction())
-        self._tools_cache = result.observation.tools
+        if isinstance(result.observation, ListToolsObservation):
+            self._tools_cache = result.observation.tools
+            return self._tools_cache
+
+        # Unexpected observation type; keep API stable with an empty tool list.
+        self._tools_cache = []
         return self._tools_cache
 
     def _step_payload(self, action: Any) -> Dict[str, Any]:
@@ -251,6 +309,35 @@ class MCPClientBase(EnvClient[Any, Observation, State]):
             step_count=payload.get("step_count", 0),
         )
 
+    async def close(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.
+        """
+        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
+
+        if self._http_client is not None:
+            try:
+                await self._http_client.aclose()
+            except Exception:
+                pass
+            finally:
+                self._http_client = None
+
+        await super().close()
+
 
 class MCPToolClient(MCPClientBase):
     """
@@ -316,6 +403,26 @@ class MCPToolClient(MCPClientBase):
             >>> result = await env.call_tool("greet", name="Claude")
             >>> print(result)  # "Hello, Claude!"
         """
+        if getattr(self, "use_production_mode", False):
+            session_id = await self._ensure_production_session()
+            data = await self._production_mcp_request(
+                "tools/call",
+                {
+                    "name": name,
+                    "arguments": kwargs,
+                    "session_id": session_id,
+                },
+            )
+
+            if "error" in data:
+                message = data.get("error", {}).get("message", "unknown error")
+                raise RuntimeError(f"Tool '{name}' failed: {message}")
+
+            result = data.get("result")
+            if isinstance(result, dict) and "data" in result:
+                return result["data"]
+            return result
+
         action = CallToolAction(tool_name=name, arguments=kwargs)
         result = await self.step(action)
         obs = result.observation

src/openenv/__init__.py

@@ -14,10 +14,18 @@ __all__ = [
     "SyncEnvClient",
 ]
 
-try:
-    __version__ = metadata.version("openenv")  # type: ignore[arg-type]
-except metadata.PackageNotFoundError:  # pragma: no cover - local dev
-    __version__ = "0.0.0"
+
+def _load_package_version() -> str:
+    """Resolve the installed distribution version for the OpenEnv package."""
+    for distribution_name in ("openenv-core", "openenv"):
+        try:
+            return metadata.version(distribution_name)
+        except metadata.PackageNotFoundError:
+            continue
+    return "0.0.0"
+
+
+__version__ = _load_package_version()
 
 
 _LAZY_MODULES = {

src/openenv/cli/templates/openenv_env/pyproject.toml

@@ -17,7 +17,7 @@ dependencies = [
     # Core OpenEnv runtime (provides FastAPI server + HTTP client types)
     # install from github
     # "openenv-core[core] @ git+https://github.com/meta-pytorch/OpenEnv.git",
-    "openenv-core[core]>=0.2.1",
+    "openenv-core[core]>=0.2.2",
     # Environment-specific dependencies
     # Add all dependencies needed for your environment here
     # Examples:

src/openenv/core/env_server/http_server.py

@@ -16,11 +16,15 @@ from __future__ import annotations
 import asyncio
 import inspect
 import json
+import logging
 import os
 import time
 import uuid
 from concurrent.futures import ThreadPoolExecutor
-from typing import Any, Callable, Dict, Optional, Type
+from contextlib import AsyncExitStack
+from typing import Any, AsyncContextManager, Callable, cast, Dict, Optional, Type
+
+_MISSING = object()
 
 from fastapi import (
     Body,
@@ -204,8 +208,9 @@ class HTTPEnvServer:
         self.observation_cls = observation_cls
 
         # Session management for WebSocket connections
-        self._sessions: Dict[str, Environment] = {}
+        self._sessions: Dict[str, Optional[Environment]] = {}
         self._session_executors: Dict[str, ThreadPoolExecutor] = {}
+        self._session_stacks: Dict[str, AsyncExitStack] = {}
         self._session_info: Dict[str, SessionInfo] = {}
         self._session_lock = asyncio.Lock()
 
@@ -213,6 +218,14 @@ class HTTPEnvServer:
         # This is needed for environments using sync libraries (e.g., Playwright)
         self._executor = ThreadPoolExecutor(max_workers=32)
 
+        # Idle session reaper configuration.
+        # Timeout is taken from ConcurrencyConfig.session_timeout;
+        # None means no timeout (default — reaper is a no-op).
+        self._session_idle_timeout_s: Optional[float] = (
+            self._concurrency_config.session_timeout
+        )
+        self._reaper_task: Optional[asyncio.Task[None]] = None
+
     def _validate_concurrency_safety(self) -> None:
         """
         Validate that the environment supports the configured concurrency level.
@@ -321,12 +334,37 @@ class HTTPEnvServer:
             )
             raise EnvironmentFactoryError(factory_name) from e
 
+        # Hold the MCP session open for the lifetime of this session,
+        # matching the WebSocket path's AsyncExitStack pattern.  This
+        # prevents per-request MCP transport teardown/reconnection and
+        # preserves FastMCP session state (ctx.set_state / ctx.get_state)
+        # across HTTP calls within the same OpenEnv session.
+        stack = AsyncExitStack()
+        try:
+            mcp_session_factory = getattr(env, "mcp_session", None)
+            if callable(mcp_session_factory):
+                mcp_session_cm = cast(AsyncContextManager[Any], mcp_session_factory())
+                await stack.enter_async_context(mcp_session_cm)
+        except Exception:
+            # MCP transport failed to start — clean up the reserved slot,
+            # the env, and the executor so they don't leak permanently
+            # against _max_concurrent_envs.
+            await stack.aclose()  # best-effort
+            async with self._session_lock:
+                self._sessions.pop(session_id, None)
+                self._session_executors.pop(session_id, None)
+                self._session_info.pop(session_id, None)
+            await self._cleanup_session_resources(env, executor)
+            raise
+
         async with self._session_lock:
             self._sessions[session_id] = env
+            self._session_stacks[session_id] = stack
+            now = time.time()
             self._session_info[session_id] = SessionInfo(
                 session_id=session_id,
                 created_at=current_time,
-                last_activity_at=current_time,
+                last_activity_at=now,
                 step_count=0,
                 environment_type=type(env).__name__,
             )
@@ -343,8 +381,27 @@ class HTTPEnvServer:
         async with self._session_lock:
             env = self._sessions.pop(session_id, None)
             executor = self._session_executors.pop(session_id, None)
+            stack = self._session_stacks.pop(session_id, None)
             self._session_info.pop(session_id, None)
 
+        await self._cleanup_session_resources(env, executor, stack)
+
+    async def _cleanup_session_resources(
+        self,
+        env: Optional[Environment],
+        executor: Optional[ThreadPoolExecutor],
+        stack: Optional[AsyncExitStack] = None,
+    ) -> None:
+        """Close an environment and shut down its executor (best-effort)."""
+        # Close the MCP session stack first — this gracefully exits the
+        # mcp_session() context (and the underlying FastMCP Client session)
+        # before we tear down the environment references.
+        if stack is not None:
+            try:
+                await stack.aclose()
+            except Exception:
+                pass  # Best effort cleanup
+
         # Run close() in the same executor where the env was created
         # This is required for thread-sensitive libraries like Playwright/greenlet
         if env is not None:
@@ -383,6 +440,51 @@ class HTTPEnvServer:
             if increment_step:
                 self._session_info[session_id].step_count += 1
 
+    async def _reap_idle_sessions(self) -> None:
+        """Background task that periodically destroys sessions idle beyond the timeout."""
+        timeout = self._session_idle_timeout_s
+        if timeout is None:
+            return  # no timeout configured — noop
+        interval = max(timeout / 4, 5.0)  # check frequently enough
+        while True:
+            try:
+                await asyncio.sleep(interval)
+                now = time.time()
+                stale_ids: list[str] = []
+                async with self._session_lock:
+                    for sid, info in self._session_info.items():
+                        if now - info.last_activity_at > timeout:
+                            stale_ids.append(sid)
+                for sid in stale_ids:
+                    # Re-check under lock: activity may have arrived since
+                    # the snapshot was taken, making this session active again.
+                    # Refresh `now` so slow _destroy_session calls don't cause
+                    # subsequent entries to be validated against a stale clock.
+                    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:
+                            continue
+                    await self._destroy_session(sid)
+            except asyncio.CancelledError:
+                break
+            except Exception as exc:
+                logging.getLogger(__name__).warning(
+                    "Idle-session reaper encountered an error (will retry): %s",
+                    exc,
+                )
+
+    def _start_reaper(self) -> None:
+        """Start the idle-session reaper if a timeout is configured."""
+        if self._session_idle_timeout_s is not None and self._reaper_task is None:
+            self._reaper_task = asyncio.create_task(self._reap_idle_sessions())
+
+    def _stop_reaper(self) -> None:
+        """Cancel the reaper background task."""
+        if self._reaper_task is not None:
+            self._reaper_task.cancel()
+            self._reaper_task = None
+
     def get_session_info(self, session_id: str) -> Optional[SessionInfo]:
         """
         Get information about a specific session.
@@ -458,6 +560,20 @@ class HTTPEnvServer:
                     f"Invalid mode: '{mode}'. Must be one of: {valid_modes}"
                 )
 
+        # Wire up idle-session reaper lifecycle via app events
+        server_ref = self
+
+        async def _start_session_reaper() -> None:
+            server_ref._start_reaper()
+
+        async def _stop_session_reaper() -> None:
+            server_ref._stop_reaper()
+
+        if not getattr(app.router, "_openenv_reaper_registered", False):
+            app.router.on_startup.append(_start_session_reaper)
+            app.router.on_shutdown.append(_stop_session_reaper)
+            app.router._openenv_reaper_registered = True  # type: ignore[attr-defined]
+
         # Helper function to handle reset endpoint
         async def reset_handler(
             request: ResetRequest = Body(default_factory=ResetRequest),
@@ -526,53 +642,214 @@ class HTTPEnvServer:
 
         # Helper function to handle MCP endpoint
         async def mcp_handler(
-            request: JsonRpcRequest, session_env: Optional[Environment] = None
+            request: JsonRpcRequest,
+            session_env: Optional[Environment] = None,
+            session_id: Optional[str] = None,
         ) -> JsonRpcResponse:
             """
             Handle MCP JSON-RPC requests.
 
-            Supports tools/list and tools/call methods in JSON-RPC 2.0 format.
+            Supports tools/list and tools/call methods in JSON-RPC 2.0 format,
+            plus OpenEnv session lifecycle methods for HTTP MCP:
+            - openenv/session/create
+            - openenv/session/close
             """
             method = request.method
             request_id = request.id
+            params = request.params
+            if not isinstance(params, dict):
+                return JsonRpcResponse.error_response(
+                    JsonRpcErrorCode.INVALID_PARAMS,
+                    "Params must be an object",
+                    request_id=request_id,
+                )
+
+            # OpenEnv extension methods for explicit MCP session management.
+            # This enables persistent MCP lifecycles over HTTP /mcp, matching WebSocket semantics.
+            if method == "openenv/session/create":
+                if session_env is not None and session_id is not None:
+                    return JsonRpcResponse.success(
+                        result={"session_id": session_id},
+                        request_id=request_id,
+                    )
+                try:
+                    created_session_id, _ = await self._create_session()
+                except SessionCapacityError as e:
+                    return JsonRpcResponse.error_response(
+                        JsonRpcErrorCode.SERVER_ERROR,
+                        str(e),
+                        request_id=request_id,
+                        data={
+                            "active_sessions": e.active_sessions,
+                            "max_sessions": e.max_sessions,
+                        },
+                    )
+                except EnvironmentFactoryError as e:
+                    return JsonRpcResponse.error_response(
+                        JsonRpcErrorCode.SERVER_ERROR,
+                        str(e),
+                        request_id=request_id,
+                        data={"factory_name": e.factory_name},
+                    )
+                return JsonRpcResponse.success(
+                    result={"session_id": created_session_id},
+                    request_id=request_id,
+                )
+
+            if method == "openenv/session/close":
+                target_session_id = params.get("session_id")
+                if not target_session_id:
+                    return JsonRpcResponse.error_response(
+                        JsonRpcErrorCode.INVALID_PARAMS,
+                        "Invalid params - 'session_id' is required",
+                        request_id=request_id,
+                    )
+
+                if session_id is not None and target_session_id == session_id:
+                    return JsonRpcResponse.error_response(
+                        JsonRpcErrorCode.INVALID_REQUEST,
+                        "Cannot close active WebSocket-managed session via MCP method",
+                        request_id=request_id,
+                    )
+
+                async with self._session_lock:
+                    env = self._sessions.pop(target_session_id, _MISSING)
+                    if env is not _MISSING:
+                        executor = self._session_executors.pop(target_session_id, None)
+                        stack = self._session_stacks.pop(target_session_id, None)
+                        self._session_info.pop(target_session_id, None)
+                    else:
+                        executor = None
+                        stack = None
+
+                if env is _MISSING:
+                    return JsonRpcResponse.error_response(
+                        JsonRpcErrorCode.INVALID_PARAMS,
+                        f"Unknown session_id: {target_session_id}",
+                        request_id=request_id,
+                    )
+
+                if env is None:
+                    # Session slot reserved but env factory still running;
+                    # re-insert the placeholder AND the executor so
+                    # _create_session can finish and the executor remains
+                    # tracked for eventual shutdown.
+                    async with self._session_lock:
+                        self._sessions[target_session_id] = None
+                        if executor is not None:
+                            self._session_executors[target_session_id] = executor
+                    return JsonRpcResponse.error_response(
+                        JsonRpcErrorCode.INVALID_REQUEST,
+                        f"Session {target_session_id} is still initializing; retry shortly",
+                        request_id=request_id,
+                    )
+
+                # env/executor/stack cleanup outside the lock
+                await self._cleanup_session_resources(env, executor, stack)
+                return JsonRpcResponse.success(
+                    result={"session_id": target_session_id, "closed": True},
+                    request_id=request_id,
+                )
+
+            requested_session_id = params.get("session_id")
+            managed_session_id = session_id
 
             # Use provided session environment or create temporary one
             if session_env is not None:
                 _env = session_env
                 should_close = False
+            elif requested_session_id:
+                async with self._session_lock:
+                    _env = self._sessions.get(requested_session_id, _MISSING)
+
+                if _env is _MISSING:
+                    return JsonRpcResponse.error_response(
+                        JsonRpcErrorCode.INVALID_PARAMS,
+                        f"Unknown session_id: {requested_session_id}",
+                        request_id=request_id,
+                    )
+
+                if _env is None:
+                    return JsonRpcResponse.error_response(
+                        JsonRpcErrorCode.INVALID_REQUEST,
+                        f"Session {requested_session_id} is still initializing; retry shortly",
+                        request_id=request_id,
+                    )
+
+                should_close = False
+                managed_session_id = requested_session_id
             else:
                 _env = self._env_factory()
                 should_close = True
             try:
+                mcp_client = getattr(_env, "mcp_client", None)
+                mcp_server = getattr(_env, "mcp_server", None)
+                mcp_session_factory = getattr(_env, "mcp_session", None)
+
                 if method == McpMethod.TOOLS_LIST:
                     # Check if environment is MCP-enabled
-                    if not hasattr(_env, "mcp_client"):
+                    if mcp_client is None and mcp_server is None:
                         return JsonRpcResponse.error_response(
                             JsonRpcErrorCode.INTERNAL_ERROR,
                             "Environment does not support MCP",
                             request_id=request_id,
                         )
 
-                    # Use async context manager for MCP client
-                    async with _env.mcp_client:
-                        tools = await _env.mcp_client.list_tools()
+                    if mcp_client:
+                        if managed_session_id and mcp_client.is_connected():
+                            # Session-managed with live transport — call
+                            # directly, no redundant re-entry.
+                            tools = await mcp_client.list_tools()
+                        elif callable(mcp_session_factory):
+                            # Stateless request, or session-managed but the
+                            # background transport was lost: (re-)open.
+                            mcp_session_cm = cast(
+                                AsyncContextManager[Any], mcp_session_factory()
+                            )
+                            async with mcp_session_cm:
+                                tools = await mcp_client.list_tools()
+                        else:
+                            async with mcp_client:
+                                tools = await mcp_client.list_tools()
+
+                        return JsonRpcResponse.success(
+                            result={
+                                "tools": [
+                                    t.model_dump()
+                                    if hasattr(t, "model_dump")
+                                    else dict(t)
+                                    for t in tools
+                                ]
+                            },
+                            request_id=request_id,
+                        )
 
-                    return JsonRpcResponse.success(
-                        result={
-                            "tools": [
-                                t.model_dump() if hasattr(t, "model_dump") else dict(t)
-                                for t in tools
-                            ]
-                        },
+                    if mcp_server:
+                        tools = []
+                        for _tool_name, tool in get_server_tools(mcp_server).items():
+                            tools.append(
+                                {
+                                    "name": tool.name,
+                                    "description": tool.description or "",
+                                    "inputSchema": tool.parameters or {},
+                                }
+                            )
+                        return JsonRpcResponse.success(
+                            result={"tools": tools},
+                            request_id=request_id,
+                        )
+
+                    return JsonRpcResponse.error_response(
+                        JsonRpcErrorCode.INTERNAL_ERROR,
+                        "MCP server not available",
                         request_id=request_id,
                     )
 
                 elif method == McpMethod.TOOLS_CALL:
-                    params = request.params
                     tool_name = params.get("name")
                     arguments = params.get("arguments", {})
 
-                    if not hasattr(_env, "mcp_client"):
+                    if mcp_client is None and mcp_server is None:
                         return JsonRpcResponse.error_response(
                             JsonRpcErrorCode.INTERNAL_ERROR,
                             "Environment does not support MCP",
@@ -581,15 +858,51 @@ class HTTPEnvServer:
 
                     if not tool_name:
                         return JsonRpcResponse.error_response(
-                            JsonRpcErrorCode.INVALID_REQUEST,
+                            JsonRpcErrorCode.INVALID_PARAMS,
                             "Missing 'name' in params",
                             request_id=request_id,
                         )
 
-                    # Use async context manager for MCP client
-                    async with _env.mcp_client:
-                        result = await _env.mcp_client.call_tool(
-                            name=tool_name, arguments=arguments
+                    if mcp_client:
+                        if managed_session_id and mcp_client.is_connected():
+                            # Session-managed with live transport.
+                            result = await mcp_client.call_tool(
+                                name=tool_name, arguments=arguments
+                            )
+                        elif callable(mcp_session_factory):
+                            # Stateless request, or session-managed but the
+                            # background transport was lost: (re-)open.
+                            mcp_session_cm = cast(
+                                AsyncContextManager[Any], mcp_session_factory()
+                            )
+                            async with mcp_session_cm:
+                                result = await mcp_client.call_tool(
+                                    name=tool_name, arguments=arguments
+                                )
+                        else:
+                            async with mcp_client:
+                                result = await mcp_client.call_tool(
+                                    name=tool_name, arguments=arguments
+                                )
+                    elif mcp_server:
+                        server_tools = get_server_tools(mcp_server)
+                        if tool_name in server_tools:
+                            tool = server_tools[tool_name]
+                            if inspect.iscoroutinefunction(tool.fn):
+                                result = await tool.fn(**arguments)
+                            else:
+                                result = tool.fn(**arguments)
+                        else:
+                            return JsonRpcResponse.error_response(
+                                JsonRpcErrorCode.INVALID_PARAMS,
+                                f"Tool not found: {tool_name}",
+                                request_id=request_id,
+                            )
+                    else:
+                        return JsonRpcResponse.error_response(
+                            JsonRpcErrorCode.INTERNAL_ERROR,
+                            "MCP server not available",
+                            request_id=request_id,
                         )
 
                     # Ensure result is JSON serializable
@@ -614,6 +927,11 @@ class HTTPEnvServer:
                     request_id=request_id,
                 )
             finally:
+                if managed_session_id:
+                    self._update_session_activity(
+                        managed_session_id,
+                        increment_step=(method == McpMethod.TOOLS_CALL),
+                    )
                 if should_close:
                     _env.close()
 
@@ -637,42 +955,59 @@ class HTTPEnvServer:
             try:
                 # Create session with dedicated environment
                 session_id, session_env = await self._create_session()
+                if session_env is None:
+                    raise RuntimeError(
+                        "Session environment not initialized for MCP websocket"
+                    )
 
-                while True:
-                    # Receive message from client
-                    raw_message = await websocket.receive_text()
-
-                    try:
-                        jsonrpc_dict = json.loads(raw_message)
-                        jsonrpc_request = JsonRpcRequest(**jsonrpc_dict)
-                    except json.JSONDecodeError as e:
-                        error_resp = JsonRpcResponse.error_response(
-                            JsonRpcErrorCode.PARSE_ERROR,
-                            f"Parse error: {e}",
-                        )
-                        await websocket.send_text(error_resp.model_dump_json())
-                        continue
-                    except ValidationError as e:
-                        error_resp = JsonRpcResponse.error_response(
-                            JsonRpcErrorCode.INVALID_REQUEST,
-                            f"Invalid request: {e}",
-                        )
-                        await websocket.send_text(error_resp.model_dump_json())
-                        continue
+                # If environment has an mcp_session context manager, hold it open
+                # for the lifetime of the websocket connection
 
-                    try:
-                        # Call mcp_handler with session environment
-                        response = await mcp_handler(
-                            jsonrpc_request, session_env=session_env
+                async with AsyncExitStack() as stack:
+                    mcp_session_factory = getattr(session_env, "mcp_session", None)
+                    if callable(mcp_session_factory):
+                        mcp_session_cm = cast(
+                            AsyncContextManager[Any], mcp_session_factory()
                         )
-                        await websocket.send_text(response.model_dump_json())
-                    except Exception as e:
-                        error_resp = JsonRpcResponse.error_response(
-                            JsonRpcErrorCode.INTERNAL_ERROR,
-                            str(e),
-                            request_id=jsonrpc_request.id,
-                        )
-                        await websocket.send_text(error_resp.model_dump_json())
+                        await stack.enter_async_context(mcp_session_cm)
+
+                    while True:
+                        # Receive message from client
+                        raw_message = await websocket.receive_text()
+
+                        try:
+                            jsonrpc_dict = json.loads(raw_message)
+                            jsonrpc_request = JsonRpcRequest(**jsonrpc_dict)
+                        except json.JSONDecodeError as e:
+                            error_resp = JsonRpcResponse.error_response(
+                                JsonRpcErrorCode.PARSE_ERROR,
+                                f"Parse error: {e}",
+                            )
+                            await websocket.send_text(error_resp.model_dump_json())
+                            continue
+                        except ValidationError as e:
+                            error_resp = JsonRpcResponse.error_response(
+                                JsonRpcErrorCode.INVALID_REQUEST,
+                                f"Invalid request: {e}",
+                            )
+                            await websocket.send_text(error_resp.model_dump_json())
+                            continue
+
+                        try:
+                            # Call mcp_handler with session environment
+                            response = await mcp_handler(
+                                jsonrpc_request,
+                                session_env=session_env,
+                                session_id=session_id,
+                            )
+                            await websocket.send_text(response.model_dump_json())
+                        except Exception as e:
+                            error_resp = JsonRpcResponse.error_response(
+                                JsonRpcErrorCode.INTERNAL_ERROR,
+                                str(e),
+                                request_id=jsonrpc_request.id,
+                            )
+                            await websocket.send_text(error_resp.model_dump_json())
 
             except WebSocketDisconnect:
                 pass
@@ -931,120 +1266,8 @@ all schema information needed to interact with the environment.
                     JsonRpcErrorCode.PARSE_ERROR
                 ).model_dump()
 
-            method = request.method
-            params = request.params
-            request_id = request.id
-
-            # Create a temporary environment for MCP access
-            _env = self._env_factory()
-
-            try:
-                # Check if environment supports MCP
-                if not hasattr(_env, "mcp_client") and not hasattr(_env, "mcp_server"):
-                    return JsonRpcResponse.error_response(
-                        JsonRpcErrorCode.INTERNAL_ERROR,
-                        "Environment does not support MCP",
-                        request_id=request_id,
-                    ).model_dump()
-
-                if method == McpMethod.TOOLS_LIST:
-                    # List tools from MCP server
-                    if hasattr(_env, "mcp_client") and _env.mcp_client:
-                        async with _env.mcp_client:
-                            tools = await _env.mcp_client.list_tools()
-                        return JsonRpcResponse.success(
-                            result={
-                                "tools": [
-                                    t.model_dump()
-                                    if hasattr(t, "model_dump")
-                                    else dict(t)
-                                    for t in tools
-                                ]
-                            },
-                            request_id=request_id,
-                        ).model_dump()
-                    elif hasattr(_env, "mcp_server") and _env.mcp_server:
-                        # Use server directly
-                        tools = []
-                        for tool_name, tool in get_server_tools(
-                            _env.mcp_server
-                        ).items():
-                            tool_dict = {
-                                "name": tool.name,
-                                "description": tool.description or "",
-                                "inputSchema": tool.parameters or {},
-                            }
-                            tools.append(tool_dict)
-                        return JsonRpcResponse.success(
-                            result={"tools": tools},
-                            request_id=request_id,
-                        ).model_dump()
-                    else:
-                        return JsonRpcResponse.error_response(
-                            JsonRpcErrorCode.INTERNAL_ERROR,
-                            "MCP server not available",
-                            request_id=request_id,
-                        ).model_dump()
-
-                elif method == McpMethod.TOOLS_CALL:
-                    tool_name = params.get("name")
-                    arguments = params.get("arguments", {})
-
-                    if not tool_name:
-                        return JsonRpcResponse.error_response(
-                            JsonRpcErrorCode.INVALID_PARAMS,
-                            "Invalid params - 'name' is required",
-                            request_id=request_id,
-                        ).model_dump()
-
-                    # Call tool via MCP
-                    if hasattr(_env, "mcp_client") and _env.mcp_client:
-                        async with _env.mcp_client:
-                            result = await _env.mcp_client.call_tool(
-                                name=tool_name, arguments=arguments
-                            )
-                    elif hasattr(_env, "mcp_server") and _env.mcp_server:
-                        # Call tool directly on FastMCP server
-                        server_tools = get_server_tools(_env.mcp_server)
-                        if tool_name in server_tools:
-                            tool = server_tools[tool_name]
-                            result = tool.fn(**arguments)
-                        else:
-                            return JsonRpcResponse.error_response(
-                                JsonRpcErrorCode.INVALID_PARAMS,
-                                f"Tool not found: {tool_name}",
-                                request_id=request_id,
-                            ).model_dump()
-                    else:
-                        return JsonRpcResponse.error_response(
-                            JsonRpcErrorCode.INTERNAL_ERROR,
-                            "MCP server not available",
-                            request_id=request_id,
-                        ).model_dump()
-
-                    # Make result JSON serializable
-                    serializable_result = _make_json_serializable(result)
-
-                    return JsonRpcResponse.success(
-                        result=serializable_result,
-                        request_id=request_id,
-                    ).model_dump()
-
-                else:
-                    return JsonRpcResponse.error_response(
-                        JsonRpcErrorCode.METHOD_NOT_FOUND,
-                        f"Method not found: {method}",
-                        request_id=request_id,
-                    ).model_dump()
-
-            except Exception as e:
-                return JsonRpcResponse.error_response(
-                    JsonRpcErrorCode.INTERNAL_ERROR,
-                    str(e),
-                    request_id=request_id,
-                ).model_dump()
-            finally:
-                _env.close()
+            response = await mcp_handler(request)
+            return response.model_dump()
 
         # Register WebSocket endpoint for persistent sessions
         @app.websocket("/ws")
@@ -1066,135 +1289,167 @@ all schema information needed to interact with the environment.
             try:
                 # Create session with dedicated environment
                 session_id, session_env = await self._create_session()
+                if session_env is None:
+                    raise RuntimeError(
+                        "Session environment not initialized for websocket"
+                    )
 
-                while True:
-                    # Receive message from client
-                    raw_message = await websocket.receive_text()
+                # Keep MCP session open for entire websocket lifetime
+                # (avoids reconnect overhead on every message)
 
-                    try:
-                        message_dict = json.loads(raw_message)
-                    except json.JSONDecodeError as e:
-                        error_resp = WSErrorResponse(
-                            data={
-                                "message": f"Invalid JSON: {e}",
-                                "code": WSErrorCode.INVALID_JSON,
-                            }
+                async with AsyncExitStack() as stack:
+                    mcp_session_factory = getattr(session_env, "mcp_session", None)
+                    if callable(mcp_session_factory):
+                        mcp_session_cm = cast(
+                            AsyncContextManager[Any], mcp_session_factory()
                         )
-                        await websocket.send_text(error_resp.model_dump_json())
-                        continue
-
-                    msg_type = message_dict.get("type", "")
+                        await stack.enter_async_context(mcp_session_cm)
+
+                    while True:
+                        # Receive message from client
+                        raw_message = await websocket.receive_text()
+
+                        try:
+                            message_dict = json.loads(raw_message)
+                        except json.JSONDecodeError as e:
+                            error_resp = WSErrorResponse(
+                                data={
+                                    "message": f"Invalid JSON: {e}",
+                                    "code": WSErrorCode.INVALID_JSON,
+                                }
+                            )
+                            await websocket.send_text(error_resp.model_dump_json())
+                            continue
 
-                    try:
-                        match msg_type:
-                            case "reset":
-                                msg = WSResetMessage(**message_dict)
+                        msg_type = message_dict.get("type", "")
 
-                                is_async = (
-                                    session_env.reset_async.__func__
-                                    is not Environment.reset_async
-                                )
+                        try:
+                            match msg_type:
+                                case "reset":
+                                    msg = WSResetMessage(**message_dict)
 
-                                if is_async:
-                                    sig = inspect.signature(session_env.reset_async)
-                                    valid_kwargs = self._get_valid_kwargs(sig, msg.data)
-                                    observation = await session_env.reset_async(
-                                        **valid_kwargs
+                                    is_async = (
+                                        session_env.reset_async.__func__
+                                        is not Environment.reset_async
                                     )
-                                else:
-                                    sig = inspect.signature(session_env.reset)
-                                    valid_kwargs = self._get_valid_kwargs(sig, msg.data)
-                                    observation = await self._run_in_session_executor(
-                                        session_id, session_env.reset, **valid_kwargs
-                                    )
-
-                                self._update_session_activity(session_id)
-
-                                response = WSObservationResponse(
-                                    data=serialize_observation(observation),
-                                )
 
-                            case "step":
-                                msg = WSStepMessage(**message_dict)
-                                action = deserialize_action(msg.data, self.action_cls)
+                                    if is_async:
+                                        sig = inspect.signature(session_env.reset_async)
+                                        valid_kwargs = self._get_valid_kwargs(
+                                            sig, msg.data
+                                        )
+                                        observation = await session_env.reset_async(
+                                            **valid_kwargs
+                                        )
+                                    else:
+                                        sig = inspect.signature(session_env.reset)
+                                        valid_kwargs = self._get_valid_kwargs(
+                                            sig, msg.data
+                                        )
+                                        observation = (
+                                            await self._run_in_session_executor(
+                                                session_id,
+                                                session_env.reset,
+                                                **valid_kwargs,
+                                            )
+                                        )
+
+                                    self._update_session_activity(session_id)
+
+                                    response = WSObservationResponse(
+                                        data=serialize_observation(observation),
+                                    )
 
-                                is_async = (
-                                    session_env.step_async.__func__
-                                    is not Environment.step_async
-                                )
+                                case "step":
+                                    msg = WSStepMessage(**message_dict)
+                                    action = deserialize_action(
+                                        msg.data, self.action_cls
+                                    )
 
-                                if is_async:
-                                    observation = await session_env.step_async(action)
-                                else:
-                                    observation = await self._run_in_session_executor(
-                                        session_id, session_env.step, action
+                                    is_async = (
+                                        session_env.step_async.__func__
+                                        is not Environment.step_async
                                     )
 
-                                self._update_session_activity(
-                                    session_id, increment_step=True
-                                )
+                                    if is_async:
+                                        observation = await session_env.step_async(
+                                            action
+                                        )
+                                    else:
+                                        observation = (
+                                            await self._run_in_session_executor(
+                                                session_id, session_env.step, action
+                                            )
+                                        )
+
+                                    self._update_session_activity(
+                                        session_id, increment_step=True
+                                    )
 
-                                response = WSObservationResponse(
-                                    data=serialize_observation(observation)
-                                )
+                                    response = WSObservationResponse(
+                                        data=serialize_observation(observation)
+                                    )
 
-                            case "state":
-                                msg = WSStateMessage(**message_dict)
-                                state = session_env.state
-                                if hasattr(state, "model_dump"):
-                                    state_data = state.model_dump()
-                                else:
-                                    state_data = dict(state) if state else {}
-
-                                response = WSStateResponse(data=state_data)
-
-                            case "close":
-                                msg = WSCloseMessage(**message_dict)
-                                break
-
-                            case "mcp":
-                                msg = WSMCPMessage(**message_dict)
-                                try:
-                                    rpc_request = JsonRpcRequest(**msg.data)
-                                except (ValidationError, Exception) as e:
-                                    rpc_response = JsonRpcResponse.error_response(
-                                        JsonRpcErrorCode.INVALID_REQUEST,
-                                        f"Invalid request: {e}",
+                                case "state":
+                                    msg = WSStateMessage(**message_dict)
+                                    state = session_env.state
+                                    if hasattr(state, "model_dump"):
+                                        state_data = state.model_dump()
+                                    else:
+                                        state_data = dict(state) if state else {}
+
+                                    response = WSStateResponse(data=state_data)
+
+                                case "close":
+                                    msg = WSCloseMessage(**message_dict)
+                                    break
+
+                                case "mcp":
+                                    msg = WSMCPMessage(**message_dict)
+                                    try:
+                                        rpc_request = JsonRpcRequest(**msg.data)
+                                    except (ValidationError, Exception) as e:
+                                        rpc_response = JsonRpcResponse.error_response(
+                                            JsonRpcErrorCode.INVALID_REQUEST,
+                                            f"Invalid request: {e}",
+                                        )
+                                    else:
+                                        rpc_response = await mcp_handler(
+                                            rpc_request,
+                                            session_env=session_env,
+                                            session_id=session_id,
+                                        )
+                                    response = WSMCPResponse(
+                                        data=rpc_response.model_dump()
                                     )
-                                else:
-                                    rpc_response = await mcp_handler(
-                                        rpc_request,
-                                        session_env=session_env,
+
+                                case _:
+                                    response = WSErrorResponse(
+                                        data={
+                                            "message": f"Unknown message type: {msg_type}",
+                                            "code": WSErrorCode.UNKNOWN_TYPE,
+                                        }
                                     )
-                                response = WSMCPResponse(data=rpc_response.model_dump())
-
-                            case _:
-                                response = WSErrorResponse(
-                                    data={
-                                        "message": f"Unknown message type: {msg_type}",
-                                        "code": WSErrorCode.UNKNOWN_TYPE,
-                                    }
-                                )
 
-                        await websocket.send_text(response.model_dump_json())
+                            await websocket.send_text(response.model_dump_json())
 
-                    except ValidationError as e:
-                        error_resp = WSErrorResponse(
-                            data={
-                                "message": "Invalid message",
-                                "code": WSErrorCode.VALIDATION_ERROR,
-                                "errors": e.errors(),
-                            }
-                        )
-                        await websocket.send_text(error_resp.model_dump_json())
-                    except Exception as e:
-                        error_resp = WSErrorResponse(
-                            data={
-                                "message": str(e),
-                                "code": WSErrorCode.EXECUTION_ERROR,
-                            }
-                        )
-                        await websocket.send_text(error_resp.model_dump_json())
+                        except ValidationError as e:
+                            error_resp = WSErrorResponse(
+                                data={
+                                    "message": "Invalid message",
+                                    "code": WSErrorCode.VALIDATION_ERROR,
+                                    "errors": e.errors(),
+                                }
+                            )
+                            await websocket.send_text(error_resp.model_dump_json())
+                        except Exception as e:
+                            error_resp = WSErrorResponse(
+                                data={
+                                    "message": str(e),
+                                    "code": WSErrorCode.EXECUTION_ERROR,
+                                }
+                            )
+                            await websocket.send_text(error_resp.model_dump_json())
 
             except WebSocketDisconnect:
                 pass
@@ -1276,7 +1531,7 @@ def create_app(
         from .web_interface import create_web_interface_app
 
         return create_web_interface_app(
-            env,
+            cast(Any, env),
             action_cls,
             observation_cls,
             env_name,

src/openenv/core/env_server/mcp_environment.py

@@ -56,6 +56,7 @@ import asyncio
 import inspect
 from abc import abstractmethod
 from collections import defaultdict
+from contextlib import asynccontextmanager
 from typing import Any, Callable, Dict, Optional
 
 from fastmcp import Client
@@ -164,6 +165,52 @@ class MCPEnvironment(Environment):
         # Track tool schemas for list_tools: {tool_name: {mode: schema}}
         self._mode_tool_schemas = defaultdict(dict)
 
+    def _require_mcp_client(self) -> Any:
+        """Return MCP client or raise if environment has been closed."""
+        if self.mcp_client is None:
+            raise RuntimeError("MCP client is not available; environment is closed")
+        return self.mcp_client
+
+    def _require_mcp_server(self) -> Any:
+        """Return MCP server or raise if environment has been closed."""
+        if self.mcp_server is None:
+            raise RuntimeError("MCP server is not available; environment is closed")
+        return self.mcp_server
+
+    @asynccontextmanager
+    async def mcp_session(self):
+        """
+        Context manager for MCP client sessions.
+
+        This wrapper serves two purposes:
+
+        1. **Null guard** — raises a clear error if ``close()`` has already
+           been called (``mcp_client`` is ``None``).
+
+        2. **AsyncExitStack adapter** — FastMCP's ``Client.__aenter__``
+           creates a background ``asyncio.Task`` for session management.
+           When entered directly via ``AsyncExitStack`` in the HTTP session
+           path (``_create_session``), this task can be cancelled by ASGI
+           harnesses (e.g. Starlette ``TestClient``) between requests,
+           corrupting session state.  Wrapping in an ``asynccontextmanager``
+           generator isolates the task lifecycle: the generator frame keeps
+           ``async with client:`` suspended at ``yield``, so cleanup only
+           runs when the stack explicitly closes the generator — not when
+           the event loop cancels orphaned tasks.
+
+        Delegates to FastMCP's ``Client`` context manager which is
+        reentrant: the first entry opens the transport and subsequent
+        (nested) entries simply increment an internal reference counter.
+        The transport is closed only when the outermost context exits.
+
+        No external lock is needed because ``Client._connect`` /
+        ``Client._disconnect`` already serialise connection state changes
+        through their own ``anyio.Lock``.
+        """
+        client = self._require_mcp_client()
+        async with client:
+            yield client
+
     @property
     def supports_code_mode(self) -> bool:
         """Check if this environment supports code mode (execute_code)."""
@@ -292,7 +339,8 @@ class MCPEnvironment(Environment):
 
             # If mode is None, register with FastMCP as usual
             if mode is None:
-                decorated_func = self.mcp_server.tool()(func)
+                mcp_server = self._require_mcp_server()
+                decorated_func = mcp_server.tool()(func)
                 self._mode_tools[tool_name][None] = func
                 return decorated_func
 
@@ -372,24 +420,49 @@ class MCPEnvironment(Environment):
             return self._step_impl(action, timeout_s=timeout_s, **kwargs)
 
     def _handle_list_tools(self) -> ListToolsObservation:
+        """Sync wrapper — delegates to the canonical async implementation."""
+        return run_async_safely(self._async_handle_list_tools())
+
+    async def _async_list_tools(self) -> list:
         """
-        Handle a ListToolsAction by querying the MCP server.
+        Async helper to list tools from the MCP client.
 
         Returns:
-            ListToolsObservation containing all available tools with their
-            names, descriptions, and input schemas, filtered by current mode.
+            List of tool objects from the MCP server.
         """
-        try:
-            # Get current mode
-            current_mode = getattr(self, "_mode", None)
+        async with self.mcp_session() as client:
+            return await client.list_tools()
 
-            # Start with tools from FastMCP server (mode=None tools)
-            tools_result = run_async_safely(self._async_list_tools())
+    def _handle_call_tool(
+        self,
+        action: CallToolAction,
+        timeout_s: Optional[float] = None,
+    ) -> CallToolObservation:
+        """Sync wrapper — delegates to the canonical async implementation."""
+        return run_async_safely(
+            self._async_handle_call_tool(action, timeout_s=timeout_s)
+        )
 
-            # Build list of Tool objects
-            tools = []
+    async def _async_call_tool(self, tool_name: str, arguments: dict) -> Any:
+        """
+        Async helper to call a tool on the MCP server.
 
-            # Add FastMCP tools that are not mode-specific
+        Args:
+            tool_name: Name of the tool to invoke.
+            arguments: Dictionary of arguments to pass to the tool.
+
+        Returns:
+            The result from the tool execution.
+        """
+        async with self.mcp_session() as client:
+            return await client.call_tool(tool_name, arguments)
+
+    async def _async_handle_list_tools(self) -> ListToolsObservation:
+        """Async version of _handle_list_tools — avoids run_async_safely."""
+        try:
+            current_mode = getattr(self, "_mode", None)
+            tools_result = await self._async_list_tools()
+            tools = []
             for tool in tools_result:
                 if tool.name not in self._mode_tool_schemas:
                     tools.append(
@@ -401,11 +474,8 @@ class MCPEnvironment(Environment):
                             else {},
                         )
                     )
-
-            # Add mode-specific tools available in current mode
             for tool_name, mode_schemas in self._mode_tool_schemas.items():
                 if None in mode_schemas:
-                    # Tool available in all modes
                     schema = mode_schemas[None]
                     tools.append(
                         Tool(
@@ -415,7 +485,6 @@ class MCPEnvironment(Environment):
                         )
                     )
                 elif current_mode in mode_schemas:
-                    # Tool available in current mode
                     schema = mode_schemas[current_mode]
                     tools.append(
                         Tool(
@@ -424,65 +493,30 @@ class MCPEnvironment(Environment):
                             input_schema=schema["input_schema"],
                         )
                     )
-
             return ListToolsObservation(tools=tools)
-
         except Exception as e:
-            # Return an observation with error in metadata
             return ListToolsObservation(
                 tools=[],
-                metadata={
-                    "error": str(e),
-                    "error_type": "list_tools_failed",
-                },
+                metadata={"error": str(e), "error_type": "list_tools_failed"},
             )
 
-    async def _async_list_tools(self) -> list:
-        """
-        Async helper to list tools from the MCP client.
-
-        Returns:
-            List of tool objects from the MCP server.
-        """
-        async with self.mcp_client:
-            return await self.mcp_client.list_tools()
-
-    def _handle_call_tool(
+    async def _async_handle_call_tool(
         self,
         action: CallToolAction,
         timeout_s: Optional[float] = None,
     ) -> CallToolObservation:
-        """
-        Handle a CallToolAction by invoking the specified tool.
-
-        Args:
-            action: The CallToolAction containing tool_name and arguments.
-            timeout_s: Timeout in seconds. Defaults to MCP_TOOL_CALL_TIMEOUT (30s).
-
-        Returns:
-            CallToolObservation with the tool's result or an error.
-        """
+        """Async version of _handle_call_tool — avoids run_async_safely."""
         timeout = timeout_s if timeout_s is not None else MCP_TOOL_CALL_TIMEOUT
-
-        # Check if this is a mode-specific tool
         tool_name = action.tool_name
         current_mode = getattr(self, "_mode", None)
 
         if tool_name in self._mode_tools:
             mode_info = self._mode_tools[tool_name]
-
-            # Check if tool is available in current mode
-            # Tool is available if:
-            # 1. It has a None mode (available in all modes), OR
-            # 2. It has an implementation for the current mode
             if None in mode_info:
-                # Use the mode-agnostic version
                 func = mode_info[None]
             elif current_mode in mode_info:
-                # Use the mode-specific version
                 func = mode_info[current_mode]
             else:
-                # Tool not available in current mode
                 return CallToolObservation(
                     tool_name=tool_name,
                     result=None,
@@ -491,16 +525,11 @@ class MCPEnvironment(Environment):
                         message=f"Tool '{tool_name}' not available in {current_mode} mode",
                     ),
                 )
-
-            # Call the mode-specific function directly
             try:
-                # Check if function is async and await if necessary
                 if inspect.iscoroutinefunction(func):
-                    result = run_async_safely(func(**action.arguments))
+                    result = await func(**action.arguments)
                 else:
                     result = func(**action.arguments)
-
-                # Wrap result in CallToolResult format to match FastMCP behavior
                 return CallToolObservation(
                     tool_name=tool_name,
                     result=CallToolResult(
@@ -521,22 +550,12 @@ class MCPEnvironment(Environment):
                     ),
                 )
 
-        # Not a mode-specific tool, use FastMCP
         try:
-            # Run the async call_tool with timeout
-            # Use run_async_safely to handle both sync and async contexts
-            result = run_async_safely(
-                asyncio.wait_for(
-                    self._async_call_tool(action.tool_name, action.arguments),
-                    timeout=timeout,
-                )
-            )
-
-            return CallToolObservation(
-                tool_name=action.tool_name,
-                result=result,
+            result = await asyncio.wait_for(
+                self._async_call_tool(action.tool_name, action.arguments),
+                timeout=timeout,
             )
-
+            return CallToolObservation(tool_name=action.tool_name, result=result)
         except asyncio.TimeoutError:
             return CallToolObservation(
                 tool_name=action.tool_name,
@@ -546,11 +565,8 @@ class MCPEnvironment(Environment):
                     message=f"Tool '{action.tool_name}' timed out after {timeout} seconds",
                 ),
             )
-
         except Exception as e:
             error_message = str(e)
-
-            # Determine error type based on the exception
             if (
                 "not found" in error_message.lower()
                 or "unknown tool" in error_message.lower()
@@ -563,29 +579,34 @@ class MCPEnvironment(Environment):
                 error_type = ToolErrorType.INVALID_ARGS
             else:
                 error_type = ToolErrorType.EXECUTION_ERROR
-
             return CallToolObservation(
                 tool_name=action.tool_name,
                 result=None,
-                error=ToolError(
-                    error_type=error_type,
-                    message=error_message,
-                ),
+                error=ToolError(error_type=error_type, message=error_message),
             )
 
-    async def _async_call_tool(self, tool_name: str, arguments: dict) -> Any:
+    async def step_async(
+        self,
+        action: Action,
+        timeout_s: Optional[float] = None,
+        **kwargs: Any,
+    ) -> Observation:
         """
-        Async helper to call a tool on the MCP server.
+        Async step that routes MCP actions without going through run_async_safely.
 
-        Args:
-            tool_name: Name of the tool to invoke.
-            arguments: Dictionary of arguments to pass to the tool.
-
-        Returns:
-            The result from the tool execution.
+        The WebSocket handler calls this directly on the outer event loop, where
+        the MCP session is already open, avoiding the thread/event-loop deadlock
+        that occurs when the sync step() path is used via run_in_executor.
         """
-        async with self.mcp_client:
-            return await self.mcp_client.call_tool(tool_name, arguments)
+        if isinstance(action, ListToolsAction):
+            return await self._async_handle_list_tools()
+        elif isinstance(action, CallToolAction):
+            return await self._async_handle_call_tool(action, timeout_s=timeout_s)
+        else:
+            loop = asyncio.get_event_loop()
+            return await loop.run_in_executor(
+                None, lambda: self._step_impl(action, timeout_s=timeout_s, **kwargs)
+            )
 
     @abstractmethod
     def _step_impl(

src/openenv/core/env_server/serialization.py

@@ -14,14 +14,28 @@ HTTP server and web interface implementations.
 
 from typing import Any, Dict, Type
 
+from .mcp_types import CallToolAction, ListToolsAction
 from .types import Action, Observation
 
+# MCP action types keyed by their "type" discriminator value.
+# These are checked before the environment's own action_cls so that
+# ListToolsAction / CallToolAction payloads are never rejected by an
+# unrelated Pydantic model.
+_MCP_ACTION_TYPES: Dict[str, Type[Action]] = {
+    "list_tools": ListToolsAction,
+    "call_tool": CallToolAction,
+}
+
 
 def deserialize_action(action_data: Dict[str, Any], action_cls: Type[Action]) -> Action:
     """
     Convert JSON dict to Action instance using Pydantic validation.
 
-    This is a basic deserialization that works for most environments.
+    MCP action types (``list_tools``, ``call_tool``) are recognised
+    automatically via the ``"type"`` discriminator field, regardless of
+    the environment's configured ``action_cls``.  All other payloads
+    fall through to ``action_cls.model_validate()``.
+
     For special cases (e.g., tensor fields, custom type conversions),
     use deserialize_action_with_preprocessing().
 
@@ -38,6 +52,17 @@ def deserialize_action(action_data: Dict[str, Any], action_cls: Type[Action]) ->
     Note:
         This uses Pydantic's model_validate() for automatic validation.
     """
+    # Route MCP action types before falling through to the env action_cls.
+    # Only intercept when action_cls is the generic Action base or itself an
+    # MCP type (i.e. the server hosts an MCP environment).  This avoids
+    # silently bypassing env-specific validation for non-MCP environments
+    # that happen to use "call_tool" / "list_tools" as a type discriminator.
+    action_type = action_data.get("type")
+    if action_type in _MCP_ACTION_TYPES:
+        mcp_cls = _MCP_ACTION_TYPES[action_type]
+        if action_cls is Action or action_cls in _MCP_ACTION_TYPES.values():
+            return mcp_cls.model_validate(action_data)
+
     return action_cls.model_validate(action_data)
 
 
@@ -62,6 +87,15 @@ def deserialize_action_with_preprocessing(
     Raises:
         ValidationError: If action_data is invalid for the action class
     """
+    # Route MCP action types before preprocessing (they don't need it).
+    # Same guard as deserialize_action: only intercept when action_cls is
+    # the generic Action base or itself an MCP type.
+    action_type = action_data.get("type")
+    if action_type in _MCP_ACTION_TYPES:
+        mcp_cls = _MCP_ACTION_TYPES[action_type]
+        if action_cls is Action or action_cls in _MCP_ACTION_TYPES.values():
+            return mcp_cls.model_validate(action_data)
+
     processed_data = {}
 
     for key, value in action_data.items():

src/openenv/core/env_server/web_interface.py

@@ -15,13 +15,15 @@ option (e.g. openenv push --enable-interface) or ENABLE_WEB_INTERFACE env var.
 from __future__ import annotations
 
 import asyncio
+import inspect
 import json
 from concurrent.futures import ThreadPoolExecutor
 from datetime import datetime
 from typing import Any, Callable, Dict, List, Optional, Type
 
 import gradio as gr
-from fastapi import FastAPI, WebSocket, WebSocketDisconnect
+from fastapi import Body, FastAPI, HTTPException, status, WebSocket, WebSocketDisconnect
+from fastapi.responses import RedirectResponse
 from pydantic import BaseModel, ConfigDict, Field
 
 from .gradio_theme import OPENENV_GRADIO_CSS, OPENENV_GRADIO_THEME
@@ -269,6 +271,28 @@ class WebInterfaceManager:
         # Thread pool for running sync code (e.g., Playwright sync API) in async context
         self._executor = ThreadPoolExecutor(max_workers=1)
 
+    @staticmethod
+    def _get_valid_kwargs(
+        sig: inspect.Signature,
+        kwargs: Dict[str, Any],
+        skip_params: Optional[set[str]] = None,
+    ) -> Dict[str, Any]:
+        """Filter kwargs to only those accepted by the target function."""
+        skip_params = skip_params or set()
+        valid_kwargs: Dict[str, Any] = {}
+        has_var_kwargs = any(
+            param.kind == inspect.Parameter.VAR_KEYWORD
+            for param in sig.parameters.values()
+        )
+
+        for key, value in kwargs.items():
+            if key in skip_params:
+                continue
+            if key in sig.parameters or has_var_kwargs:
+                valid_kwargs[key] = value
+
+        return valid_kwargs
+
     async def _run_sync_in_thread_pool(self, func, *args, **kwargs):
         """Run a synchronous function in the thread pool executor.
 
@@ -317,11 +341,24 @@ class WebInterfaceManager:
         for client in disconnected_clients:
             self.connected_clients.remove(client)
 
-    async def reset_environment(self) -> Dict[str, Any]:
+    async def reset_environment(
+        self, reset_kwargs: Optional[Dict[str, Any]] = None
+    ) -> Dict[str, Any]:
         """Reset the environment and update state."""
-        # Run sync reset in thread pool to avoid blocking event loop
-        # and to support environments using sync libraries (e.g., Playwright)
-        observation: Observation = await self._run_sync_in_thread_pool(self.env.reset)
+        reset_kwargs = reset_kwargs or {}
+
+        is_async = self.env.reset_async.__func__ is not Environment.reset_async
+        sig = inspect.signature(self.env.reset_async if is_async else self.env.reset)
+        valid_kwargs = self._get_valid_kwargs(sig, reset_kwargs)
+
+        if is_async:
+            observation = await self.env.reset_async(**valid_kwargs)
+        else:
+            # Run sync reset in thread pool to avoid blocking event loop
+            # and to support environments using sync libraries (e.g., Playwright)
+            observation = await self._run_sync_in_thread_pool(
+                self.env.reset, **valid_kwargs
+            )
         state: State = self.env.state
 
         # Serialize observation once using shared utility
@@ -428,6 +465,16 @@ def create_web_interface_app(
     web_manager = WebInterfaceManager(env, action_cls, observation_cls, metadata)
 
     # Web API routes first (so they take precedence over Gradio mount at /web)
+    @app.get("/", include_in_schema=False)
+    async def web_root():
+        """Redirect the app root to the Gradio interface."""
+        return RedirectResponse(url="/web/")
+
+    @app.get("/web", include_in_schema=False)
+    async def web_root_no_slash():
+        """Redirect /web to /web/ for mounted Gradio deployments behind proxies."""
+        return RedirectResponse(url="/web/")
+
     @app.get("/web/metadata")
     async def web_metadata():
         """Get environment metadata."""
@@ -449,9 +496,9 @@ def create_web_interface_app(
             await web_manager.disconnect_websocket(websocket)
 
     @app.post("/web/reset")
-    async def web_reset():
+    async def web_reset(request: Optional[Dict[str, Any]] = Body(default=None)):
         """Reset endpoint for web interface."""
-        return await web_manager.reset_environment()
+        return await web_manager.reset_environment(request)
 
     @app.post("/web/step")
     async def web_step(request: Dict[str, Any]):
@@ -475,7 +522,13 @@ def create_web_interface_app(
     @app.get("/web/state")
     async def web_state():
         """State endpoint for web interface."""
-        return web_manager.get_state()
+        try:
+            return web_manager.get_state()
+        except RuntimeError as exc:
+            raise HTTPException(
+                status_code=status.HTTP_409_CONFLICT,
+                detail=str(exc),
+            ) from exc
 
     action_fields = _extract_action_fields(action_cls)
     is_chat_env = _is_chat_env(action_cls)
@@ -505,7 +558,7 @@ def create_web_interface_app(
             )
         gradio_blocks = gr.TabbedInterface(
             [default_blocks, custom_blocks],
-            tab_names=["Playground", "Visualization"],
+            tab_names=["Playground", "Custom"],
             title=get_gradio_display_title(metadata),
         )
     else:

src/openenv/core/mcp_client.py

@@ -52,6 +52,7 @@ Example (sync wrapper):
     ...     result = env.call_tool("echo_message", message="Hello!")
 """
 
+import asyncio
 from typing import Any, Dict, List, Optional
 
 from .client_types import StepResult
@@ -118,6 +119,66 @@ class MCPClientBase(EnvClient[Any, Observation, State]):
         )
         self._tools_cache: Optional[List[Tool]] = None
         self.use_production_mode = False
+        self._production_session_id: Optional[str] = None
+        self._production_session_lock = asyncio.Lock()
+        self._jsonrpc_request_id = 0
+        self._http_client: Optional[Any] = None  # lazily-created httpx.AsyncClient
+
+    def _next_request_id(self) -> int:
+        """Generate a monotonically increasing JSON-RPC request id."""
+        self._jsonrpc_request_id += 1
+        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"
+
+    async def _get_http_client(self) -> Any:
+        """Return a shared httpx.AsyncClient, creating one lazily."""
+        if self._http_client is None:
+            import httpx
+
+            self._http_client = httpx.AsyncClient()
+        return self._http_client
+
+    async def _production_mcp_request(
+        self, method: str, params: Optional[Dict[str, Any]] = None
+    ) -> Dict[str, Any]:
+        """Send a JSON-RPC request to HTTP /mcp and return parsed JSON response."""
+        client = await self._get_http_client()
+        response = await client.post(
+            self._production_mcp_url(),
+            json={
+                "jsonrpc": "2.0",
+                "method": method,
+                "params": params or {},
+                "id": self._next_request_id(),
+            },
+            timeout=self._message_timeout,
+        )
+        response.raise_for_status()
+        return response.json()
+
+    async def _ensure_production_session(self) -> str:
+        """Create and cache a persistent HTTP MCP session id if needed."""
+        async with self._production_session_lock:
+            if self._production_session_id is not None:
+                return self._production_session_id
+
+            data = await self._production_mcp_request("openenv/session/create")
+            if "error" in data:
+                message = data.get("error", {}).get("message", "unknown error")
+                raise RuntimeError(f"Failed to create MCP session: {message}")
+
+            session_id = data.get("result", {}).get("session_id")
+            if not session_id:
+                raise RuntimeError("Failed to create MCP session: missing session_id")
+
+            self._production_session_id = session_id
+            return session_id
 
     async def list_tools(self, use_cache: bool = True) -> List[Tool]:
         """
@@ -138,26 +199,18 @@ class MCPClientBase(EnvClient[Any, Observation, State]):
         if use_cache and self._tools_cache is not None:
             return self._tools_cache
 
-        # Use production mode HTTP endpoint if enabled
-        if self.use_production_mode:
-            import requests
-
-            # Convert ws:// URL to http:// URL
-            url = self._ws_url.replace("ws://", "http://").replace("wss://", "https://")
-            # Remove /ws suffix if present and add /mcp
-            url = url.rstrip("/ws").rstrip("/") + "/mcp"
-
+        # Use production mode HTTP endpoint if enabled.
+        # Some tests instantiate with __new__ and skip __init__, so default missing flag to False.
+        if getattr(self, "use_production_mode", False):
             try:
-                response = requests.post(
-                    url,
-                    json={
-                        "jsonrpc": "2.0",
-                        "method": "tools/list",
-                        "params": {},
-                        "id": 1,
-                    },
+                session_id = await self._ensure_production_session()
+                data = await self._production_mcp_request(
+                    "tools/list",
+                    {"session_id": session_id},
                 )
-                data = response.json()
+                if "error" in data:
+                    message = data.get("error", {}).get("message", "unknown error")
+                    raise RuntimeError(f"list_tools failed: {message}")
                 if "result" in data and "tools" in data["result"]:
                     tools = [
                         Tool(
@@ -177,7 +230,12 @@ class MCPClientBase(EnvClient[Any, Observation, State]):
             return []
 
         result = await self.step(ListToolsAction())
-        self._tools_cache = result.observation.tools
+        if isinstance(result.observation, ListToolsObservation):
+            self._tools_cache = result.observation.tools
+            return self._tools_cache
+
+        # Unexpected observation type; keep API stable with an empty tool list.
+        self._tools_cache = []
         return self._tools_cache
 
     def _step_payload(self, action: Any) -> Dict[str, Any]:
@@ -251,6 +309,35 @@ class MCPClientBase(EnvClient[Any, Observation, State]):
             step_count=payload.get("step_count", 0),
         )
 
+    async def close(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.
+        """
+        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
+
+        if self._http_client is not None:
+            try:
+                await self._http_client.aclose()
+            except Exception:
+                pass
+            finally:
+                self._http_client = None
+
+        await super().close()
+
 
 class MCPToolClient(MCPClientBase):
     """
@@ -316,6 +403,26 @@ class MCPToolClient(MCPClientBase):
             >>> result = await env.call_tool("greet", name="Claude")
             >>> print(result)  # "Hello, Claude!"
         """
+        if getattr(self, "use_production_mode", False):
+            session_id = await self._ensure_production_session()
+            data = await self._production_mcp_request(
+                "tools/call",
+                {
+                    "name": name,
+                    "arguments": kwargs,
+                    "session_id": session_id,
+                },
+            )
+
+            if "error" in data:
+                message = data.get("error", {}).get("message", "unknown error")
+                raise RuntimeError(f"Tool '{name}' failed: {message}")
+
+            result = data.get("result")
+            if isinstance(result, dict) and "data" in result:
+                return result["data"]
+            return result
+
         action = CallToolAction(tool_name=name, arguments=kwargs)
         result = await self.step(action)
         obs = result.observation

src/openenv_core.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: openenv-core
-Version: 0.2.2.dev0
+Version: 0.2.3
 Summary: A unified framework for reinforcement learning environments
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
@@ -19,6 +19,7 @@ Requires-Dist: tomli-w>=1.2.0
 Requires-Dist: websockets>=15.0.1
 Requires-Dist: fastmcp>=3.0.0
 Requires-Dist: gradio>=4.0.0
+Requires-Dist: httpx>=0.28.1
 Provides-Extra: core
 Requires-Dist: fastapi>=0.104.0; extra == "core"
 Requires-Dist: pydantic>=2.0.0; extra == "core"
@@ -61,7 +62,7 @@ Dynamic: license-file
 
 An e2e framework for creating, deploying and using isolated execution environments for agentic RL training, built using Gymnasium style simple APIs.
 
-[![PyPI](https://img.shields.io/pypi/v/openenv?color=blue)](https://pypi.org/project/openenv/)
+[![PyPI](https://img.shields.io/pypi/v/openenv-core?color=blue)](https://pypi.org/project/openenv-core/)
 [![Discord](https://img.shields.io/badge/Discord-OpenEnv-7289da?style=flat&logo=discord&logoColor=white)](https://discord.gg/YsTYBh6PD9)
 [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/meta-pytorch/OpenEnv/blob/main/examples/OpenEnv_Tutorial.ipynb)
 [![Docs](https://img.shields.io/badge/Docs-Explore-blue?logo=readthedocs&logoColor=white)](https://meta-pytorch.org/OpenEnv/)

src/openenv_core.egg-info/SOURCES.txt

@@ -1,4 +1,5 @@
 LICENSE
+MANIFEST.in
 README.md
 pyproject.toml
 src/openenv/__init__.py
@@ -19,8 +20,6 @@ src/openenv/cli/commands/serve.py
 src/openenv/cli/commands/skills.py
 src/openenv/cli/commands/validate.py
 src/openenv/cli/templates/__init__.py
-src/openenv/cli/templates/__pycache__/__init__.cpython-311.pyc
-src/openenv/cli/templates/__pycache__/__init__.cpython-313.pyc
 src/openenv/cli/templates/openenv_env/README.md
 src/openenv/cli/templates/openenv_env/__init__.py
 src/openenv/cli/templates/openenv_env/client.py

src/openenv_core.egg-info/requires.txt

@@ -12,6 +12,7 @@ tomli-w>=1.2.0
 websockets>=15.0.1
 fastmcp>=3.0.0
 gradio>=4.0.0
+httpx>=0.28.1
 
 [all]
 openenv-core[core]