Auto sync env clients (#874)

* fix: auto-dispatch env client sync calls

* fix: preserve async dispatch in sync loop

* fix: repair env client dispatch modes

Author: burtenshaw

src/openenv/core/env_client.py

@@ -40,7 +40,7 @@
 import time
 from abc import ABC, abstractmethod
 from contextlib import suppress
-from typing import Any, Dict, Generic, Optional, Type, TYPE_CHECKING, TypeVar
+from typing import Any, Callable, Dict, Generic, Optional, Type, TYPE_CHECKING, TypeVar
 from urllib.parse import urlsplit
 
 from .client_types import StateT, StepResult
@@ -58,10 +58,49 @@
 ActT = TypeVar("ActT")
 ObsT = TypeVar("ObsT")
 EnvClientT = TypeVar("EnvClientT", bound="EnvClient")
+ResultT = TypeVar("ResultT")
 
 _VALID_CLIENT_MODES = ("simulation", "production")
 
 
+class _AutoAsyncResult(Generic[ResultT]):
+    """Awaitable result that can also be resolved by synchronous access."""
+
+    def __init__(
+        self,
+        client: "EnvClient[Any, Any, Any]",
+        coro_factory: Callable[[], Any],
+    ):
+        self._client = client
+        self._coro_factory = coro_factory
+        self._resolved = False
+        self._result: ResultT | None = None
+
+    def __await__(self):
+        async def _await_result():
+            self._client._claim_execution_mode("async")
+            return await self._coro_factory()
+
+        return _await_result().__await__()
+
+    def _sync_result(self) -> ResultT:
+        if not self._resolved:
+            self._result = self._client._run_sync(self._coro_factory)
+            self._resolved = True
+        return self._result
+
+    def __getattr__(self, name: str) -> Any:
+        return getattr(self._sync_result(), name)
+
+    def __bool__(self) -> bool:
+        return bool(self._sync_result())
+
+    def __repr__(self) -> str:
+        if self._resolved:
+            return repr(self._result)
+        return f"<{type(self).__name__} pending>"
+
+
 def _normalize_mode(mode: Optional[str]) -> str:
     """Resolve and validate the client communication mode."""
     raw_mode = (
@@ -212,6 +251,8 @@ def __init__(
         self._start_provider_on_connect = base_url is None
         self._child_clients: list[EnvClient[Any, Any, Any]] = []
         self._ws: Optional[ClientConnection] = None
+        self._execution_mode: Optional[str] = None
+        self._sync_client: Optional["SyncEnvClient[ActT, ObsT, StateT]"] = None
         self._ws_loop: Optional[asyncio.AbstractEventLoop] = None
         if base_url is not None:
             self._set_base_url(base_url)
@@ -295,7 +336,56 @@ def __setattr__(self, name: str, value: Any) -> None:
             raise AttributeError("Cannot modify mode after initialization")
         super().__setattr__(name, value)
 
-    async def connect(self) -> "EnvClient":
+    def _claim_execution_mode(self, mode: str) -> None:
+        """Lock the client to sync or async execution on first use."""
+        if self._execution_mode is None:
+            self._execution_mode = mode
+        elif self._execution_mode != mode:
+            raise RuntimeError(
+                f"EnvClient is already being used in {self._execution_mode} mode. "
+                "Create a separate client instance when mixing sync and async code."
+            )
+
+    def _run_sync(
+        self,
+        coro_factory: Callable[[], Any],
+        *,
+        allow_async_handoff: bool = False,
+    ) -> Any:
+        """Run an async operation through the sync wrapper."""
+        if allow_async_handoff and self._execution_mode == "async":
+            self._execution_mode = "sync"
+        else:
+            self._claim_execution_mode("sync")
+        if self._sync_client is None:
+            self._sync_client = self.sync()
+        return self._sync_client._run(coro_factory())
+
+    def _dispatch(self, coro_factory: Callable[[], Any]) -> Any:
+        """Return an awaitable in async code and a concrete result in sync code."""
+        try:
+            running_loop = asyncio.get_running_loop()
+        except RuntimeError:
+            running_loop = None
+
+        if self._execution_mode == "sync":
+            sync_loop = (
+                self._sync_client._loop if self._sync_client is not None else None
+            )
+            if running_loop is sync_loop:
+                return coro_factory()
+            return self._run_sync(coro_factory)
+        if running_loop is not None:
+            if self._execution_mode == "async":
+                return coro_factory()
+            return _AutoAsyncResult(self, coro_factory)
+
+        return self._run_sync(coro_factory, allow_async_handoff=True)
+
+    def connect(self) -> Any:
+        return self._dispatch(self._connect_async)
+
+    async def _connect_async(self) -> "EnvClient":
         """
         Establish WebSocket connection to the server.
 
@@ -352,7 +442,10 @@ async def connect(self) -> "EnvClient":
 
         return self
 
-    async def disconnect(self) -> None:
+    def disconnect(self) -> Any:
+        return self._dispatch(self._disconnect_async)
+
+    async def _disconnect_async(self) -> None:
         """Close the WebSocket connection."""
         if self._ws is not None:
             ws = self._ws
@@ -374,17 +467,13 @@ async def disconnect(self) -> None:
     async def _ensure_connected(self) -> None:
         """Ensure WebSocket connection is established on the current loop.
 
-        Always delegates to `connect()` rather than pre-checking `self._ws is
-        None`: `connect()` itself is the one that knows whether an existing
-        `_ws` is reusable (same event loop) or stale (a different one, e.g.
-        from a prior `from_env()` call now being driven through `.sync()`'s
-        own loop). A pre-check here that only looked at `_ws is None` would
-        skip `connect()` entirely whenever `_ws` is already set -- including
-        the stale-loop case -- so the reconnect logic would never run for
-        callers that never explicitly call `.connect()` themselves (e.g.
-        `client.sync().reset()` right after `from_env()`).
+        Always delegates to `_connect_async()` rather than pre-checking
+        `self._ws is None`: `_connect_async()` itself is the one that knows
+        whether an existing `_ws` is reusable (same event loop) or stale (a
+        different one, e.g. from a prior `from_env()` call now being driven
+        through `.sync()`'s own loop).
         """
-        await self.connect()
+        await self._connect_async()
 
     async def _send(self, message: Dict[str, Any]) -> None:
         """Send a message over the WebSocket."""
@@ -570,7 +659,10 @@ def _parse_state(self, payload: Dict[str, Any]) -> StateT:
         """Convert a JSON response from the state endpoint to a State object."""
         raise NotImplementedError
 
-    async def reset(self, **kwargs: Any) -> StepResult[ObsT]:
+    def reset(self, **kwargs: Any) -> Any:
+        return self._dispatch(lambda: self._reset_async(**kwargs))
+
+    async def _reset_async(self, **kwargs: Any) -> StepResult[ObsT]:
         """
         Reset the environment with optional parameters.
 
@@ -588,7 +680,10 @@ async def reset(self, **kwargs: Any) -> StepResult[ObsT]:
         response = await self._send_and_receive(message)
         return self._parse_result(response.get("data", {}))
 
-    async def step(self, action: ActT, **kwargs: Any) -> StepResult[ObsT]:
+    def step(self, action: ActT, **kwargs: Any) -> Any:
+        return self._dispatch(lambda: self._step_async(action, **kwargs))
+
+    async def _step_async(self, action: ActT, **kwargs: Any) -> StepResult[ObsT]:
         """
         Execute an action in the environment.
 
@@ -608,7 +703,10 @@ async def step(self, action: ActT, **kwargs: Any) -> StepResult[ObsT]:
         response = await self._send_and_receive(message)
         return self._parse_result(response.get("data", {}))
 
-    async def state(self) -> StateT:
+    def state(self) -> Any:
+        return self._dispatch(self._state_async)
+
+    async def _state_async(self) -> StateT:
         """
         Get the current environment state from the server.
 
@@ -619,7 +717,10 @@ async def state(self) -> StateT:
         response = await self._send_and_receive(message)
         return self._parse_state(response.get("data", {}))
 
-    async def close(self) -> None:
+    def close(self) -> Any:
+        return self._dispatch(self._close_async)
+
+    async def _close_async(self) -> None:
         """
         Close the WebSocket connection and clean up resources.
 
@@ -632,7 +733,7 @@ async def close(self) -> None:
         self._child_clients.clear()
 
         try:
-            await self.disconnect()
+            await self._disconnect_async()
         finally:
             try:
                 if self._provider is not None:
@@ -648,6 +749,7 @@ async def close(self) -> None:
 
     async def __aenter__(self) -> "EnvClient":
         """Enter async context manager, ensuring connection is established."""
+        self._claim_execution_mode("async")
         await self.connect()
         return self
 
@@ -656,17 +758,13 @@ async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
         await self.close()
 
     def __enter__(self) -> "EnvClient":
-        """Sync context manager entry - raises error suggesting async usage."""
-        raise TypeError(
-            "EnvClient is async by default. Use 'async with' instead of 'with', "
-            "or call .sync() to get a synchronous wrapper:\n"
-            "  async with client:  # async usage\n"
-            "  with client.sync():  # sync wrapper"
-        )
+        """Enter sync context manager, ensuring connection is established."""
+        self._run_sync(self._connect_async)
+        return self
 
     def __exit__(self, exc_type, exc_val, exc_tb) -> None:
-        """Sync context manager exit - should not be reached."""
-        pass  # pragma: no cover
+        """Exit sync context manager, closing connection."""
+        self._run_sync(self._close_async)
 
     def sync(self) -> "SyncEnvClient":
         """
@@ -694,4 +792,6 @@ def sync(self) -> "SyncEnvClient":
         """
         from .sync_client import SyncEnvClient
 
-        return SyncEnvClient(self)
+        if self._sync_client is None:
+            self._sync_client = SyncEnvClient(self)
+        return self._sync_client

src/openenv/core/sync_client.py

@@ -127,8 +127,13 @@ def _ensure_loop(self) -> asyncio.AbstractEventLoop:
             assert self._loop is not None
             return self._loop
 
+    def _claim_sync_mode(self) -> None:
+        if hasattr(self._async, "_claim_execution_mode"):
+            self._async._claim_execution_mode("sync")
+
     def _run(self, coro: Any) -> Any:
         """Run coroutine on dedicated loop and block for result."""
+        self._claim_sync_mode()
         loop = self._ensure_loop()
         future: concurrent.futures.Future[Any] = asyncio.run_coroutine_threadsafe(
             coro, loop
@@ -162,12 +167,14 @@ def connect(self) -> "SyncEnvClient[ActT, ObsT, StateT]":
         Returns:
             self for method chaining
         """
-        self._run(self._async.connect())
+        self._claim_sync_mode()
+        self._run(self._async._connect_async())
         return self
 
     def disconnect(self) -> None:
         """Close the connection."""
-        self._run(self._async.disconnect())
+        self._claim_sync_mode()
+        self._run(self._async._disconnect_async())
 
     def reset(self, **kwargs: Any) -> StepResult[ObsT]:
         """
@@ -180,7 +187,8 @@ def reset(self, **kwargs: Any) -> StepResult[ObsT]:
         Returns:
             StepResult containing initial observation
         """
-        return self._run(self._async.reset(**kwargs))
+        self._claim_sync_mode()
+        return self._run(self._async._reset_async(**kwargs))
 
     def step(self, action: ActT, **kwargs: Any) -> StepResult[ObsT]:
         """
@@ -195,7 +203,8 @@ def step(self, action: ActT, **kwargs: Any) -> StepResult[ObsT]:
         Returns:
             StepResult containing observation, reward, and done status
         """
-        return self._run(self._async.step(action, **kwargs))
+        self._claim_sync_mode()
+        return self._run(self._async._step_async(action, **kwargs))
 
     def state(self) -> StateT:
         """
@@ -204,7 +213,8 @@ def state(self) -> StateT:
         Returns:
             State object with environment state information
         """
-        return self._run(self._async.state())
+        self._claim_sync_mode()
+        return self._run(self._async._state_async())
 
     def close(self) -> None:
         """Close the connection and clean up resources."""
@@ -213,7 +223,8 @@ def close(self) -> None:
                 with suppress(Exception):
                     child.close()
             self._child_clients.clear()
-            self._run(self._async.close())
+            self._claim_sync_mode()
+            self._run(self._async._close_async())
         finally:
             self._stop_loop()
 

tests/envs/test_repl_env.py

@@ -964,7 +964,7 @@ async def fake_send_and_receive(message):
                 return {"data": {}}
             raise AssertionError(f"Unexpected message type: {message['type']}")
 
-        monkeypatch.setattr(env.async_client, "connect", fake_connect)
+        monkeypatch.setattr(env.async_client, "_connect_async", fake_connect)
         monkeypatch.setattr(
             env.async_client, "_send_and_receive", fake_send_and_receive
         )