From ab0d3fac3d4cde14333d4107d8b5ae1475399a04 Mon Sep 17 00:00:00 2001 From: Seppe Gadeyne Date: Fri, 17 Jul 2026 10:18:01 +0200 Subject: [PATCH] fix(mcp): keep drain and stop on loop thread --- tests/tools/test_mcp_stability.py | 81 ++++++++++++++++++++++--------- tools/mcp_tool.py | 31 ++++++++++-- 2 files changed, 86 insertions(+), 26 deletions(-) diff --git a/tests/tools/test_mcp_stability.py b/tests/tools/test_mcp_stability.py index 712a80c04eb..cf43a0d5785 100644 --- a/tests/tools/test_mcp_stability.py +++ b/tests/tools/test_mcp_stability.py @@ -796,20 +796,32 @@ class TestMCPLoopDrainOnStop: assert state["started"], "task never started on the MCP loop" stop_saw_cleanup = [] + call_soon = loop.call_soon call_soon_threadsafe = loop.call_soon_threadsafe - def record_stop_order(callback, *args, **kwargs): + def record_stop_order(schedule, callback, *args, **kwargs): if ( getattr(callback, "__self__", None) is loop and getattr(callback, "__name__", None) == "stop" ): stop_saw_cleanup.append(state["cleanup_ran"]) - return call_soon_threadsafe(callback, *args, **kwargs) + return schedule(callback, *args, **kwargs) - with patch.object( - loop, - "call_soon_threadsafe", - side_effect=record_stop_order, + with ( + patch.object( + loop, + "call_soon", + side_effect=lambda callback, *args, **kwargs: record_stop_order( + call_soon, callback, *args, **kwargs + ), + ), + patch.object( + loop, + "call_soon_threadsafe", + side_effect=lambda callback, *args, **kwargs: record_stop_order( + call_soon_threadsafe, callback, *args, **kwargs + ), + ), ): mcp_mod._stop_mcp_loop() @@ -843,7 +855,8 @@ class TestMCPLoopDrainOnStop: await asyncio.sleep(0) try: with caplog.at_level("WARNING", logger=mcp_mod.logger.name): - await mcp_mod._drain_mcp_loop_tasks(timeout=0.01) + async with asyncio.timeout(0.5): + await mcp_mod._drain_mcp_loop_tasks(timeout=0.01) assert not task.done(), "drain waited indefinitely for resistant task" finally: release.set() @@ -858,34 +871,56 @@ class TestMCPLoopDrainOnStop: for record in caplog.records ) - def test_stop_warns_when_drain_wait_times_out(self, caplog): - """A failed final drain must be visible instead of silently no-oping.""" + def test_outer_timeout_still_allows_loop_owned_drain_before_stop( + self, caplog, monkeypatch + ): + """A blocked loop must drain after it resumes, not stop ahead of the drain.""" + import threading import tools.mcp_tool as mcp_mod - class TimedOutFuture: - def result(self, timeout): - assert timeout > 0 - raise TimeoutError("simulated drain timeout") + parked_started = threading.Event() + cleanup_ran = threading.Event() + blocker_started = threading.Event() + release_blocker = threading.Event() - def cancel(self): - return True + async def parked_task(): + parked_started.set() + try: + await asyncio.Event().wait() + finally: + cleanup_ran.set() - def report_timeout(coro, _loop, **_kwargs): - coro.close() - return TimedOutFuture() + def block_loop(): + blocker_started.set() + release_blocker.wait(timeout=5) with mcp_mod._lock: mcp_mod._servers.clear() mcp_mod._server_connecting.clear() mcp_mod._ensure_mcp_loop() + with mcp_mod._lock: + loop = mcp_mod._mcp_loop + assert loop is not None + + future = asyncio.run_coroutine_threadsafe(parked_task(), loop) + assert parked_started.wait(timeout=2) + loop.call_soon_threadsafe(block_loop) + assert blocker_started.wait(timeout=2) + + monkeypatch.setattr(mcp_mod, "_MCP_LOOP_DRAIN_TIMEOUT", 0.01) + release_timer = threading.Timer(1.2, release_blocker.set) + release_timer.start() try: with caplog.at_level("WARNING", logger=mcp_mod.logger.name): - with patch( - "agent.async_utils.safe_schedule_threadsafe", - side_effect=report_timeout, - ): - mcp_mod._stop_mcp_loop() + mcp_mod._stop_mcp_loop() + + assert cleanup_ran.is_set(), "drain was overtaken by loop.stop" + assert future.done(), "parked task remained pending after loop resumed" + assert loop.is_closed() finally: + release_timer.cancel() + release_blocker.set() + future.cancel() with mcp_mod._lock: mcp_mod._servers.clear() mcp_mod._server_connecting.clear() diff --git a/tools/mcp_tool.py b/tools/mcp_tool.py index a5eb720a24a..dd3f353a906 100644 --- a/tools/mcp_tool.py +++ b/tools/mcp_tool.py @@ -6660,6 +6660,21 @@ async def _drain_mcp_loop_tasks( ) +async def _drain_and_stop_mcp_loop() -> None: + """Drain pending tasks, then stop the loop from its owning thread. + + Keeping both operations in one loop-owned sequence matters when the caller + times out waiting for a blocked loop. Queuing ``loop.stop`` separately from + the caller can overtake the scheduled drain before it receives a loop cycle, + leaving the drain coroutine itself pending when the loop is closed. + """ + loop = asyncio.get_running_loop() + try: + await _drain_mcp_loop_tasks(timeout=_MCP_LOOP_DRAIN_TIMEOUT) + finally: + loop.call_soon(loop.stop) + + def _stop_mcp_loop(*, only_if_idle: bool = False) -> bool: """Stop the background event loop and join its thread.""" global _mcp_loop, _mcp_thread @@ -6677,27 +6692,37 @@ def _stop_mcp_loop(*, only_if_idle: bool = False) -> bool: # to run cleanup against a loop that is already closed -> "Event loop # is closed" (#60197). ``shutdown_mcp_servers`` only reaps servers held # in ``_servers``, so anything else left on this loop ends up here. + stop_owned_by_loop = False if loop.is_running(): from agent.async_utils import safe_schedule_threadsafe future = safe_schedule_threadsafe( - _drain_mcp_loop_tasks(timeout=_MCP_LOOP_DRAIN_TIMEOUT), loop, + _drain_and_stop_mcp_loop(), loop, logger=logger, log_message="MCP loop drain: failed to schedule", log_level=logging.WARNING, ) if future is not None: + stop_owned_by_loop = True try: future.result(timeout=_MCP_LOOP_DRAIN_TIMEOUT + 1) except TimeoutError: - future.cancel() logger.warning( "Timed out waiting for MCP loop drain after %.1fs", _MCP_LOOP_DRAIN_TIMEOUT + 1, ) except BaseException as exc: logger.warning("Error draining MCP loop tasks: %s", exc) - loop.call_soon_threadsafe(loop.stop) + elif not loop.is_closed(): + try: + loop.run_until_complete( + _drain_mcp_loop_tasks(timeout=_MCP_LOOP_DRAIN_TIMEOUT) + ) + except BaseException as exc: + logger.warning("Error draining stopped MCP loop tasks: %s", exc) + + if not stop_owned_by_loop and loop.is_running(): + loop.call_soon_threadsafe(loop.stop) if thread is not None: thread.join(timeout=5) if thread.is_alive():