openenv/coding_env
21
1# Copyright (c) Meta Platforms, Inc. and affiliates.2# All rights reserved.3#4# This source code is licensed under the BSD-style license found in the5# LICENSE file in the root directory of this source tree.6 7"""8HTTP server wrapper for Environment instances.9 10This module provides utilities to wrap any Environment subclass and expose it11over HTTP and WebSocket endpoints that EnvClient can consume.12"""13 14from __future__ import annotations15 16import asyncio17import inspect18import json19import logging20import os21import time22import uuid23from concurrent.futures import ThreadPoolExecutor24from contextlib import AsyncExitStack25from typing import Any, AsyncContextManager, Callable, cast, Dict, Optional, Type26 27_MISSING = object()28 29from fastapi import (30 Body,31 FastAPI,32 HTTPException,33 Request,34 status,35 WebSocket,36 WebSocketDisconnect,37)38from pydantic import ValidationError39 40from .interfaces import Environment41from .mcp_environment import get_server_tools42from .mcp_types import (43 JsonRpcErrorCode,44 JsonRpcRequest,45 JsonRpcResponse,46 McpMethod,47 WSMCPMessage,48 WSMCPResponse,49)50from .route_config import GetEndpointConfig, register_get_endpoints51from .serialization import deserialize_action, serialize_observation52from .types import (53 Action,54 ConcurrencyConfig,55 EnvironmentMetadata,56 HealthResponse,57 HealthStatus,58 Observation,59 ResetRequest,60 ResetResponse,61 SchemaResponse,62 ServerCapacityStatus,63 ServerMode,64 SessionInfo,65 State,66 StepRequest,67 StepResponse,68 WSCloseMessage,69 WSErrorCode,70 WSErrorResponse,71 WSObservationResponse,72 WSResetMessage,73 WSStateMessage,74 WSStateResponse,75 WSStepMessage,76)77 78 79def _make_json_serializable(obj: Any) -> Any:80 """81 Convert an object to a JSON-serializable form.82 83 Handles Pydantic models, dataclasses, and other common types.84 85 Args:86 obj: The object to convert87 88 Returns:89 A JSON-serializable representation of the object90 """91 if obj is None:92 return None93 if isinstance(obj, (str, int, float, bool)):94 return obj95 if isinstance(obj, (list, tuple)):96 return [_make_json_serializable(item) for item in obj]97 if isinstance(obj, dict):98 return {k: _make_json_serializable(v) for k, v in obj.items()}99 if hasattr(obj, "model_dump"):100 # Pydantic model101 return obj.model_dump()102 if hasattr(obj, "__dict__"):103 # Object with __dict__104 return {k: _make_json_serializable(v) for k, v in obj.__dict__.items()}105 # Fallback to string representation106 return str(obj)107 108 109from .exceptions import (110 ConcurrencyConfigurationError,111 EnvironmentFactoryError,112 SessionCapacityError,113)114 115 116class HTTPEnvServer:117 """118 HTTP server wrapper for Environment instances.119 120 This class wraps an Environment and exposes its reset(), step(), and state121 methods as HTTP and WebSocket endpoints compatible with EnvClient.122 123 The server expects:124 - Action deserialization: Converts JSON dict to Action subclass125 - Observation serialization: Converts Observation subclass to JSON dict126 127 Example:128 >>> from core.env_server import HTTPEnvServer129 >>> from envs.coding_env.server import CodeExecutionEnvironment130 >>> from envs.coding_env.models import CodeAction, CodeObservation131 >>>132 >>> # Pass environment class (factory pattern)133 >>> server = HTTPEnvServer(134 ... env=CodeExecutionEnvironment,135 ... action_cls=CodeAction,136 ... observation_cls=CodeObservation,137 ... max_concurrent_envs=4,138 ... )139 >>>140 >>> # Register routes with FastAPI141 >>> from fastapi import FastAPI142 >>> app = FastAPI()143 >>> server.register_routes(app)144 """145 146 def __init__(147 self,148 env: Callable[[], Environment],149 action_cls: Type[Action],150 observation_cls: Type[Observation],151 max_concurrent_envs: Optional[int] = None,152 concurrency_config: Optional[ConcurrencyConfig] = None,153 ):154 """155 Initialize HTTP server wrapper.156 157 Args:158 env: Environment factory (callable) that creates new instances.159 Will be called to create a new environment for each WebSocket session.160 action_cls: The Action subclass this environment expects161 observation_cls: The Observation subclass this environment returns162 max_concurrent_envs: Maximum number of concurrent WebSocket sessions.163 Mutually exclusive with concurrency_config.164 concurrency_config: Optional ConcurrencyConfig for advanced concurrency settings.165 Mutually exclusive with max_concurrent_envs.166 167 Raises:168 ValueError: If both max_concurrent_envs and concurrency_config are provided.169 ConcurrencyConfigurationError: If max_concurrent_envs > 1 for an170 environment that is not marked as SUPPORTS_CONCURRENT_SESSIONS.171 """172 # Validate that env is callable173 if not callable(env):174 raise TypeError(175 f"env must be a callable (class or factory function), got {type(env)}. "176 f"Pass the environment class (e.g., MyEnvironment) not an instance (e.g., MyEnvironment())."177 )178 179 self._env_factory: Callable[[], Environment] = env180 181 # Handle concurrency configuration182 if max_concurrent_envs is not None and concurrency_config is not None:183 raise ValueError(184 "Cannot specify both 'max_concurrent_envs' and 'concurrency_config'. "185 "Please use only one method to configure concurrency."186 )187 188 if concurrency_config is not None:189 self._concurrency_config = concurrency_config190 elif max_concurrent_envs is not None:191 self._concurrency_config = ConcurrencyConfig(192 max_concurrent_envs=max_concurrent_envs,193 session_timeout=None,194 )195 else:196 # Default configuration197 self._concurrency_config = ConcurrencyConfig(198 max_concurrent_envs=1,199 session_timeout=None,200 )201 202 self._max_concurrent_envs = self._concurrency_config.max_concurrent_envs203 204 # Validate concurrency configuration205 self._validate_concurrency_safety()206 207 self.action_cls = action_cls208 self.observation_cls = observation_cls209 210 # Session management for WebSocket connections211 self._sessions: Dict[str, Optional[Environment]] = {}212 self._session_executors: Dict[str, ThreadPoolExecutor] = {}213 self._session_stacks: Dict[str, AsyncExitStack] = {}214 self._session_info: Dict[str, SessionInfo] = {}215 self._session_lock = asyncio.Lock()216 217 # Create thread pool for running sync code in async context218 # This is needed for environments using sync libraries (e.g., Playwright)219 self._executor = ThreadPoolExecutor(max_workers=32)220 221 # Idle session reaper configuration.222 # Timeout is taken from ConcurrencyConfig.session_timeout;223 # None means no timeout (default — reaper is a no-op).224 self._session_idle_timeout_s: Optional[float] = (225 self._concurrency_config.session_timeout226 )227 self._reaper_task: Optional[asyncio.Task[None]] = None228 229 def _validate_concurrency_safety(self) -> None:230 """231 Validate that the environment supports the configured concurrency level.232 233 Raises:234 ConcurrencyConfigurationError: If max_concurrent_envs > 1 for an235 environment that is not marked as SUPPORTS_CONCURRENT_SESSIONS.236 """237 if self._max_concurrent_envs <= 1:238 return239 240 if inspect.isclass(self._env_factory):241 env_cls = self._env_factory242 else:243 _temp_env = self._env_factory()244 env_cls = type(_temp_env)245 _temp_env.close()246 del _temp_env247 248 if not getattr(env_cls, "SUPPORTS_CONCURRENT_SESSIONS", False):249 raise ConcurrencyConfigurationError(250 environment_name=env_cls.__name__,251 max_concurrent_envs=self._max_concurrent_envs,252 )253 254 def get_capacity_status(self) -> ServerCapacityStatus:255 """256 Get the current capacity status of the server.257 258 Returns:259 ServerCapacityStatus with current session counts and availability.260 """261 return ServerCapacityStatus.from_counts(262 active=len(self._sessions),263 max_sessions=self._max_concurrent_envs,264 )265 266 async def _run_sync_in_thread_pool(267 self, func: Callable[..., Observation], *args, **kwargs268 ) -> Observation:269 """Run a synchronous function in the thread pool executor."""270 loop = asyncio.get_event_loop()271 return await loop.run_in_executor(self._executor, lambda: func(*args, **kwargs))272 273 def _get_valid_kwargs(274 self,275 sig: inspect.Signature,276 kwargs: Dict[str, Any],277 skip_params: Optional[set[str]] = None,278 ) -> Dict[str, Any]:279 """Filter kwargs to only include parameters accepted by the function signature."""280 if skip_params is None:281 skip_params = set()282 283 valid_kwargs = {}284 285 has_kwargs = any(286 p.kind == inspect.Parameter.VAR_KEYWORD for p in sig.parameters.values()287 )288 289 for k, v in kwargs.items():290 if k in sig.parameters or has_kwargs:291 if k not in skip_params:292 valid_kwargs[k] = v293 294 return valid_kwargs295 296 async def _create_session(self) -> tuple[str, Environment]:297 """298 Create a new WebSocket session with its own environment instance.299 300 Returns:301 Tuple of (session_id, environment)302 303 Raises:304 SessionCapacityError: If max concurrent sessions reached305 EnvironmentFactoryError: If the factory fails to create an environment306 """307 async with self._session_lock:308 if len(self._sessions) >= self._max_concurrent_envs:309 raise SessionCapacityError(310 active_sessions=len(self._sessions),311 max_sessions=self._max_concurrent_envs,312 )313 314 session_id = str(uuid.uuid4())315 current_time = time.time()316 317 # Create executor and reserve slot so capacity is not exceeded while318 # we create the env outside the lock (avoids blocking other sessions)319 executor = ThreadPoolExecutor(max_workers=1)320 self._session_executors[session_id] = executor321 self._sessions[session_id] = None # placeholder until env is ready322 323 try:324 # Create environment in the executor thread (outside lock)325 loop = asyncio.get_event_loop()326 env = await loop.run_in_executor(executor, self._env_factory)327 except Exception as e:328 async with self._session_lock:329 executor.shutdown(wait=False)330 self._session_executors.pop(session_id, None)331 self._sessions.pop(session_id, None)332 factory_name = getattr(333 self._env_factory, "__name__", str(self._env_factory)334 )335 raise EnvironmentFactoryError(factory_name) from e336 337 # Hold the MCP session open for the lifetime of this session,338 # matching the WebSocket path's AsyncExitStack pattern. This339 # prevents per-request MCP transport teardown/reconnection and340 # preserves FastMCP session state (ctx.set_state / ctx.get_state)341 # across HTTP calls within the same OpenEnv session.342 stack = AsyncExitStack()343 try:344 mcp_session_factory = getattr(env, "mcp_session", None)345 if callable(mcp_session_factory):346 mcp_session_cm = cast(AsyncContextManager[Any], mcp_session_factory())347 await stack.enter_async_context(mcp_session_cm)348 except Exception:349 # MCP transport failed to start — clean up the reserved slot,350 # the env, and the executor so they don't leak permanently351 # against _max_concurrent_envs.352 await stack.aclose() # best-effort353 async with self._session_lock:354 self._sessions.pop(session_id, None)355 self._session_executors.pop(session_id, None)356 self._session_info.pop(session_id, None)357 await self._cleanup_session_resources(env, executor)358 raise359 360 async with self._session_lock:361 self._sessions[session_id] = env362 self._session_stacks[session_id] = stack363 now = time.time()364 self._session_info[session_id] = SessionInfo(365 session_id=session_id,366 created_at=current_time,367 last_activity_at=now,368 step_count=0,369 environment_type=type(env).__name__,370 )371 372 return session_id, env373 374 async def _destroy_session(self, session_id: str) -> None:375 """376 Destroy a WebSocket session and cleanup resources.377 378 Args:379 session_id: The session ID to destroy380 """381 async with self._session_lock:382 env = self._sessions.pop(session_id, None)383 executor = self._session_executors.pop(session_id, None)384 stack = self._session_stacks.pop(session_id, None)385 self._session_info.pop(session_id, None)386 387 await self._cleanup_session_resources(env, executor, stack)388 389 async def _cleanup_session_resources(390 self,391 env: Optional[Environment],392 executor: Optional[ThreadPoolExecutor],393 stack: Optional[AsyncExitStack] = None,394 ) -> None:395 """Close an environment and shut down its executor (best-effort)."""396 # Close the MCP session stack first — this gracefully exits the397 # mcp_session() context (and the underlying FastMCP Client session)398 # before we tear down the environment references.399 if stack is not None:400 try:401 await stack.aclose()402 except Exception:403 pass # Best effort cleanup404 405 # Run close() in the same executor where the env was created406 # This is required for thread-sensitive libraries like Playwright/greenlet407 if env is not None:408 if executor is not None:409 try:410 loop = asyncio.get_event_loop()411 await loop.run_in_executor(executor, env.close)412 except Exception:413 # If executor close fails, try direct close as fallback414 try:415 env.close()416 except Exception:417 pass # Best effort cleanup418 else:419 try:420 env.close()421 except Exception:422 pass # Best effort cleanup423 424 # Shutdown executor after close is done425 if executor is not None:426 executor.shutdown(wait=False)427 428 def _update_session_activity(429 self, session_id: str, increment_step: bool = False430 ) -> None:431 """432 Update session activity timestamp and optionally increment step count.433 434 Args:435 session_id: The session ID to update436 increment_step: If True, increment the step count437 """438 if session_id in self._session_info:439 self._session_info[session_id].last_activity_at = time.time()440 if increment_step:441 self._session_info[session_id].step_count += 1442 443 async def _reap_idle_sessions(self) -> None:444 """Background task that periodically destroys sessions idle beyond the timeout."""445 timeout = self._session_idle_timeout_s446 if timeout is None:447 return # no timeout configured — noop448 interval = max(timeout / 4, 5.0) # check frequently enough449 while True:450 try:451 await asyncio.sleep(interval)452 now = time.time()453 stale_ids: list[str] = []454 async with self._session_lock:455 for sid, info in self._session_info.items():456 if now - info.last_activity_at > timeout:457 stale_ids.append(sid)458 for sid in stale_ids:459 # Re-check under lock: activity may have arrived since460 # the snapshot was taken, making this session active again.461 # Refresh `now` so slow _destroy_session calls don't cause462 # subsequent entries to be validated against a stale clock.463 now = time.time()464 async with self._session_lock:465 info = self._session_info.get(sid)466 if info is None or (now - info.last_activity_at) <= timeout:467 continue468 await self._destroy_session(sid)469 except asyncio.CancelledError:470 break471 except Exception as exc:472 logging.getLogger(__name__).warning(473 "Idle-session reaper encountered an error (will retry): %s",474 exc,475 )476 477 def _start_reaper(self) -> None:478 """Start the idle-session reaper if a timeout is configured."""479 if self._session_idle_timeout_s is not None and self._reaper_task is None:480 self._reaper_task = asyncio.create_task(self._reap_idle_sessions())481 482 def _stop_reaper(self) -> None:483 """Cancel the reaper background task."""484 if self._reaper_task is not None:485 self._reaper_task.cancel()486 self._reaper_task = None487 488 def get_session_info(self, session_id: str) -> Optional[SessionInfo]:489 """490 Get information about a specific session.491 492 Args:493 session_id: The session ID to query494 495 Returns:496 SessionInfo if the session exists, None otherwise497 """498 return self._session_info.get(session_id)499 500 async def _run_in_session_executor(501 self, session_id: str, func: Callable[..., Observation], *args, **kwargs502 ) -> Observation:503 """Run a synchronous function in the session's thread pool executor."""504 executor = self._session_executors.get(session_id, self._executor)505 loop = asyncio.get_event_loop()506 return await loop.run_in_executor(executor, lambda: func(*args, **kwargs))507 508 @property509 def active_sessions(self) -> int:510 """Return the number of active WebSocket sessions."""511 return len(self._sessions)512 513 @property514 def max_concurrent_envs(self) -> int:515 """Return the maximum number of concurrent environments."""516 return self._max_concurrent_envs517 518 @property519 def is_concurrency_safe(self) -> bool:520 """Return whether the environment is marked as concurrency safe."""521 import inspect522 523 if inspect.isclass(self._env_factory):524 return getattr(self._env_factory, "SUPPORTS_CONCURRENT_SESSIONS", False)525 else:526 _temp_env = self._env_factory()527 result = getattr(_temp_env, "SUPPORTS_CONCURRENT_SESSIONS", False)528 _temp_env.close()529 del _temp_env530 return result531 532 @property533 def concurrency_config(self) -> ConcurrencyConfig:534 """Return the concurrency configuration."""535 return self._concurrency_config536 537 def register_routes(538 self, app: FastAPI, mode: ServerMode | str = ServerMode.SIMULATION539 ) -> None:540 """541 Register HTTP routes on a FastAPI application.542 543 Args:544 app: FastAPI application instance545 mode: Server mode - either SIMULATION or PRODUCTION (or string equivalents).546 In production mode, simulation control endpoints (/reset, /step, /state)547 are NOT registered. Only safe endpoints (/health, /schema, /metadata, /ws)548 are available. Defaults to SIMULATION for backwards compatibility.549 550 Raises:551 ValueError: If mode is not a valid ServerMode or string equivalent.552 """553 # Convert string to ServerMode enum for backwards compatibility554 if isinstance(mode, str):555 try:556 mode = ServerMode(mode.lower())557 except ValueError:558 valid_modes = [m.value for m in ServerMode]559 raise ValueError(560 f"Invalid mode: '{mode}'. Must be one of: {valid_modes}"561 )562 563 # Wire up idle-session reaper lifecycle via app events564 server_ref = self565 566 async def _start_session_reaper() -> None:567 server_ref._start_reaper()568 569 async def _stop_session_reaper() -> None:570 server_ref._stop_reaper()571 572 if not getattr(app.router, "_openenv_reaper_registered", False):573 app.router.on_startup.append(_start_session_reaper)574 app.router.on_shutdown.append(_stop_session_reaper)575 app.router._openenv_reaper_registered = True # type: ignore[attr-defined]576 577 # Helper function to handle reset endpoint578 async def reset_handler(579 request: ResetRequest = Body(default_factory=ResetRequest),580 ) -> ResetResponse:581 """Reset endpoint - returns initial observation."""582 _env = self._env_factory()583 584 try:585 kwargs = request.model_dump(exclude_unset=True)586 587 is_async = _env.reset_async.__func__ is not Environment.reset_async588 589 if is_async:590 sig = inspect.signature(_env.reset_async)591 else:592 sig = inspect.signature(_env.reset)593 valid_kwargs = self._get_valid_kwargs(sig, kwargs)594 595 if is_async:596 observation = await _env.reset_async(**valid_kwargs)597 else:598 observation = await self._run_sync_in_thread_pool(599 _env.reset, **valid_kwargs600 )601 return ResetResponse(**serialize_observation(observation))602 finally:603 _env.close()604 605 # Helper function to handle step endpoint606 async def step_handler(request: StepRequest) -> StepResponse:607 """Step endpoint - executes action and returns observation."""608 action_data = request.action609 610 try:611 action = deserialize_action(action_data, self.action_cls)612 except ValidationError as e:613 raise HTTPException(614 status_code=status.HTTP_422_UNPROCESSABLE_CONTENT, detail=e.errors()615 )616 617 _env = self._env_factory()618 619 try:620 kwargs = request.model_dump(exclude_unset=True, exclude={"action"})621 622 is_async = _env.step_async.__func__ is not Environment.step_async623 624 if is_async:625 sig = inspect.signature(_env.step_async)626 else:627 sig = inspect.signature(_env.step)628 valid_kwargs = self._get_valid_kwargs(629 sig, kwargs, skip_params={"action"}630 )631 632 if is_async:633 observation = await _env.step_async(action, **valid_kwargs)634 else:635 observation = await self._run_sync_in_thread_pool(636 _env.step, action, **valid_kwargs637 )638 639 return StepResponse(**serialize_observation(observation))640 finally:641 _env.close()642 643 # Helper function to handle MCP endpoint644 async def mcp_handler(645 request: JsonRpcRequest,646 session_env: Optional[Environment] = None,647 session_id: Optional[str] = None,648 ) -> JsonRpcResponse:649 """650 Handle MCP JSON-RPC requests.651 652 Supports tools/list and tools/call methods in JSON-RPC 2.0 format,653 plus OpenEnv session lifecycle methods for HTTP MCP:654 - openenv/session/create655 - openenv/session/close656 """657 method = request.method658 request_id = request.id659 params = request.params660 if not isinstance(params, dict):661 return JsonRpcResponse.error_response(662 JsonRpcErrorCode.INVALID_PARAMS,663 "Params must be an object",664 request_id=request_id,665 )666 667 # OpenEnv extension methods for explicit MCP session management.668 # This enables persistent MCP lifecycles over HTTP /mcp, matching WebSocket semantics.669 if method == "openenv/session/create":670 if session_env is not None and session_id is not None:671 return JsonRpcResponse.success(672 result={"session_id": session_id},673 request_id=request_id,674 )675 try:676 created_session_id, _ = await self._create_session()677 except SessionCapacityError as e:678 return JsonRpcResponse.error_response(679 JsonRpcErrorCode.SERVER_ERROR,680 str(e),681 request_id=request_id,682 data={683 "active_sessions": e.active_sessions,684 "max_sessions": e.max_sessions,685 },686 )687 except EnvironmentFactoryError as e:688 return JsonRpcResponse.error_response(689 JsonRpcErrorCode.SERVER_ERROR,690 str(e),691 request_id=request_id,692 data={"factory_name": e.factory_name},693 )694 return JsonRpcResponse.success(695 result={"session_id": created_session_id},696 request_id=request_id,697 )698 699 if method == "openenv/session/close":700 target_session_id = params.get("session_id")701 if not target_session_id:702 return JsonRpcResponse.error_response(703 JsonRpcErrorCode.INVALID_PARAMS,704 "Invalid params - 'session_id' is required",705 request_id=request_id,706 )707 708 if session_id is not None and target_session_id == session_id:709 return JsonRpcResponse.error_response(710 JsonRpcErrorCode.INVALID_REQUEST,711 "Cannot close active WebSocket-managed session via MCP method",712 request_id=request_id,713 )714 715 async with self._session_lock:716 env = self._sessions.pop(target_session_id, _MISSING)717 if env is not _MISSING:718 executor = self._session_executors.pop(target_session_id, None)719 stack = self._session_stacks.pop(target_session_id, None)720 self._session_info.pop(target_session_id, None)721 else:722 executor = None723 stack = None724 725 if env is _MISSING:726 return JsonRpcResponse.error_response(727 JsonRpcErrorCode.INVALID_PARAMS,728 f"Unknown session_id: {target_session_id}",729 request_id=request_id,730 )731 732 if env is None:733 # Session slot reserved but env factory still running;734 # re-insert the placeholder AND the executor so735 # _create_session can finish and the executor remains736 # tracked for eventual shutdown.737 async with self._session_lock:738 self._sessions[target_session_id] = None739 if executor is not None:740 self._session_executors[target_session_id] = executor741 return JsonRpcResponse.error_response(742 JsonRpcErrorCode.INVALID_REQUEST,743 f"Session {target_session_id} is still initializing; retry shortly",744 request_id=request_id,745 )746 747 # env/executor/stack cleanup outside the lock748 await self._cleanup_session_resources(env, executor, stack)749 return JsonRpcResponse.success(750 result={"session_id": target_session_id, "closed": True},751 request_id=request_id,752 )753 754 requested_session_id = params.get("session_id")755 managed_session_id = session_id756 757 # Use provided session environment or create temporary one758 if session_env is not None:759 _env = session_env760 should_close = False761 elif requested_session_id:762 async with self._session_lock:763 _env = self._sessions.get(requested_session_id, _MISSING)764 765 if _env is _MISSING:766 return JsonRpcResponse.error_response(767 JsonRpcErrorCode.INVALID_PARAMS,768 f"Unknown session_id: {requested_session_id}",769 request_id=request_id,770 )771 772 if _env is None:773 return JsonRpcResponse.error_response(774 JsonRpcErrorCode.INVALID_REQUEST,775 f"Session {requested_session_id} is still initializing; retry shortly",776 request_id=request_id,777 )778 779 should_close = False780 managed_session_id = requested_session_id781 else:782 _env = self._env_factory()783 should_close = True784 try:785 mcp_client = getattr(_env, "mcp_client", None)786 mcp_server = getattr(_env, "mcp_server", None)787 mcp_session_factory = getattr(_env, "mcp_session", None)788 789 if method == McpMethod.TOOLS_LIST:790 # Check if environment is MCP-enabled791 if mcp_client is None and mcp_server is None:792 return JsonRpcResponse.error_response(793 JsonRpcErrorCode.INTERNAL_ERROR,794 "Environment does not support MCP",795 request_id=request_id,796 )797 798 if mcp_client:799 if managed_session_id and mcp_client.is_connected():800 # Session-managed with live transport — call801 # directly, no redundant re-entry.802 tools = await mcp_client.list_tools()803 elif callable(mcp_session_factory):804 # Stateless request, or session-managed but the805 # background transport was lost: (re-)open.806 mcp_session_cm = cast(807 AsyncContextManager[Any], mcp_session_factory()808 )809 async with mcp_session_cm:810 tools = await mcp_client.list_tools()811 else:812 async with mcp_client:813 tools = await mcp_client.list_tools()814 815 return JsonRpcResponse.success(816 result={817 "tools": [818 t.model_dump()819 if hasattr(t, "model_dump")820 else dict(t)821 for t in tools822 ]823 },824 request_id=request_id,825 )826 827 if mcp_server:828 tools = []829 for _tool_name, tool in get_server_tools(mcp_server).items():830 tools.append(831 {832 "name": tool.name,833 "description": tool.description or "",834 "inputSchema": tool.parameters or {},835 }836 )837 return JsonRpcResponse.success(838 result={"tools": tools},839 request_id=request_id,840 )841 842 return JsonRpcResponse.error_response(843 JsonRpcErrorCode.INTERNAL_ERROR,844 "MCP server not available",845 request_id=request_id,846 )847 848 elif method == McpMethod.TOOLS_CALL:849 tool_name = params.get("name")850 arguments = params.get("arguments", {})851 852 if mcp_client is None and mcp_server is None:853 return JsonRpcResponse.error_response(854 JsonRpcErrorCode.INTERNAL_ERROR,855 "Environment does not support MCP",856 request_id=request_id,857 )858 859 if not tool_name:860 return JsonRpcResponse.error_response(861 JsonRpcErrorCode.INVALID_PARAMS,862 "Missing 'name' in params",863 request_id=request_id,864 )865 866 if mcp_client:867 if managed_session_id and mcp_client.is_connected():868 # Session-managed with live transport.869 result = await mcp_client.call_tool(870 name=tool_name, arguments=arguments871 )872 elif callable(mcp_session_factory):873 # Stateless request, or session-managed but the874 # background transport was lost: (re-)open.875 mcp_session_cm = cast(876 AsyncContextManager[Any], mcp_session_factory()877 )878 async with mcp_session_cm:879 result = await mcp_client.call_tool(880 name=tool_name, arguments=arguments881 )882 else:883 async with mcp_client:884 result = await mcp_client.call_tool(885 name=tool_name, arguments=arguments886 )887 elif mcp_server:888 server_tools = get_server_tools(mcp_server)889 if tool_name in server_tools:890 tool = server_tools[tool_name]891 if inspect.iscoroutinefunction(tool.fn):892 result = await tool.fn(**arguments)893 else:894 result = tool.fn(**arguments)895 else:896 return JsonRpcResponse.error_response(897 JsonRpcErrorCode.INVALID_PARAMS,898 f"Tool not found: {tool_name}",899 request_id=request_id,900 )901 else:902 return JsonRpcResponse.error_response(903 JsonRpcErrorCode.INTERNAL_ERROR,904 "MCP server not available",905 request_id=request_id,906 )907 908 # Ensure result is JSON serializable909 serializable_result = _make_json_serializable(result)910 911 return JsonRpcResponse.success(912 result=serializable_result,913 request_id=request_id,914 )915 916 else:917 return JsonRpcResponse.error_response(918 JsonRpcErrorCode.METHOD_NOT_FOUND,919 f"Method not found: {method}",920 request_id=request_id,921 )922 923 except Exception as e:924 return JsonRpcResponse.error_response(925 JsonRpcErrorCode.INTERNAL_ERROR,926 str(e),927 request_id=request_id,928 )929 finally:930 if managed_session_id:931 self._update_session_activity(932 managed_session_id,933 increment_step=(method == McpMethod.TOOLS_CALL),934 )935 if should_close:936 _env.close()937 938 # Register MCP WebSocket endpoint (available in both production and simulation modes)939 @app.websocket("/mcp")940 async def mcp_websocket_endpoint(websocket: WebSocket):941 """942 WebSocket endpoint for MCP JSON-RPC requests.943 944 Each WebSocket connection gets its own environment instance for MCP operations.945 946 Message Protocol:947 - Client sends: JSON-RPC 2.0 request (tools/list, tools/call)948 - Server responds: JSON-RPC 2.0 response (result or error)949 """950 await websocket.accept()951 952 session_id = None953 session_env = None954 955 try:956 # Create session with dedicated environment957 session_id, session_env = await self._create_session()958 if session_env is None:959 raise RuntimeError(960 "Session environment not initialized for MCP websocket"961 )962 963 # If environment has an mcp_session context manager, hold it open964 # for the lifetime of the websocket connection965 966 async with AsyncExitStack() as stack:967 mcp_session_factory = getattr(session_env, "mcp_session", None)968 if callable(mcp_session_factory):969 mcp_session_cm = cast(970 AsyncContextManager[Any], mcp_session_factory()971 )972 await stack.enter_async_context(mcp_session_cm)973 974 while True:975 # Receive message from client976 raw_message = await websocket.receive_text()977 978 try:979 jsonrpc_dict = json.loads(raw_message)980 jsonrpc_request = JsonRpcRequest(**jsonrpc_dict)981 except json.JSONDecodeError as e:982 error_resp = JsonRpcResponse.error_response(983 JsonRpcErrorCode.PARSE_ERROR,984 f"Parse error: {e}",985 )986 await websocket.send_text(error_resp.model_dump_json())987 continue988 except ValidationError as e:989 error_resp = JsonRpcResponse.error_response(990 JsonRpcErrorCode.INVALID_REQUEST,991 f"Invalid request: {e}",992 )993 await websocket.send_text(error_resp.model_dump_json())994 continue995 996 try:997 # Call mcp_handler with session environment998 response = await mcp_handler(999 jsonrpc_request,1000 session_env=session_env,1001 session_id=session_id,1002 )1003 await websocket.send_text(response.model_dump_json())1004 except Exception as e:1005 error_resp = JsonRpcResponse.error_response(1006 JsonRpcErrorCode.INTERNAL_ERROR,1007 str(e),1008 request_id=jsonrpc_request.id,1009 )1010 await websocket.send_text(error_resp.model_dump_json())1011 1012 except WebSocketDisconnect:1013 pass1014 except SessionCapacityError as e:1015 error_resp = JsonRpcResponse.error_response(1016 JsonRpcErrorCode.SERVER_ERROR,1017 str(e),1018 data={1019 "active_sessions": e.active_sessions,1020 "max_sessions": e.max_sessions,1021 },1022 )1023 await websocket.send_text(error_resp.model_dump_json())1024 except EnvironmentFactoryError as e:1025 error_resp = JsonRpcResponse.error_response(1026 JsonRpcErrorCode.SERVER_ERROR,1027 str(e),1028 data={"factory_name": e.factory_name},1029 )1030 await websocket.send_text(error_resp.model_dump_json())1031 except Exception as e:1032 error_resp = JsonRpcResponse.error_response(1033 JsonRpcErrorCode.SERVER_ERROR,1034 str(e),1035 )1036 await websocket.send_text(error_resp.model_dump_json())1037 finally:1038 if session_id:1039 await self._destroy_session(session_id)1040 try:1041 await websocket.close()1042 except RuntimeError:1043 pass1044 1045 # Register simulation control routes only in simulation mode1046 if mode == ServerMode.SIMULATION:1047 1048 @app.post(1049 "/reset",1050 response_model=ResetResponse,1051 tags=["Environment Control"],1052 summary="Reset the environment",1053 description="""1054Reset the environment to its initial state and return the first observation.1055 1056You can optionally provide a seed for reproducibility and an episode_id for tracking.1057 """,1058 responses={1059 200: {1060 "description": "Environment reset successfully",1061 "content": {1062 "application/json": {1063 "example": {1064 "observation": {"status": "ready", "data": {}},1065 "reward": None,1066 "done": False,1067 }1068 }1069 },1070 }1071 },1072 )1073 async def reset(1074 request: ResetRequest = Body(default_factory=ResetRequest),1075 ) -> ResetResponse:1076 return await reset_handler(request)1077 1078 @app.post(1079 "/step",1080 response_model=StepResponse,1081 tags=["Environment Control"],1082 summary="Execute an action in the environment",1083 description="""1084Execute an action in the environment and receive the resulting observation.1085 1086The action must conform to the environment's action schema, which can be1087retrieved from the `/schema` endpoint. If the action is invalid,1088the endpoint will return HTTP 422 with detailed validation errors.1089 1090The response includes:1091- **observation**: The environment's response to the action1092- **reward**: Optional reward signal (float or None)1093- **done**: Boolean indicating if the episode has terminated1094 """,1095 responses={1096 200: {1097 "description": "Action executed successfully",1098 "content": {1099 "application/json": {1100 "example": {1101 "observation": {"status": "success", "data": {}},1102 "reward": 1.0,1103 "done": False,1104 }1105 }1106 },1107 },1108 422: {1109 "description": "Validation error - invalid action format or values",1110 "content": {1111 "application/json": {1112 "example": {1113 "detail": [1114 {1115 "type": "string_too_short",1116 "loc": ["body", "action", "message"],1117 "msg": "String should have at least 1 character",1118 "input": "",1119 }1120 ]1121 }1122 }1123 },1124 },1125 500: {1126 "description": "Internal server error during action execution"1127 },1128 },1129 )1130 async def step(request: StepRequest) -> StepResponse:1131 return await step_handler(request)1132 1133 def get_state_handler() -> State:1134 _env = self._env_factory()1135 try:1136 return _env.state1137 finally:1138 _env.close()1139 1140 def get_metadata_handler() -> EnvironmentMetadata:1141 _env = self._env_factory()1142 try:1143 return _env.get_metadata()1144 finally:1145 _env.close()1146 1147 # Build list of GET endpoints based on mode1148 get_endpoints = [1149 GetEndpointConfig(1150 path="/metadata",1151 handler=get_metadata_handler,1152 response_model=EnvironmentMetadata,1153 tag="Environment Info",1154 summary="Get environment metadata",1155 description="""1156Get metadata about this environment.1157 1158Returns information about the environment including name, description,1159version, author, and documentation links.1160 """,1161 ),1162 GetEndpointConfig(1163 path="/health",1164 handler=lambda: HealthResponse(status=HealthStatus.HEALTHY),1165 response_model=HealthResponse,1166 tag="Health",1167 summary="Health check",1168 description="Check if the environment server is running and healthy.",1169 ),1170 ]1171 1172 # Only register /state endpoint in simulation mode1173 if mode == ServerMode.SIMULATION:1174 get_endpoints.insert(1175 0,1176 GetEndpointConfig(1177 path="/state",1178 handler=get_state_handler,1179 response_model=State,1180 tag="State Management",1181 summary="Get current environment state",1182 description="""1183Retrieve the current internal state of the environment.1184 1185The structure of the state object is defined by the environment's State model.1186 """,1187 ),1188 )1189 1190 register_get_endpoints(app, get_endpoints)1191 1192 # Register combined schema endpoint1193 @app.get(1194 "/schema",1195 response_model=SchemaResponse,1196 tags=["Schema"],1197 summary="Get all JSON schemas",1198 description="""1199Get JSON schemas for actions, observations, and state in a single response.1200