fix(mcp): close the shutdown drain race + close the owned loop

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.
This commit is contained in:
Patrick Buckley
2026-06-11 18:07:47 -07:00
parent 6c48af1900
commit 3f5ee333fb
2 changed files with 54 additions and 8 deletions
+32 -3
View File
@@ -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
+22 -5
View File
@@ -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()