From bbc0468e7782f498f85c817e6253ccc1b67c9b3e Mon Sep 17 00:00:00 2001 From: burtenshaw Date: Mon, 29 Jun 2026 10:17:15 +0200 Subject: [PATCH 1/3] fix: auto-dispatch env client sync calls --- src/openenv/core/env_client.py | 126 +++++++++++++++++++++++++++----- src/openenv/core/sync_client.py | 23 ++++-- 2 files changed, 123 insertions(+), 26 deletions(-) diff --git a/src/openenv/core/env_client.py b/src/openenv/core/env_client.py index 7bfd33be2..b73bebf54 100644 --- a/src/openenv/core/env_client.py +++ b/src/openenv/core/env_client.py @@ -41,7 +41,7 @@ import json import os from abc import ABC, abstractmethod -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 @@ -59,10 +59,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 = ( @@ -182,6 +221,8 @@ def __init__( ) # Convert MB to bytes self._provider = provider self._ws: Optional[ClientConnection] = None + self._execution_mode: Optional[str] = None + self._sync_client: Optional["SyncEnvClient[ActT, ObsT, StateT]"] = None def __setattr__(self, name: str, value: Any) -> None: """Prevent modification of _mode after initialization.""" @@ -189,7 +230,40 @@ 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]) -> Any: + """Run an async operation through the sync wrapper.""" + 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.""" + if self._execution_mode == "async": + return coro_factory() + if self._execution_mode == "sync": + return self._run_sync(coro_factory) + + try: + asyncio.get_running_loop() + except RuntimeError: + return self._run_sync(coro_factory) + return _AutoAsyncResult(self, coro_factory) + + def connect(self) -> Any: + return self._dispatch(self._connect_async) + + async def _connect_async(self) -> "EnvClient": """ Establish WebSocket connection to the server. @@ -222,7 +296,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: try: @@ -239,7 +316,7 @@ async def disconnect(self) -> None: async def _ensure_connected(self) -> None: """Ensure WebSocket connection is established.""" if self._ws is None: - await self.connect() + await self._connect_async() async def _send(self, message: Dict[str, Any]) -> None: """Send a message over the WebSocket.""" @@ -407,7 +484,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. @@ -425,7 +505,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. @@ -445,7 +528,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. @@ -456,14 +542,17 @@ 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. If this client was created via from_docker_image() or from_env(), this will also stop and remove the associated container/process. """ - await self.disconnect() + await self._disconnect_async() if self._provider is not None: # Handle both ContainerProvider and RuntimeProvider @@ -474,25 +563,22 @@ async def close(self) -> None: async def __aenter__(self) -> "EnvClient": """Enter async context manager, ensuring connection is established.""" - await self.connect() + self._claim_execution_mode("async") + await self._connect_async() return self async def __aexit__(self, exc_type, exc_val, exc_tb) -> None: """Exit async context manager, closing connection.""" - await self.close() + await self._close_async() 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": """ diff --git a/src/openenv/core/sync_client.py b/src/openenv/core/sync_client.py index c2bf24753..73af5dfca 100644 --- a/src/openenv/core/sync_client.py +++ b/src/openenv/core/sync_client.py @@ -129,8 +129,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 @@ -164,12 +169,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]: """ @@ -182,7 +189,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]: """ @@ -197,7 +205,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: """ @@ -206,12 +215,14 @@ 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.""" try: - self._run(self._async.close()) + self._claim_sync_mode() + self._run(self._async._close_async()) finally: self._stop_loop() From 216a3ba1281c693a58a75d195c993a9117875445 Mon Sep 17 00:00:00 2001 From: burtenshaw Date: Wed, 1 Jul 2026 11:08:12 +0200 Subject: [PATCH 2/3] fix: preserve async dispatch in sync loop --- src/openenv/core/env_client.py | 22 ++++++++++++++------- tests/envs/test_repl_env.py | 2 +- tests/test_core/test_generic_client.py | 27 ++++++++++++++++---------- 3 files changed, 33 insertions(+), 18 deletions(-) diff --git a/src/openenv/core/env_client.py b/src/openenv/core/env_client.py index 30af83940..64ca7c0fd 100644 --- a/src/openenv/core/env_client.py +++ b/src/openenv/core/env_client.py @@ -259,16 +259,24 @@ def _run_sync(self, coro_factory: Callable[[], Any]) -> Any: def _dispatch(self, coro_factory: Callable[[], Any]) -> Any: """Return an awaitable in async code and a concrete result in sync code.""" + try: + asyncio.get_running_loop() + except RuntimeError: + running_loop = False + else: + running_loop = True + + if running_loop: + if self._execution_mode is None: + return _AutoAsyncResult(self, coro_factory) + return coro_factory() + if self._execution_mode == "async": return coro_factory() if self._execution_mode == "sync": return self._run_sync(coro_factory) - try: - asyncio.get_running_loop() - except RuntimeError: - return self._run_sync(coro_factory) - return _AutoAsyncResult(self, coro_factory) + return self._run_sync(coro_factory) def connect(self) -> Any: return self._dispatch(self._connect_async) @@ -618,12 +626,12 @@ async def _close_async(self) -> None: async def __aenter__(self) -> "EnvClient": """Enter async context manager, ensuring connection is established.""" self._claim_execution_mode("async") - await self._connect_async() + await self.connect() return self async def __aexit__(self, exc_type, exc_val, exc_tb) -> None: """Exit async context manager, closing connection.""" - await self._close_async() + await self.close() def __enter__(self) -> "EnvClient": """Enter sync context manager, ensuring connection is established.""" diff --git a/tests/envs/test_repl_env.py b/tests/envs/test_repl_env.py index b8cb8311e..861ff703a 100644 --- a/tests/envs/test_repl_env.py +++ b/tests/envs/test_repl_env.py @@ -968,7 +968,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 ) diff --git a/tests/test_core/test_generic_client.py b/tests/test_core/test_generic_client.py index 7011a303f..067179935 100644 --- a/tests/test_core/test_generic_client.py +++ b/tests/test_core/test_generic_client.py @@ -578,25 +578,32 @@ async def test_async_context_manager_enter_exit(self): mock_close.assert_called_once() - def test_sync_context_manager_raises_error(self): - """Test that sync context manager raises helpful error.""" - client = GenericEnvClient(base_url="http://localhost:8000") + def test_sync_context_manager_enter_exit(self): + """Test that sync context manager works correctly.""" + with ( + patch.object( + GenericEnvClient, "_connect_async", new_callable=AsyncMock + ) as mock_connect, + patch.object( + GenericEnvClient, "_close_async", new_callable=AsyncMock + ) as mock_close, + ): + client = GenericEnvClient(base_url="http://localhost:8000") - with pytest.raises(TypeError) as exc_info: - with client: - pass + with client as active_client: + assert active_client is client + mock_connect.assert_called_once() - assert "async by default" in str(exc_info.value) - assert ".sync()" in str(exc_info.value) + mock_close.assert_called_once() def test_sync_wrapper_context_manager(self): """Test SyncEnvClient context manager works correctly.""" with ( patch.object( - GenericEnvClient, "connect", new_callable=AsyncMock + GenericEnvClient, "_connect_async", new_callable=AsyncMock ) as mock_connect, patch.object( - GenericEnvClient, "close", new_callable=AsyncMock + GenericEnvClient, "_close_async", new_callable=AsyncMock ) as mock_close, ): async_client = GenericEnvClient(base_url="http://localhost:8000") From d8316ee97f522076bda266e61611794319f92042 Mon Sep 17 00:00:00 2001 From: burtenshaw Date: Thu, 2 Jul 2026 11:29:54 +0200 Subject: [PATCH 3/3] fix: repair env client dispatch modes --- src/openenv/core/env_client.py | 42 +++++++++------ tests/test_core/test_generic_client.py | 72 ++++++++++++++++++++++++++ 2 files changed, 98 insertions(+), 16 deletions(-) diff --git a/src/openenv/core/env_client.py b/src/openenv/core/env_client.py index e77386c57..d4d30877a 100644 --- a/src/openenv/core/env_client.py +++ b/src/openenv/core/env_client.py @@ -350,9 +350,17 @@ def _claim_execution_mode(self, mode: str) -> None: "Create a separate client instance when mixing sync and async code." ) - def _run_sync(self, coro_factory: Callable[[], Any]) -> Any: + def _run_sync( + self, + coro_factory: Callable[[], Any], + *, + allow_async_handoff: bool = False, + ) -> Any: """Run an async operation through the sync wrapper.""" - self._claim_execution_mode("sync") + 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()) @@ -360,23 +368,23 @@ def _run_sync(self, coro_factory: Callable[[], Any]) -> Any: def _dispatch(self, coro_factory: Callable[[], Any]) -> Any: """Return an awaitable in async code and a concrete result in sync code.""" try: - asyncio.get_running_loop() + running_loop = asyncio.get_running_loop() except RuntimeError: - running_loop = False - else: - running_loop = True - - if running_loop: - if self._execution_mode is None: - return _AutoAsyncResult(self, coro_factory) - return coro_factory() + running_loop = None - if self._execution_mode == "async": - return coro_factory() 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) + return self._run_sync(coro_factory, allow_async_handoff=True) def connect(self) -> Any: return self._dispatch(self._connect_async) @@ -729,7 +737,7 @@ async def _close_async(self) -> None: self._child_clients.clear() try: - await self.disconnect() + await self._disconnect_async() finally: try: if self._provider is not None: @@ -788,4 +796,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 diff --git a/tests/test_core/test_generic_client.py b/tests/test_core/test_generic_client.py index 85a545f41..5ed358651 100644 --- a/tests/test_core/test_generic_client.py +++ b/tests/test_core/test_generic_client.py @@ -692,6 +692,78 @@ def test_sync_close_stops_provider_when_child_close_raises(self, mock_provider): assert client._child_clients == [] mock_provider.stop_container.assert_called_once_with() + def test_sync_call_after_async_connect_blocks_for_result(self, monkeypatch): + """Sync calls after async setup return concrete results, not coroutines.""" + client = GenericEnvClient(base_url="http://localhost:8000") + + async def connect_client(): + with patch.object( + GenericEnvClient, "_connect_async", new_callable=AsyncMock + ) as mock_connect: + await client.connect() + mock_connect.assert_awaited_once() + + asyncio.run(connect_client()) + assert client._execution_mode == "async" + + expected = StepResult( + observation={"ready": True}, + reward=0.0, + done=False, + ) + + async def fake_reset(**kwargs): + return expected + + monkeypatch.setattr(client, "_reset_async", fake_reset) + + result = client.reset() + if asyncio.iscoroutine(result): + result.close() + pytest.fail("client.reset() returned a coroutine in synchronous code") + + assert result is expected + + @pytest.mark.asyncio + async def test_sync_locked_dispatch_inside_running_loop_uses_sync_loop( + self, monkeypatch + ): + """Sync-locked clients keep work on the sync wrapper loop.""" + client = GenericEnvClient(base_url="http://localhost:8000") + sync_client = client.sync() + + async def fake_connect(): + return client + + monkeypatch.setattr(client, "_connect_async", fake_connect) + sync_client.connect() + + expected = StepResult( + observation={"ready": True}, + reward=0.0, + done=False, + ) + observed = {} + + async def fake_reset(**kwargs): + observed["loop"] = asyncio.get_running_loop() + return expected + + monkeypatch.setattr(client, "_reset_async", fake_reset) + + try: + current_loop = asyncio.get_running_loop() + result = client.reset() + if asyncio.iscoroutine(result): + result.close() + pytest.fail("sync-locked client.reset() returned a coroutine") + + assert result is expected + assert observed["loop"] is sync_client._loop + assert observed["loop"] is not current_loop + finally: + sync_client.close() + # ============================================================================ # Context Manager Tests