diff --git a/src/openenv/core/harness/bridge.py b/src/openenv/core/harness/bridge.py index 5bef30739..7167690ac 100644 --- a/src/openenv/core/harness/bridge.py +++ b/src/openenv/core/harness/bridge.py @@ -149,6 +149,8 @@ def start(self, timeout_s: float = 10.0) -> str: If the server fails to start within the timeout. """ if self._thread is not None and self._thread.is_alive(): + if self._uvicorn_server.should_exit: + raise HarnessError("MCP bridge is still stopping") assert self._url is not None return self._url @@ -202,13 +204,29 @@ def stop(self, timeout_s: float = 5.0) -> None: Args: timeout_s (`float`, *optional*, defaults to `5.0`): Maximum time to wait for the server thread to exit. + + Raises: + [`~openenv.core.harness.adapter.HarnessError`]: + If shutdown times out. The server remains tracked so callers + can retry `stop()`; `start()` rejects it while still stopping. """ server = self._uvicorn_server thread = self._thread + deadline = time.monotonic() + max(timeout_s, 0.0) if server is not None: + # Cancel lingering HTTP streams before the thread join expires. + # Reserve time for ASGI lifespan cleanup and forced shutdown. + server.config.timeout_graceful_shutdown = max(timeout_s, 0.0) / 4 server.should_exit = True if thread is not None and thread.is_alive(): - thread.join(timeout=timeout_s) + thread.join(timeout=max(timeout_s, 0.0) / 2) + if thread.is_alive() and server is not None: + server.force_exit = True + thread.join(timeout=max(0.0, deadline - time.monotonic())) + if thread.is_alive(): + # Keep ownership so callers can retry, and start() cannot + # replace a server that still holds resources. + raise HarnessError(f"MCP bridge did not stop within {timeout_s}s") self._teardown() def _teardown(self) -> None: diff --git a/src/openenv/core/harness/environment.py b/src/openenv/core/harness/environment.py index 6ce2f3850..437178ee5 100644 --- a/src/openenv/core/harness/environment.py +++ b/src/openenv/core/harness/environment.py @@ -137,9 +137,16 @@ async def reset_async( # died on its own still holds reapable resources (pipes, reader threads, # an unwaited process) that is_alive() reports nothing about. bridge_start: Optional[asyncio.Task[str]] = None + bridge_stop_failed = False try: await self.adapter.stop() - await self._stop_bridge() + if self._bridge is not None: + # A failed shutdown must not be replaced by another live bridge. + try: + await asyncio.to_thread(self._bridge.stop) + except Exception: + bridge_stop_failed = True + raise tools = await self._collect_injectable_tools() resolved = resolve_tool_conflicts(tools, self.adapter.BUILTIN_TOOL_NAMES) @@ -182,7 +189,10 @@ async def reset_async( ) return self._apply_transform(observation) except BaseException: - await self._cleanup_episode(bridge_start) + # The adapter is already stopped when bridge shutdown fails. Keep + # the bridge for a later retry without spending its timeout twice. + if not bridge_stop_failed: + await self._cleanup_episode(bridge_start) raise def reset( @@ -403,7 +413,7 @@ async def _stop_bridge(self) -> None: def close(self) -> None: """Stop the harness process and release environment resources.""" - if self._closed: + if self._closed and self._bridge is None: return self._closed = True self._episode_active = False @@ -416,7 +426,8 @@ def close(self) -> None: self._bridge.stop() except Exception: pass - self._bridge = None + else: + self._bridge = None super().close() diff --git a/tests/core/test_agentic_harness_bridge.py b/tests/core/test_agentic_harness_bridge.py index 155ff066d..d20bbc927 100644 --- a/tests/core/test_agentic_harness_bridge.py +++ b/tests/core/test_agentic_harness_bridge.py @@ -6,6 +6,7 @@ import asyncio import threading +import time from contextlib import asynccontextmanager from typing import AsyncIterator, Optional @@ -19,6 +20,7 @@ AgenticHarnessAdapter, HarnessConfig, HarnessEnvironment, + HarnessError, HarnessEvent, HarnessEventType, HarnessMCPBridge, @@ -61,6 +63,72 @@ async def is_alive(self) -> bool: class TestBridgeStandalone: + async def test_stop_closes_active_http_stream(self): + from types import SimpleNamespace + + import httpx + from starlette.applications import Starlette + from starlette.responses import StreamingResponse + from starlette.routing import Route + + finished = threading.Event() + + async def stream(request): + async def chunks(): + try: + yield b"ready\n" + await asyncio.Event().wait() + finally: + finished.set() + + return StreamingResponse(chunks()) + + app = Starlette(routes=[Route("/mcp", stream)]) + bridge = HarnessMCPBridge(SimpleNamespace(http_app=lambda: app)) + url = bridge.start() + thread = bridge._thread + try: + async with httpx.AsyncClient() as client: + async with client.stream("GET", url) as response: + assert await anext(response.aiter_lines()) == "ready" + started = time.monotonic() + await asyncio.to_thread(bridge.stop, 2.0) + assert time.monotonic() - started < 3.0 + assert not thread.is_alive() + assert finished.is_set() + assert bridge.url is None + bridge.start() + assert bridge._thread is not thread + finally: + await asyncio.to_thread(bridge.stop) + + def test_failed_stop_keeps_server_until_thread_exits(self): + from types import SimpleNamespace + + release = threading.Event() + thread = threading.Thread(target=release.wait, daemon=True) + server = SimpleNamespace( + should_exit=False, force_exit=False, config=SimpleNamespace() + ) + bridge = HarnessMCPBridge(make_mcp()) + bridge._thread = thread + bridge._uvicorn_server = server + bridge._url = "http://127.0.0.1:9/mcp" + thread.start() + try: + with pytest.raises(HarnessError, match="stop"): + bridge.stop(timeout_s=0.05) + assert bridge._thread is thread + assert bridge._uvicorn_server is server + assert server.force_exit + with pytest.raises(HarnessError, match="stopping"): + bridge.start() + finally: + release.set() + thread.join(timeout=1.0) + bridge.stop() + assert bridge.url is None + async def test_serves_tools_over_http(self): bridge = HarnessMCPBridge(make_mcp()) url = bridge.start() diff --git a/tests/core/test_agentic_harness_environment.py b/tests/core/test_agentic_harness_environment.py index c66d6217a..3cea46273 100644 --- a/tests/core/test_agentic_harness_environment.py +++ b/tests/core/test_agentic_harness_environment.py @@ -100,6 +100,7 @@ def __init__(self, mcp_server): self.mcp_server = mcp_server self.started = 0 self.stopped = 0 + self.fail_on_stop = False FakeBridge.instances.append(self) def start(self, timeout_s: float = 10.0) -> str: @@ -108,6 +109,8 @@ def start(self, timeout_s: float = 10.0) -> str: def stop(self, timeout_s: float = 5.0) -> None: self.stopped += 1 + if self.fail_on_stop: + raise HarnessError("bridge shutdown timed out") @pytest.fixture(autouse=True) @@ -174,6 +177,55 @@ async def test_reset_stops_live_adapter_first(self): await env.reset_async() assert adapter.calls == ["stop", "inject_tools", "start"] + async def test_reset_does_not_replace_bridge_after_shutdown_failure(self): + env, adapter = make_env() + await env.reset_async() + bridge = env._bridge + bridge.fail_on_stop = True + stop_count = bridge.stopped + adapter.calls.clear() + + with pytest.raises(HarnessError, match="bridge shutdown timed out"): + await env.reset_async() + + assert env._bridge is bridge + assert bridge.stopped == stop_count + 1 + assert FakeBridge.instances == [bridge] + assert "inject_tools" not in adapter.calls + assert "start" not in adapter.calls + assert adapter.alive is False + with pytest.raises(HarnessNotRunningError): + await env.step_async(HarnessAction(message="go")) + + bridge.fail_on_stop = False + await env.reset_async() + assert env._bridge is not bridge + assert adapter.alive is True + env.close() + + async def test_cancelled_bridge_stop_still_cleans_up(self, monkeypatch): + env, adapter = make_env() + await env.reset_async() + bridge = env._bridge + stop = bridge.stop + + def cancel_once(): + stop() + if bridge.stopped == 1: + raise asyncio.CancelledError() + + monkeypatch.setattr(bridge, "stop", cancel_once) + with pytest.raises(asyncio.CancelledError): + await env.reset_async() + + assert bridge.stopped == 2 + assert env._bridge is bridge + assert FakeBridge.instances == [bridge] + assert adapter.alive is False + with pytest.raises(HarnessNotRunningError): + await env.step_async(HarnessAction(message="go")) + env.close() + async def test_reset_clears_trajectory_and_resets_rubric(self): rubric = SpyRubric() env, adapter = make_env(rubric=rubric) @@ -350,6 +402,23 @@ def test_close_is_idempotent_and_stops_adapter_and_bridge(self): env.close() assert adapter.calls.count("stop") == stop_count + def test_close_retries_bridge_shutdown_after_failure(self): + env, _ = make_env() + env.reset() + bridge = env._bridge + bridge.fail_on_stop = True + + env.close() + + assert env._bridge is bridge + stop_count = bridge.stopped + bridge.fail_on_stop = False + env.close() + assert bridge.stopped == stop_count + 1 + assert env._bridge is None + env.close() + assert bridge.stopped == stop_count + 1 + def test_constructor_starts_nothing(self): env, adapter = make_env() assert adapter.calls == []