fix(mcp): keep drain and stop on loop thread

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()

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():