From 3f5ee333fbcc4c024f3dcdab92139b91264059da Mon Sep 17 00:00:00 2001
From: Patrick Buckley <eous@users.noreply.github.com>
Date: Thu, 11 Jun 2026 18:07:47 -0700
Subject: [PATCH] fix(mcp): close the shutdown drain race + close the owned
 loop
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Review feedback: (1) gating the drain on a main-thread truthiness check
of _background_tasks could skip cancellation when a spawn queued via
call_soon_threadsafe had not reached the set yet — submit whenever the
loop is RUNNING and snapshot on the loop, where FIFO callback order
guarantees earlier-queued spawns have landed; (2) shutdown stopped the
loop thread but never closed the loop or cleared _loop/_thread, leaking
selector resources for embedders that cycle managers — close + clear
when we own the thread and it actually stopped (loud warning when it
does not); unowned loops (tests wiring _loop directly) stay untouched;
(3) the bare await-in-suppress drain loops become
asyncio.gather(return_exceptions=True) in both the shutdown drain and
the test fixture.
---
 tests/test_mcp_client.py     | 35 ++++++++++++++++++++++++++++++++---
 turnstone/core/mcp_client.py | 27 ++++++++++++++++++++++-----
 2 files changed, 54 insertions(+), 8 deletions(-)

diff --git a/tests/test_mcp_client.py b/tests/test_mcp_client.py
index 81297d34..ce5909ae 100644
--- a/tests/test_mcp_client.py
+++ b/tests/test_mcp_client.py
@@ -163,9 +163,7 @@ def running_loop_mgr():
             tasks = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()]
             for t in tasks:
                 t.cancel()
-            for t in tasks:
-                with suppress(BaseException):
-                    await t
+            await asyncio.gather(*tasks, return_exceptions=True)
 
         with suppress(Exception):
             asyncio.run_coroutine_threadsafe(_cancel_pending(), loop).result(timeout=5)
@@ -2042,6 +2040,37 @@ class TestShutdownCleanup:
         assert mgr._resource_map == {}
         assert mgr._prompt_map == {}
 
+    def test_shutdown_closes_owned_loop_and_clears_refs(self):
+        """When the manager owns the loop thread, shutdown must close the loop
+        (selector resources leak otherwise) and drop both refs; a second
+        shutdown is then a clean no-op."""
+        import threading as _threading
+
+        mgr = MCPClientManager({})
+        loop = asyncio.new_event_loop()
+        thread = _threading.Thread(target=loop.run_forever, daemon=True)
+        thread.start()
+        mgr._loop = loop
+        mgr._thread = thread
+
+        mgr.shutdown()
+        assert loop.is_closed()
+        assert mgr._loop is None
+        assert mgr._thread is None
+        mgr.shutdown()  # idempotent
+
+    def test_shutdown_leaves_unowned_loop_open(self):
+        """Tests (and any embedder) that wire ``_loop`` directly without a
+        thread own the loop's lifecycle — shutdown must not close it."""
+        mgr = MCPClientManager({})
+        loop = asyncio.new_event_loop()
+        mgr._loop = loop
+        try:
+            mgr.shutdown()
+            assert not loop.is_closed()
+        finally:
+            loop.close()
+
 
 # ---------------------------------------------------------------------------
 # TCP probe and unreachable server handling
diff --git a/turnstone/core/mcp_client.py b/turnstone/core/mcp_client.py
index fe594b03..dcd22cf6 100644
--- a/turnstone/core/mcp_client.py
+++ b/turnstone/core/mcp_client.py
@@ -2721,16 +2721,21 @@ class MCPClientManager:
         """Close all MCP sessions and stop the background loop."""
         # Cancel tracked background tasks (catalog refreshes etc.) FIRST —
         # they are pure auxiliaries, and draining them up front means the
-        # stack teardown below can't race an in-flight refresh.
-        if self._loop and self._background_tasks:
+        # stack teardown below can't race an in-flight refresh.  Submitted
+        # whenever a loop exists — NOT gated on a main-thread truthiness
+        # check of ``_background_tasks``: a spawn queued via
+        # call_soon_threadsafe may not have reached the set yet, but ready
+        # callbacks run in FIFO order, so by the time the drain coroutine
+        # snapshots the set ON the loop, every earlier-queued spawn has
+        # landed.  ``is_running()`` guard: on a stopped loop nothing can
+        # execute the drain — submitting would just stall on the future.
+        if self._loop is not None and self._loop.is_running():
 
             async def _cancel_background() -> None:
                 tasks = list(self._background_tasks)
                 for task in tasks:
                     task.cancel()
-                for task in tasks:
-                    with contextlib.suppress(BaseException):
-                        await task
+                await asyncio.gather(*tasks, return_exceptions=True)
 
             future = asyncio.run_coroutine_threadsafe(_cancel_background(), self._loop)
             try:
@@ -2793,6 +2798,18 @@ class MCPClientManager:
             self._loop.call_soon_threadsafe(self._loop.stop)
         if self._thread:
             self._thread.join(timeout=5)
+            if self._thread.is_alive():
+                # Closing a still-running loop raises; the daemon thread dies
+                # with the process, so leaving the loop open is the lesser
+                # evil.  Loud because a stuck loop thread is itself a bug.
+                log.warning("MCP loop thread did not stop within 5s; loop left open")
+            else:
+                if self._loop is not None:
+                    self._loop.close()
+                self._loop = None
+                self._thread = None
+        # When no thread was started (tests wire ``_loop`` directly) the loop
+        # is not ours to close — the stop above is all the owner needs.
 
         # Clear all state
         self._background_tasks.clear()