openenv/echo_env
6
1# SPDX-License-Identifier: BSD-3-Clause2 3"""4HTTP server wrapper for Environment instances.5 6This module provides utilities to wrap any Environment subclass and expose it7over HTTP and WebSocket endpoints that EnvClient can consume.8"""9 10from __future__ import annotations11 12import asyncio13import inspect14import json15import logging16import os17import time18import uuid19from concurrent.futures import ThreadPoolExecutor20from contextlib import AsyncExitStack21from typing import (22 Any,23 AsyncContextManager,24 Awaitable,25 Callable,26 cast,27 Dict,28 Optional,29 Type,30 TypeVar,31)32 33_MISSING = object()34_MCPResult = TypeVar("_MCPResult")35 36from fastapi import (37 Body,38 FastAPI,39 HTTPException,40 Request,41 status,42 WebSocket,43 WebSocketDisconnect,44)45from pydantic import ValidationError46 47from ._utils import overrides_method48from .interfaces import Environment49from .mcp_environment import get_server_tools50from .mcp_types import (51 CallToolAction,52 CallToolObservation,53 JsonRpcErrorCode,54 JsonRpcRequest,55 JsonRpcResponse,56 ListToolsAction,57 ListToolsObservation,58 McpMethod,59 WSMCPMessage,60 WSMCPResponse,61)62from .route_config import GetEndpointConfig, register_get_endpoints63from .serialization import deserialize_action, serialize_observation64from .types import (65 Action,66 ConcurrencyConfig,67 EnvironmentMetadata,68 GetTaskRangeRequest,69 GetTaskRequest,70 HealthResponse,71 HealthStatus,72 ListTasksRequest,73 NumTasksRequest,74 Observation,75 ResetRequest,76 ResetResponse,77 SchemaResponse,78 ServerCapacityStatus,79 ServerMode,80 SessionInfo,81 State,82 StepRequest,83 StepResponse,84 WSCloseMessage,85 WSErrorCode,86 WSErrorResponse,87 WSObservationResponse,88 WSResetMessage,89 WSStateMessage,90 WSStateResponse,91 WSStepMessage,92)93 94 95def _make_json_serializable(obj: Any) -> Any:96 """97 Convert an object to a JSON-serializable form.98 99 Handles Pydantic models, dataclasses, and other common types.100 101 Args:102 obj (`Any`):103 The object to convert.104 105 Returns:106 `Any`: A JSON-serializable representation of the object.107 """108 if obj is None:109 return None110 if isinstance(obj, (str, int, float, bool)):111 return obj112 if isinstance(obj, (list, tuple)):113 return [_make_json_serializable(item) for item in obj]114 if isinstance(obj, dict):115 return {k: _make_json_serializable(v) for k, v in obj.items()}116 if hasattr(obj, "model_dump"):117 # Pydantic model118 return obj.model_dump()119 if hasattr(obj, "__dict__"):120 # Object with __dict__121 return {k: _make_json_serializable(v) for k, v in obj.__dict__.items()}122 # Fallback to string representation123 return str(obj)124 125 126async def _maybe_await(value: Any) -> Any:127 """Await values returned by async task APIs while preserving sync APIs."""128 if inspect.isawaitable(value):129 return await value130 return value131 132 133from .exceptions import (134 ConcurrencyConfigurationError,135 EnvironmentFactoryError,136 SessionCapacityError,137)138 139 140class HTTPEnvServer:141 """142 HTTP server wrapper for Environment instances.143 144 This class wraps an Environment and exposes its reset(), step(), and state145 methods as HTTP and WebSocket endpoints compatible with EnvClient.146 147 The server expects:148 - Action deserialization: Converts JSON dict to Action subclass149 - Observation serialization: Converts Observation subclass to JSON dict150 151 Examples:152 153 ```python154 from core.env_server import HTTPEnvServer155 from envs.coding_env.server import CodeExecutionEnvironment156 from envs.coding_env.models import CodeAction, CodeObservation157 158 # Pass environment class (factory pattern)159 server = HTTPEnvServer(160 env=CodeExecutionEnvironment,161 action_cls=CodeAction,162 observation_cls=CodeObservation,163 max_concurrent_envs=4,164 )165 166 # Register routes with FastAPI167 from fastapi import FastAPI168 app = FastAPI()169 server.register_routes(app)170 ```171 """172 173 def __init__(174 self,175 env: Callable[[], Environment],176 action_cls: Type[Action],177 observation_cls: Type[Observation],178 max_concurrent_envs: Optional[int] = None,179 concurrency_config: Optional[ConcurrencyConfig] = None,180 env_name: Optional[str] = None,181 state_cls: Type[State] = State,182 *,183 reset_observation_cls: Optional[Type[Observation]] = None,184 ):185 """186 Initialize HTTP server wrapper.187 188 Args:189 env (`Callable[[], Environment]`):190 Environment factory (callable) that creates new instances. Will be called191 to create a new environment for each WebSocket session.192 action_cls (`Type[Action]`):193 The `Action` subclass this environment expects.194 observation_cls (`Type[Observation]`):195 The `Observation` subclass returned by step.196 reset_observation_cls (`Type[Observation]`, *optional*):197 The reset observation model published in `/schema`. Defaults to198 `observation_cls`; declare a distinct model when reset differs.199 max_concurrent_envs (`int`, *optional*):200 Maximum number of concurrent WebSocket sessions. Mutually exclusive with201 `concurrency_config`.202 concurrency_config (`ConcurrencyConfig`, *optional*):203 Advanced concurrency settings. Mutually exclusive with204 `max_concurrent_envs`.205 env_name (`str`, *optional*):206 Public environment name used by task/split endpoints.207 state_cls (`Type[State]`, *optional*, defaults to `State`):208 The `State` subclass this environment reports. Used for the `/state`209 response model and the `state` entry of `/schema`, so that fields210 declared by the subclass are published and serialized.211 212 Raises:213 `ValueError`: If both `max_concurrent_envs` and `concurrency_config` are provided.214 `ConcurrencyConfigurationError`: If `max_concurrent_envs` > 1 for an environment215 that is not marked as `SUPPORTS_CONCURRENT_SESSIONS`.216 """217 # Validate that env is callable218 if not callable(env):219 raise TypeError(220 f"env must be a callable (class or factory function), got {type(env)}. "221 f"Pass the environment class (e.g., MyEnvironment) not an instance (e.g., MyEnvironment())."222 )223 224 self._env_factory: Callable[[], Environment] = env225 226 # Handle concurrency configuration227 if max_concurrent_envs is not None and concurrency_config is not None:228 raise ValueError(229 "Cannot specify both 'max_concurrent_envs' and 'concurrency_config'. "230 "Please use only one method to configure concurrency."231 )232 233 if concurrency_config is not None:234 self._concurrency_config = concurrency_config235 elif max_concurrent_envs is not None:236 self._concurrency_config = ConcurrencyConfig(237 max_concurrent_envs=max_concurrent_envs,238 session_timeout=None,239 )240 else:241 # Default configuration242 self._concurrency_config = ConcurrencyConfig(243 max_concurrent_envs=1,244 session_timeout=None,245 )246 247 self._max_concurrent_envs = self._concurrency_config.max_concurrent_envs248 249 # Validate concurrency configuration250 self._validate_concurrency_safety()251 252 self.action_cls = action_cls253 self.observation_cls = observation_cls254 self.reset_observation_cls = reset_observation_cls or observation_cls255 self.state_cls = state_cls256 self.env_name = env_name or self._default_env_name()257 258 # Session management for WebSocket connections259 self._sessions: Dict[str, Optional[Environment]] = {}260 self._session_executors: Dict[str, ThreadPoolExecutor] = {}261 self._session_stacks: Dict[str, AsyncExitStack] = {}262 self._session_info: Dict[str, SessionInfo] = {}263 self._session_websocket_attachments: set[str] = set()264 self._session_pending_closes: set[str] = set()265 self._session_lock = asyncio.Lock()266 267 # Create thread pool for running sync code in async context268 # This is needed for environments using sync libraries (e.g., Playwright)269 self._executor = ThreadPoolExecutor(max_workers=32)270 271 self._requires_single_thread_executor = self._detect_single_thread_requirement()272 self._shared_session_executor: Optional[ThreadPoolExecutor] = None273 if self._requires_single_thread_executor:274 self._shared_session_executor = ThreadPoolExecutor(max_workers=1)275 276 # Idle session reaper configuration.277 # Timeout is taken from ConcurrencyConfig.session_timeout;278 # None means no timeout (default — reaper is a no-op).279 self._session_idle_timeout_s: Optional[float] = (280 self._concurrency_config.session_timeout281 )282 self._reaper_task: Optional[asyncio.Task[None]] = None283 284 def _default_env_name(self) -> str:285 factory = self._env_factory286 if inspect.isclass(factory):287 return factory.__name__288 return getattr(factory, "__name__", "environment")289 290 def _validate_concurrency_safety(self) -> None:291 """292 Validate that the environment supports the configured concurrency level.293 294 Raises:295 `ConcurrencyConfigurationError`: If `max_concurrent_envs` > 1 for an environment296 that is not marked as `SUPPORTS_CONCURRENT_SESSIONS`.297 """298 import functools299 300 if self._max_concurrent_envs <= 1:301 return302 303 factory = self._env_factory304 if isinstance(factory, functools.partial):305 factory = factory.func306 307 if inspect.isclass(factory):308 env_cls = factory309 else:310 _temp_env = self._env_factory()311 env_cls = type(_temp_env)312 _temp_env.close()313 del _temp_env314 315 if not getattr(env_cls, "SUPPORTS_CONCURRENT_SESSIONS", False):316 raise ConcurrencyConfigurationError(317 environment_name=env_cls.__name__,318 max_concurrent_envs=self._max_concurrent_envs,319 )320 321 def _detect_single_thread_requirement(self) -> bool:322 import functools323 324 factory = self._env_factory325 if isinstance(factory, functools.partial):326 factory = factory.func327 if inspect.isclass(factory):328 return getattr(factory, "REQUIRES_SINGLE_THREAD_EXECUTOR", False)329 return False330 331 def get_capacity_status(self) -> ServerCapacityStatus:332 """333 Get the current capacity status of the server.334 335 Returns:336 [`ServerCapacityStatus`] with current session counts and availability.337 """338 return ServerCapacityStatus.from_counts(339 active=len(self._sessions),340 max_sessions=self._max_concurrent_envs,341 )342 343 async def _run_sync_in_thread_pool(344 self, func: Callable[..., Observation], *args, **kwargs345 ) -> Observation:346 """Run a synchronous function in the thread pool executor."""347 loop = asyncio.get_event_loop()348 return await loop.run_in_executor(self._executor, lambda: func(*args, **kwargs))349 350 def _get_valid_kwargs(351 self,352 sig: inspect.Signature,353 kwargs: Dict[str, Any],354 skip_params: Optional[set[str]] = None,355 ) -> Dict[str, Any]:356 """Filter kwargs to only include parameters accepted by the function signature."""357 if skip_params is None:358 skip_params = set()359 360 valid_kwargs = {}361 362 has_kwargs = any(363 p.kind == inspect.Parameter.VAR_KEYWORD for p in sig.parameters.values()364 )365 366 for k, v in kwargs.items():367 if k in sig.parameters or has_kwargs:368 if k not in skip_params:369 valid_kwargs[k] = v370 371 return valid_kwargs372 373 async def _create_session(self) -> tuple[str, Environment]:374 """375 Create a new WebSocket session with its own environment instance.376 377 Returns:378 `tuple[str, Environment]`: Tuple of (session_id, environment).379 380 Raises:381 `SessionCapacityError`: If max concurrent sessions reached.382 `EnvironmentFactoryError`: If the factory fails to create an environment.383 """384 async with self._session_lock:385 if len(self._sessions) >= self._max_concurrent_envs:386 raise SessionCapacityError(387 active_sessions=len(self._sessions),388 max_sessions=self._max_concurrent_envs,389 )390 391 session_id = str(uuid.uuid4())392 current_time = time.time()393 394 # Create executor and reserve slot so capacity is not exceeded while395 # we create the env outside the lock (avoids blocking other sessions)396 if self._shared_session_executor is not None:397 executor = self._shared_session_executor398 else:399 executor = ThreadPoolExecutor(max_workers=1)400 self._session_executors[session_id] = executor401 self._sessions[session_id] = None # placeholder until env is ready402 403 try:404 # Create environment in the executor thread (outside lock)405 loop = asyncio.get_event_loop()406 env = await loop.run_in_executor(executor, self._env_factory)407 except Exception as e:408 async with self._session_lock:409 if executor is not self._shared_session_executor:410 executor.shutdown(wait=False)411 self._session_executors.pop(session_id, None)412 self._sessions.pop(session_id, None)413 factory_name = getattr(414 self._env_factory, "__name__", str(self._env_factory)415 )416 raise EnvironmentFactoryError(factory_name) from e417 418 # Hold the MCP session open for the lifetime of this session,419 # matching the WebSocket path's AsyncExitStack pattern. This420 # prevents per-request MCP transport teardown/reconnection and421 # preserves FastMCP session state (ctx.set_state / ctx.get_state)422 # across HTTP calls within the same OpenEnv session.423 stack = AsyncExitStack()424 try:425 mcp_session_factory = getattr(env, "mcp_session", None)426 if callable(mcp_session_factory):427 mcp_session_cm = cast(AsyncContextManager[Any], mcp_session_factory())428 await stack.enter_async_context(mcp_session_cm)429 except Exception:430 # MCP transport failed to start — clean up the reserved slot,431 # the env, and the executor so they don't leak permanently432 # against _max_concurrent_envs.433 await stack.aclose() # best-effort434 async with self._session_lock:435 self._sessions.pop(session_id, None)436 self._session_executors.pop(session_id, None)437 self._session_info.pop(session_id, None)438 await self._cleanup_session_resources(env, executor)439 raise440 441 async with self._session_lock:442 self._sessions[session_id] = env443 self._session_stacks[session_id] = stack444 now = time.time()445 self._session_info[session_id] = SessionInfo(446 session_id=session_id,447 created_at=current_time,448 last_activity_at=now,449 step_count=0,450 environment_type=type(env).__name__,451 )452 453 return session_id, env454 455 async def _destroy_session(self, session_id: str) -> None:456 """457 Destroy a WebSocket session and cleanup resources.458 459 Args:460 session_id (`str`):461 The session ID to destroy.462 """463 async with self._session_lock:464 env = self._sessions.pop(session_id, None)465 executor = self._session_executors.pop(session_id, None)466 stack = self._session_stacks.pop(session_id, None)467 self._session_info.pop(session_id, None)468 self._session_websocket_attachments.discard(session_id)469 self._session_pending_closes.discard(session_id)470 471 await self._cleanup_session_resources(env, executor, stack)472 473 async def _cleanup_session_resources(474 self,475 env: Optional[Environment],476 executor: Optional[ThreadPoolExecutor],477 stack: Optional[AsyncExitStack] = None,478 ) -> None:479 """Close an environment and shut down its executor (best-effort)."""480 # Close the MCP session stack first — this gracefully exits the481 # mcp_session() context (and the underlying FastMCP Client session)482 # before we tear down the environment references.483 if stack is not None:484 try:485 await stack.aclose()486 except Exception:487 pass # Best effort cleanup488 489 # Run close() in the same executor where the env was created490 # This is required for thread-sensitive libraries like Playwright/greenlet491 if env is not None:492 if executor is not None:493 try:494 loop = asyncio.get_event_loop()495 await loop.run_in_executor(executor, env.close)496 except Exception:497 # If executor close fails, try direct close as fallback498 try:499 env.close()500 except Exception:501 pass # Best effort cleanup502 else:503 try:504 env.close()505 except Exception:506 pass # Best effort cleanup507 508 if executor is not None and executor is not self._shared_session_executor:509 executor.shutdown(wait=False)510 511 def _update_session_activity(512 self, session_id: str, increment_step: bool = False513 ) -> None:514 """515 Update session activity timestamp and optionally increment step count.516 517 Args:518 session_id (`str`):519 The session ID to update.520 increment_step (`bool`, *optional*, defaults to `False`):521 If `True`, increment the step count.522 """523 if session_id in self._session_info:524 self._session_info[session_id].last_activity_at = time.time()525 if increment_step:526 self._session_info[session_id].step_count += 1527 528 async def _reap_idle_sessions(self) -> None:529 """Background task that periodically destroys sessions idle beyond the timeout."""530 timeout = self._session_idle_timeout_s531 if timeout is None:532 return # no timeout configured — noop533 interval = max(timeout / 4, 5.0) # check frequently enough534 while True:535 try:536 await asyncio.sleep(interval)537 now = time.time()538 stale_ids: list[str] = []539 async with self._session_lock:540 for sid, info in self._session_info.items():541 if (542 sid not in self._session_websocket_attachments543 and now - info.last_activity_at > timeout544 ):545 stale_ids.append(sid)546 for sid in stale_ids:547 # Re-check under lock: activity may have arrived since548 # the snapshot was taken, making this session active again.549 # Refresh `now` so slow _destroy_session calls don't cause550 # subsequent entries to be validated against a stale clock.551 now = time.time()552 async with self._session_lock:553 info = self._session_info.get(sid)554 if (555 info is None556 or sid in self._session_websocket_attachments557 or (now - info.last_activity_at) <= timeout558 ):559 continue560 await self._destroy_session(sid)561 except asyncio.CancelledError:562 break563 except Exception as exc:564 logging.getLogger(__name__).warning(565 "Idle-session reaper encountered an error (will retry): %s",566 exc,567 )568 569 def _start_reaper(self) -> None:570 """Start the idle-session reaper if a timeout is configured."""571 if self._session_idle_timeout_s is not None and self._reaper_task is None:572 self._reaper_task = asyncio.create_task(self._reap_idle_sessions())573 574 def _stop_reaper(self) -> None:575 """Cancel the reaper background task."""576 if self._reaper_task is not None:577 self._reaper_task.cancel()578 self._reaper_task = None579 580 def get_session_info(self, session_id: str) -> Optional[SessionInfo]:581 """582 Get information about a specific session.583 584 Args:585 session_id (`str`):586 The session ID to query.587 588 Returns:589 [`SessionInfo`] if the session exists, `None` otherwise.590 """591 return self._session_info.get(session_id)592 593 async def _run_in_session_executor(594 self, session_id: str, func: Callable[..., Observation], *args, **kwargs595 ) -> Observation:596 """Run a synchronous function in the session's thread pool executor."""597 executor = self._session_executors.get(session_id, self._executor)598 loop = asyncio.get_event_loop()599 return await loop.run_in_executor(executor, lambda: func(*args, **kwargs))600 601 async def _run_mcp_client_operation(602 self,603 mcp_client: Any,604 mcp_session_factory: Any,605 managed_session_id: Optional[str],606 operation: Callable[[], Awaitable[_MCPResult]],607 ) -> _MCPResult:608 """Run an MCP client operation with the appropriate transport lifecycle."""609 if managed_session_id and mcp_client.is_connected():610 return await operation()611 612 if callable(mcp_session_factory):613 mcp_session_cm = cast(AsyncContextManager[Any], mcp_session_factory())614 async with mcp_session_cm:615 return await operation()616 617 async with mcp_client:618 return await operation()619 620 @property621 def active_sessions(self) -> int:622 """Return the number of active WebSocket sessions."""623 return len(self._sessions)624 625 @property626 def max_concurrent_envs(self) -> int:627 """Return the maximum number of concurrent environments."""628 return self._max_concurrent_envs629 630 @property631 def is_concurrency_safe(self) -> bool:632 """Return whether the environment is marked as concurrency safe."""633 import inspect634 635 if inspect.isclass(self._env_factory):636 return getattr(self._env_factory, "SUPPORTS_CONCURRENT_SESSIONS", False)637 else:638 _temp_env = self._env_factory()639 result = getattr(_temp_env, "SUPPORTS_CONCURRENT_SESSIONS", False)640 _temp_env.close()641 del _temp_env642 return result643 644 @property645 def concurrency_config(self) -> ConcurrencyConfig:646 """Return the concurrency configuration."""647 return self._concurrency_config648 649 def _factory_produces_harness_env(self) -> bool:650 """Return whether the env factory produces a HarnessEnvironment."""651 import inspect652 653 # Lazy import: openenv.core.harness imports env_server modules, so a654 # top-level import here would be circular.655 from ..harness.environment import HarnessEnvironment656 657 if inspect.isclass(self._env_factory):658 return issubclass(self._env_factory, HarnessEnvironment)659 _temp_env = self._env_factory()660 try:661 return isinstance(_temp_env, HarnessEnvironment)662 finally:663 _temp_env.close()664 665 def _register_harness_route(self, app: FastAPI) -> None:666 """667 Register the production `/harness` WebSocket route (RFC 005).668 669 Each connection gets its own environment session: connecting resets670 the environment (which starts the harness process and injects tools),671 and each `{"type": "message", "content": ...}` frame runs one672 conversational turn, streamed back as `HarnessEvent` JSON frames673 ending with a `turn_complete` event. Malformed client frames receive a674 recoverable `protocol_error` response without starting a turn; the675 connection remains usable. Terminal failures use `error`.676 """677 # Lazy import to avoid a circular import with openenv.core.harness.678 from ..harness.adapter import HarnessNotRunningError679 from ..harness.events import (680 HarnessClientMessage,681 HarnessEvent,682 HarnessEventType,683 HarnessProtocolError,684 )685 686 @app.websocket("/harness")687 async def harness_websocket_endpoint(websocket: WebSocket):688 await websocket.accept()689 690 session_id = None691 session_env = None692 693 async def send_error(message: str, code: WSErrorCode) -> None:694 error_response = WSErrorResponse(695 data={"message": message, "code": code}696 )697 await websocket.send_text(error_response.model_dump_json())698 699 async def send_protocol_error(message: str, code: WSErrorCode) -> None:700 error_response = HarnessProtocolError(701 data={"message": message, "code": code}702 )703 await websocket.send_text(error_response.model_dump_json())704 705 async def send_harness_error(message: str) -> None:706 """Emit a terminal ERROR event in the harness event stream."""707 error_event = HarnessEvent(708 type=HarnessEventType.ERROR,709 data={"message": message, "recoverable": False},710 )711 await websocket.send_text(error_event.model_dump_json())712 713 try:714 session_id, session_env = await self._create_session()715 # Protect the live harness from idle reaping and HTTP session716 # close, including startup and turns that emit no events.717 self._session_websocket_attachments.add(session_id)718 719 async with AsyncExitStack() as stack:720 mcp_session_factory = getattr(session_env, "mcp_session", None)721 if callable(mcp_session_factory):722 mcp_session_cm = cast(723 AsyncContextManager[Any], mcp_session_factory()724 )725 await stack.enter_async_context(mcp_session_cm)726 727 # Starts the harness process and injects environment tools728 await session_env.reset_async()729 await websocket.send_text(730 json.dumps(731 {732 "type": "session_started",733 "data": {734 "session_id": session_id,735 "harness": session_env.adapter.config.name,736 },737 }738 )739 )740 741 while True:742 raw_message = await websocket.receive_text()743 744 try:745 message_dict = json.loads(raw_message)746 except json.JSONDecodeError as e:747 await send_protocol_error(748 f"Invalid JSON: {e}", WSErrorCode.INVALID_JSON749 )750 continue751 try:752 client_message = HarnessClientMessage(**message_dict)753 except (ValidationError, TypeError) as e:754 await send_protocol_error(755 f"Invalid message: {e}",756 WSErrorCode.VALIDATION_ERROR,757 )758 continue759 760 self._update_session_activity(session_id, increment_step=True)761 762 async def stream_turn(content: str) -> bool:763 """Stream one turn; True if it ended with TURN_COMPLETE."""764 saw_terminal = False765 adapter = session_env.adapter766 if not await adapter.is_alive():767 raise HarnessNotRunningError(768 "harness process is not running"769 )770 async for event in adapter.send_message_streaming(content):771 await websocket.send_text(event.model_dump_json())772 # Record progress throughout the turn.773 self._update_session_activity(session_id)774 saw_terminal = (775 event.type is HarnessEventType.TURN_COMPLETE776 )777 return saw_terminal778 779 # Bound the turn in wall-clock time, matching what780 # simulation mode does in HarnessEnvironment._run_turn.781 # Without this a hung harness holds the session open782 # forever, and the server sits at capacity.783 turn_timeout_s = session_env.adapter.config.session_timeout_s784 try:785 completed = await asyncio.wait_for(786 stream_turn(client_message.content),787 turn_timeout_s,788 )789 except asyncio.TimeoutError:790 await send_harness_error(791 f"harness turn exceeded {turn_timeout_s} seconds"792 )793 break794 except HarnessNotRunningError:795 await send_harness_error("harness process is not running")796 break797 except Exception:798 # Harness state after a crash is undefined; end799 # the session so a reconnect gets a fresh one.800 # Adapter exceptions can contain credentials or801 # subprocess output; do not expose them to clients.802 await send_harness_error("harness turn failed")803 break804 805 if not completed:806 # send_message() raises HarnessError here; the807 # socket equivalent is to say so and end the808 # session, rather than leaving a client that809 # blocks on the terminal event waiting forever.810 await send_harness_error(811 "harness event stream ended without a "812 "TURN_COMPLETE event"813 )814 break815 816 except WebSocketDisconnect:817 pass818 except SessionCapacityError as e:819 await send_error(str(e), WSErrorCode.CAPACITY_REACHED)820 except EnvironmentFactoryError as e:821 await send_error(str(e), WSErrorCode.FACTORY_ERROR)822 except Exception:823 try:824 await send_error(825 "harness session failed", WSErrorCode.SESSION_ERROR826 )827 except (RuntimeError, WebSocketDisconnect):828 pass829 finally:830 if session_id:831 # Release ownership without an await so cancellation cannot832 # leave a session permanently exempt from idle reaping.833 self._session_websocket_attachments.discard(session_id)834 cleanup = asyncio.create_task(self._destroy_session(session_id))835 try:836 await asyncio.shield(cleanup)837 except asyncio.CancelledError:838 # ASGI cancellation must not orphan a running harness.839 # Finish teardown before propagating cancellation, even840 # when the request's cancel scope cancels us repeatedly.841 while not cleanup.done():842 try:843 await asyncio.shield(cleanup)844 except asyncio.CancelledError:845 pass846 cleanup.result()847 raise848 try:849 await websocket.close()850 except (RuntimeError, WebSocketDisconnect):851 # TestClient raises RuntimeError, real ASGI servers raise852 # WebSocketDisconnect when the client is already gone.853 pass854 855 def register_routes(856 self, app: FastAPI, mode: ServerMode | str = ServerMode.SIMULATION857 ) -> None:858 """859 Register HTTP routes on a FastAPI application.860 861 Args:862 app (`FastAPI`):863 FastAPI application instance.864 mode (`ServerMode` or `str`, *optional*, defaults to `ServerMode.SIMULATION`):865 Server mode. In production mode, simulation control endpoints (/reset, /step,866 /state) are NOT registered. Only safe endpoints (/health, /schema, /metadata,867 /ws) are available.868 869 Raises:870 `ValueError`: If `mode` is not a valid `ServerMode` or string equivalent.871 """872 # Convert string to ServerMode enum for backwards compatibility873 if isinstance(mode, str):874 try:875 mode = ServerMode(mode.lower())876 except ValueError:877 valid_modes = [m.value for m in ServerMode]878 raise ValueError(879 f"Invalid mode: '{mode}'. Must be one of: {valid_modes}"880 )881 882 # Wire up idle-session reaper lifecycle via app events883 server_ref = self884 885 async def _start_session_reaper() -> None:886 server_ref._start_reaper()887 888 async def _stop_session_reaper() -> None:889 server_ref._stop_reaper()890 if server_ref._shared_session_executor is not None:891 server_ref._shared_session_executor.shutdown(wait=True)892 893 if not getattr(app.router, "_openenv_reaper_registered", False):894 app.router.on_startup.append(_start_session_reaper)895 app.router.on_shutdown.append(_stop_session_reaper)896 app.router._openenv_reaper_registered = True # type: ignore[attr-defined]897 898 # Helper function to handle reset endpoint899 async def reset_handler(900 request: ResetRequest = Body(default_factory=ResetRequest),901 ) -> ResetResponse:902 """Reset endpoint - returns initial observation."""903 _env = self._env_factory()904 905 try:906 kwargs = request.model_dump(exclude_unset=True)907 908 is_async = overrides_method(_env.reset_async, Environment.reset_async)909 910 if is_async:911 sig = inspect.signature(_env.reset_async)912 else:913 sig = inspect.signature(_env.reset)914 valid_kwargs = self._get_valid_kwargs(sig, kwargs)915 916 if is_async:917 observation = await _env.reset_async(**valid_kwargs)918 else:919 observation = await self._run_sync_in_thread_pool(920 _env.reset, **valid_kwargs921 )922 return ResetResponse(**serialize_observation(observation))923 finally:924 _env.close()925 926 # Helper function to handle step endpoint927 async def step_handler(request: StepRequest) -> StepResponse:928 """Step endpoint - executes action and returns observation."""929 action_data = request.action930 931 try:932 action = deserialize_action(action_data, self.action_cls)933 except ValidationError as e:934 raise HTTPException(935 status_code=status.HTTP_422_UNPROCESSABLE_CONTENT, detail=e.errors()936 )937 938 _env = self._env_factory()939 940 try:941 kwargs = request.model_dump(exclude_unset=True, exclude={"action"})942 943 is_async = overrides_method(_env.step_async, Environment.step_async)944 945 if is_async:946 sig = inspect.signature(_env.step_async)947 else:948 sig = inspect.signature(_env.step)949 valid_kwargs = self._get_valid_kwargs(950 sig, kwargs, skip_params={"action"}951 )952 953 if is_async:954 observation = await _env.step_async(action, **valid_kwargs)955 else:956 observation = await self._run_sync_in_thread_pool(957 _env.step, action, **valid_kwargs958 )959 960 return StepResponse(**serialize_observation(observation))961 finally:962 _env.close()963 964 # Helper function to handle MCP endpoint965 async def mcp_handler(966 request: JsonRpcRequest,967 session_env: Optional[Environment] = None,968 session_id: Optional[str] = None,969 ) -> JsonRpcResponse:970 """971 Handle MCP JSON-RPC requests.972 973 Supports tools/list and tools/call methods in JSON-RPC 2.0 format,974 plus OpenEnv session lifecycle methods for HTTP MCP (openenv/session/create,975 openenv/session/close).976 """977 method = request.method978 request_id = request.id979 params = request.params980 if not isinstance(params, dict):981 return JsonRpcResponse.error_response(982 JsonRpcErrorCode.INVALID_PARAMS,983 "Params must be an object",984 request_id=request_id,985 )986 987 # OpenEnv extension methods for explicit MCP session management.988 # This enables persistent MCP lifecycles over HTTP /mcp, matching WebSocket semantics.989 if method == "openenv/session/create":990 if session_env is not None and session_id is not None:991 return JsonRpcResponse.success(992 result={"session_id": session_id},993 request_id=request_id,994 )995 try:996 created_session_id, _ = await self._create_session()997 except SessionCapacityError as e:998 return JsonRpcResponse.error_response(999 JsonRpcErrorCode.SERVER_ERROR,1000 str(e),1001 request_id=request_id,1002 data={1003 "active_sessions": e.active_sessions,1004 "max_sessions": e.max_sessions,1005 },1006 )1007 except EnvironmentFactoryError as e:1008 return JsonRpcResponse.error_response(1009 JsonRpcErrorCode.SERVER_ERROR,1010 str(e),1011 request_id=request_id,1012 data={"factory_name": e.factory_name},1013 )1014 return JsonRpcResponse.success(1015 result={"session_id": created_session_id},1016 request_id=request_id,1017 )1018 1019 if method == "openenv/session/close":1020 target_session_id = params.get("session_id")1021 if not target_session_id:1022 return JsonRpcResponse.error_response(1023 JsonRpcErrorCode.INVALID_PARAMS,1024 "Invalid params - 'session_id' is required",1025 request_id=request_id,1026 )1027 1028 if session_id is not None and target_session_id == session_id:1029 return JsonRpcResponse.error_response(1030 JsonRpcErrorCode.INVALID_REQUEST,1031 "Cannot close active WebSocket-managed session via MCP method",1032 request_id=request_id,1033 )1034 1035 async with self._session_lock:1036 if target_session_id in self._session_websocket_attachments:1037 env = _MISSING1038 attached = True1039 self._session_pending_closes.add(target_session_id)1040 executor = None1041 stack = None1042 else:1043 attached = False1044 env = self._sessions.pop(target_session_id, _MISSING)1045 if not attached and env is not _MISSING:1046 executor = self._session_executors.pop(target_session_id, None)1047 stack = self._session_stacks.pop(target_session_id, None)1048 self._session_info.pop(target_session_id, None)1049 elif not attached:1050 executor = None1051 stack = None1052 1053 if attached:1054 return JsonRpcResponse.success(1055 result={1056 "session_id": target_session_id,1057 "closed": False,1058 "closing": True,1059 },1060 request_id=request_id,1061 )1062 1063 if env is _MISSING:1064 return JsonRpcResponse.error_response(1065 JsonRpcErrorCode.INVALID_PARAMS,1066 f"Unknown session_id: {target_session_id}",1067 request_id=request_id,1068 )1069 1070 if env is None:1071 # Session slot reserved but env factory still running;1072 # re-insert the placeholder AND the executor so1073 # _create_session can finish and the executor remains1074 # tracked for eventual shutdown.1075 async with self._session_lock:1076 self._sessions[target_session_id] = None1077 if executor is not None:1078 self._session_executors[target_session_id] = executor1079 return JsonRpcResponse.error_response(1080 JsonRpcErrorCode.INVALID_REQUEST,1081 f"Session {target_session_id} is still initializing; retry shortly",1082 request_id=request_id,1083 )1084 1085 # env/executor/stack cleanup outside the lock1086 await self._cleanup_session_resources(env, executor, stack)1087 return JsonRpcResponse.success(1088 result={"session_id": target_session_id, "closed": True},1089 request_id=request_id,1090 )1091 1092 requested_session_id = params.get("session_id")1093 managed_session_id = session_id1094 1095 # Use provided session environment or create temporary one1096 if session_env is not None:1097 _env = session_env1098 should_close = False1099 elif requested_session_id:1100 async with self._session_lock:1101 _env = self._sessions.get(requested_session_id, _MISSING)1102 1103 if _env is _MISSING:1104 return JsonRpcResponse.error_response(1105 JsonRpcErrorCode.INVALID_PARAMS,1106 f"Unknown session_id: {requested_session_id}",1107 request_id=request_id,1108 )1109 1110 if _env is None:1111 return JsonRpcResponse.error_response(1112 JsonRpcErrorCode.INVALID_REQUEST,1113 f"Session {requested_session_id} is still initializing; retry shortly",1114 request_id=request_id,1115 )1116 1117 should_close = False1118 managed_session_id = requested_session_id1119 else:1120 _env = self._env_factory()1121 should_close = True1122 try:1123 mcp_client = getattr(_env, "mcp_client", None)1124 mcp_server = getattr(_env, "mcp_server", None)1125 mcp_session_factory = getattr(_env, "mcp_session", None)1126 1127 async def call_mcp_style_step(action: Action) -> Observation:1128 is_async = overrides_method(_env.step_async, Environment.step_async)1129 if is_async:1130 return await _env.step_async(action)1131 if managed_session_id:1132 return await self._run_in_session_executor(1133 managed_session_id,1134 _env.step,1135 action,1136 )1137 return await self._run_sync_in_thread_pool(_env.step, action)1138 1139 supports_mcp_style_actions = self.action_cls in {1140 CallToolAction,1141 ListToolsAction,1142 }1143 1144 if method == McpMethod.TOOLS_LIST:1145 # Check if environment is MCP-enabled1146 if mcp_client is None and mcp_server is None:1147 if supports_mcp_style_actions:1148 observation = await call_mcp_style_step(ListToolsAction())1149 if isinstance(observation, ListToolsObservation):1150 return JsonRpcResponse.success(1151 result={1152 "tools": [1153 tool.model_dump()1154 for tool in observation.tools1155 ]1156 },1157 request_id=request_id,1158 )1159 return JsonRpcResponse.error_response(1160 JsonRpcErrorCode.INTERNAL_ERROR,1161 "MCP-style tools/list step returned "1162 f"{type(observation).__name__}, expected "1163 "ListToolsObservation",1164 request_id=request_id,1165 )1166 return JsonRpcResponse.error_response(1167 JsonRpcErrorCode.INTERNAL_ERROR,1168 "Environment does not support MCP",1169 request_id=request_id,1170 )1171 1172 if mcp_client:1173 tools = await self._run_mcp_client_operation(1174 mcp_client,1175 mcp_session_factory,1176 managed_session_id,1177 mcp_client.list_tools,1178 )1179 1180 return JsonRpcResponse.success(1181 result={1182 "tools": [1183 t.model_dump()1184 if hasattr(t, "model_dump")1185 else dict(t)1186 for t in tools1187 ]1188 },1189 request_id=request_id,1190 )1191 1192 if mcp_server:1193 tools = []1194 for _tool_name, tool in get_server_tools(mcp_server).items():1195 tools.append(1196 {1197 "name": tool.name,1198 "description": tool.description or "",1199 "inputSchema": tool.parameters or {},1200 }