openenv/echo_env

1# SPDX-License-Identifier: BSD-3-Clause2 3"""4Synchronous wrapper for async EnvClient.5 6This module provides a SyncEnvClient that wraps an async EnvClient,7allowing synchronous usage while the underlying client uses async I/O.8 9Examples:10 11    ```python12    from openenv.core import GenericEnvClient13 14    # Create async client and get sync wrapper15    async_client = GenericEnvClient(base_url="http://localhost:8000")16    sync_client = async_client.sync()17 18    # Use synchronous API19    with sync_client:20        result = sync_client.reset()21        result = sync_client.step({"code": "print('hello')"})22    ```23"""24 25from __future__ import annotations26 27import asyncio28import concurrent.futures29import inspect30import threading31from contextlib import suppress32from typing import Any, Dict, Generic, TYPE_CHECKING, TypeVar33 34from .client_types import StateT, StepResult35 36if TYPE_CHECKING:37    from .env_client import EnvClient38 39ActT = TypeVar("ActT")40ObsT = TypeVar("ObsT")41 42 43class SyncEnvClient(Generic[ActT, ObsT, StateT]):44    """45    Synchronous wrapper around an async EnvClient.46 47    This class provides a synchronous interface to an async EnvClient,48    making it easier to use in synchronous code or to stop async from49    "infecting" the entire call stack.50 51    The wrapper executes async operations on a dedicated background event loop52    so connection state remains bound to a single loop.53 54    For guaranteed resource cleanup, use `with SyncEnvClient(...)` or call55    `close()` explicitly. `__del__` is best-effort only and may not run56    reliably (for example, during interpreter shutdown).57 58    Examples:59 60        ```python61        # From an async client62        async_client = GenericEnvClient(base_url="http://localhost:8000")63        sync_client = async_client.sync()64 65        # Use synchronous context manager66        with sync_client:67            result = sync_client.reset()68            result = sync_client.step({"action": "test"})69        ```70 71    Attributes:72        _async: The wrapped async EnvClient instance73    """74 75    def __init__(self, async_client: "EnvClient[ActT, ObsT, StateT]"):76        """77        Initialize sync wrapper around an async client.78 79        Args:80            async_client (`EnvClient`):81                The async client to wrap.82        """83        self._async = async_client84        self._loop: asyncio.AbstractEventLoop | None = None85        self._loop_thread: threading.Thread | None = None86        self._loop_ready = threading.Event()87        self._loop_init_lock = threading.Lock()88        self._async_wrapper_cache: Dict[str, Any] = {}89        self._child_clients: list[SyncEnvClient[ActT, ObsT, StateT]] = []90 91    def _run_loop_forever(self) -> None:92        """Run a dedicated event loop for this sync client."""93        loop = asyncio.new_event_loop()94        self._loop = loop95        asyncio.set_event_loop(loop)96        self._loop_ready.set()97        loop.run_forever()98        loop.close()99 100    def _ensure_loop(self) -> asyncio.AbstractEventLoop:101        """Start background loop thread on first use."""102        if (103            self._loop is not None104            and self._loop_thread105            and self._loop_thread.is_alive()106        ):107            return self._loop108 109        # Protect loop initialization when multiple threads race on first use.110        with self._loop_init_lock:111            if (112                self._loop is not None113                and self._loop_thread114                and self._loop_thread.is_alive()115            ):116                return self._loop117 118            self._loop_ready.clear()119            self._loop_thread = threading.Thread(120                target=self._run_loop_forever,121                name="openenv-sync-client-loop",122                daemon=True,123            )124            self._loop_thread.start()125            if not self._loop_ready.wait(timeout=5):126                raise RuntimeError("Timed out starting sync client event loop")127            assert self._loop is not None128            return self._loop129 130    def _claim_sync_mode(self) -> None:131        if hasattr(self._async, "_claim_execution_mode"):132            self._async._claim_execution_mode("sync")133 134    def _run(self, coro: Any) -> Any:135        """Run coroutine on dedicated loop and block for result."""136        self._claim_sync_mode()137        loop = self._ensure_loop()138        future: concurrent.futures.Future[Any] = asyncio.run_coroutine_threadsafe(139            coro, loop140        )141        return future.result()142 143    def _stop_loop(self) -> None:144        """Stop and join background loop thread."""145        loop = self._loop146        thread = self._loop_thread147        if loop is None:148            return149 150        if loop.is_running():151            loop.call_soon_threadsafe(loop.stop)152        if thread is not None:153            thread.join(timeout=5)154 155        self._loop = None156        self._loop_thread = None157 158    @property159    def async_client(self) -> "EnvClient[ActT, ObsT, StateT]":160        """Access the underlying async client."""161        return self._async162 163    def connect(self) -> "SyncEnvClient[ActT, ObsT, StateT]":164        """165        Establish connection to the server.166 167        Returns:168            self for method chaining169        """170        self._claim_sync_mode()171        self._run(self._async._connect_async())172        return self173 174    def disconnect(self) -> None:175        """Close the connection."""176        self._claim_sync_mode()177        self._run(self._async._disconnect_async())178 179    def reset(self, **kwargs: Any) -> StepResult[ObsT]:180        """181        Reset the environment.182 183        Args:184            **kwargs:185                Optional parameters passed to the environment's reset method.186 187        Returns:188            StepResult containing initial observation189        """190        self._claim_sync_mode()191        return self._run(self._async._reset_async(**kwargs))192 193    def step(self, action: ActT, **kwargs: Any) -> StepResult[ObsT]:194        """195        Execute an action in the environment.196 197        Args:198            action:199                The action to execute.200            **kwargs:201                Optional parameters.202 203        Returns:204            StepResult containing observation, reward, and done status205        """206        self._claim_sync_mode()207        return self._run(self._async._step_async(action, **kwargs))208 209    def state(self) -> StateT:210        """211        Get the current environment state.212 213        Returns:214            State object with environment state information215        """216        self._claim_sync_mode()217        return self._run(self._async._state_async())218 219    def close(self) -> None:220        """Close the connection and clean up resources."""221        try:222            for child in list(self._child_clients):223                with suppress(Exception):224                    child.close()225            self._child_clients.clear()226            self._claim_sync_mode()227            self._run(self._async._close_async())228        finally:229            self._stop_loop()230 231    def new_session(self) -> "SyncEnvClient[ActT, ObsT, StateT]":232        """233        Create a new synchronous session against the same environment server.234 235        Returns:236            `SyncEnvClient`: A connected child wrapper around a child async237            client of the same concrete type.238 239        The child session is tracked by this parent and closed when the parent240        is closed. Call this after the parent has connected, because the child241        reuses the parent's current base URL. Server-side capacity still242        applies: when the server is at `MAX_CONCURRENT_ENVS`, opening the child243        WebSocket can fail and is surfaced as a connection error.244        """245        async_client = self._async._create_session_client()246        client = SyncEnvClient(async_client)247        client.connect()248        self._child_clients.append(client)249        return client250 251    def __enter__(self) -> "SyncEnvClient[ActT, ObsT, StateT]":252        """Enter context manager, establishing connection."""253        self.connect()254        return self255 256    def __exit__(self, exc_type, exc_val, exc_tb) -> None:257        """Exit context manager, closing connection."""258        self.close()259 260    def __del__(self) -> None:261        """262        Best-effort cleanup for background loop thread.263 264        Do not rely on this for deterministic cleanup; prefer context-manager265        usage or an explicit `close()` call.266        """267        try:268            self._stop_loop()269        except Exception:270            pass271 272    def __getattr__(self, name: str) -> Any:273        """274        Delegate unknown attributes to the async client.275 276        Async methods are wrapped to run on the sync client's dedicated loop.277        """278        attr = getattr(self._async, name)279 280        if inspect.iscoroutinefunction(attr):281            cached = self._async_wrapper_cache.get(name)282            if cached is not None:283                return cached284 285            def sync_wrapper(*args: Any, **kwargs: Any) -> Any:286                method = getattr(self._async, name)287                return self._run(method(*args, **kwargs))288 289            self._async_wrapper_cache[name] = sync_wrapper290            return sync_wrapper291 292        return attr293 294    # Delegate abstract method implementations to the wrapped client295    def _step_payload(self, action: ActT) -> Dict[str, Any]:296        """Delegate to async client's _step_payload."""297        return self._async._step_payload(action)298 299    def _parse_result(self, payload: Dict[str, Any]) -> StepResult[ObsT]:300        """Delegate to async client's _parse_result."""301        return self._async._parse_result(payload)302 303    def _parse_state(self, payload: Dict[str, Any]) -> StateT:304        """Delegate to async client's _parse_state."""305        return self._async._parse_state(payload)306