From 6c48af19006f15050ca71b42e543fc73f5cfb77c Mon Sep 17 00:00:00 2001 From: Patrick Buckley Date: Thu, 11 Jun 2026 17:02:15 -0700 Subject: [PATCH] fix(mcp): track fire-and-forget background tasks; harden loop teardown MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The post-reconnect catalog refresh was scheduled as a bare asyncio.create_task: no strong reference (the task could be GC'd mid-flight, so the refresh might silently never run) and no exception retrieval (failures surfaced as "Task exception was never retrieved" at GC time — in CI, onto an already-closed pytest capture stream, the "I/O operation on closed file" spew; a suspected contributor to the flaky 60-minute CI hangs via cross-test loop/task state bleed). - _spawn_background(coro, label): tracked-task set + done-callback that retrieves and logs failures at warning; discard runs LAST so set-emptiness means "done AND reported" - shutdown() drains tracked tasks FIRST, so stack teardown can't race an in-flight refresh; same run_coroutine_threadsafe idiom and timeouts as the existing close steps - running_loop_mgr fixture: cancel-pending -> drain -> stop -> join(5) with a loud assert -> loop.close() (was stop + silent join(2), never closed) - the false-property test ("swallows refresh failure" — nothing swallowed it) now waits for completion and asserts the logged warning via the patched module logger (structlog; caplog cannot observe it), polling inside the patch context --- tests/test_mcp_client.py | 65 ++++++++++++++++++++++++++++++---- turnstone/core/mcp_client.py | 67 +++++++++++++++++++++++++++++++++--- 2 files changed, 121 insertions(+), 11 deletions(-) diff --git a/tests/test_mcp_client.py b/tests/test_mcp_client.py index f43ab42b..81297d34 100644 --- a/tests/test_mcp_client.py +++ b/tests/test_mcp_client.py @@ -6,7 +6,7 @@ import asyncio import concurrent.futures import json import time -from contextlib import AsyncExitStack +from contextlib import AsyncExitStack, suppress from typing import Any from unittest.mock import AsyncMock, MagicMock, patch @@ -153,8 +153,26 @@ def running_loop_mgr(): try: yield mgr, loop, thread finally: + # Drain BEFORE stopping: a task left pending (or finished-but- + # unretrieved) on a stopped loop becomes cross-test global state — + # asyncio reports it at GC time, mid-suite, onto whatever stream + # pytest has attached THEN (the "I/O operation on closed file" + # spew), and a silently-abandoned loop thread keeps running + # manager code against torn-down mocks. + async def _cancel_pending() -> None: + 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 + + with suppress(Exception): + asyncio.run_coroutine_threadsafe(_cancel_pending(), loop).result(timeout=5) loop.call_soon_threadsafe(loop.stop) - thread.join(timeout=2) + thread.join(timeout=5) + assert not thread.is_alive(), "mcp test loop thread failed to stop within 5s" + loop.close() # --------------------------------------------------------------------------- @@ -2706,11 +2724,22 @@ class TestCBAutoReconnectRefresh: patch.object(mgr, "_refresh_server", side_effect=_refresh), ): session = mgr._cb_auto_reconnect("srv") - # Wait for the scheduled refresh task to actually run on the loop. + # Wait for the scheduled refresh task to actually run on the loop, + # then for the tracked task to DRAIN — exiting the patch context + # while the task is still in flight would hand the un-patched + # method to its tail. assert refresh_event.wait(timeout=5), "refresh task was not scheduled" + deadline = time.time() + 5 + while mgr._background_tasks and time.time() < deadline: + time.sleep(0.02) + assert not mgr._background_tasks, "background refresh task never drained" assert session is new_session - def test_auto_reconnect_swallows_refresh_failure(self, running_loop_mgr): + def test_auto_reconnect_retrieves_and_logs_refresh_failure(self, running_loop_mgr): + """A refresh failure must be RETRIEVED and logged by the task's + done-callback — not abandoned for asyncio to report as "Task exception + was never retrieved" at GC time (which lands on whatever stream pytest + has attached by then: the closed-file CI spew).""" import threading as _threading mgr, _loop, _thread = running_loop_mgr @@ -2728,10 +2757,32 @@ class TestCBAutoReconnectRefresh: with ( patch.object(mgr, "_connect_one", side_effect=_connect_one), patch.object(mgr, "_refresh_server", side_effect=_refresh_failing), + patch("turnstone.core.mcp_client.log") as mock_log, ): - # Must not raise — refresh failures are non-fatal. + # Must not raise — refresh failures are non-fatal to the caller. session = mgr._cb_auto_reconnect("srv") - # Background refresh actually started and exception was swallowed - # by the task without affecting the synchronous caller. assert refresh_started.wait(timeout=5), "refresh task was not scheduled" + # Poll for the WARNING while the patch is still active — gating on + # set-emptiness alone would race the un-patch (review-caught: the + # warning could land on the restored real logger). + deadline = time.time() + 5 + warn_calls = [] + while not warn_calls and time.time() < deadline: + warn_calls = [ + c for c in mock_log.warning.call_args_list if "MCP background" in str(c.args[0]) + ] + time.sleep(0.02) + # The tracked task must also fully drain (emptiness now implies + # "done AND reported" — discard is the callback's LAST step). + deadline = time.time() + 5 + while mgr._background_tasks and time.time() < deadline: + time.sleep(0.02) + assert not mgr._background_tasks, "background refresh task never drained" assert session is new_session + assert warn_calls, ( + "the refresh failure must be logged by the done-callback, not left " + "for GC-time reporting" + ) + exc = warn_calls[0].kwargs.get("exc_info") + assert isinstance(exc, RuntimeError) + assert "catalog fetch broke" in str(exc) diff --git a/turnstone/core/mcp_client.py b/turnstone/core/mcp_client.py index 86573473..fe594b03 100644 --- a/turnstone/core/mcp_client.py +++ b/turnstone/core/mcp_client.py @@ -57,7 +57,7 @@ from turnstone.core.mcp_oauth import ( ) if TYPE_CHECKING: - from collections.abc import Awaitable, Callable + from collections.abc import Awaitable, Callable, Coroutine log = get_logger("turnstone.mcp") @@ -607,6 +607,13 @@ class MCPClientManager: # static-only deployments). self._user_pool_eviction_task: asyncio.Task[None] | None = None + # Strong references to fire-and-forget background tasks (catalog + # refreshes etc.). ``create_task`` alone keeps only a weak ref — an + # untracked task can be GC'd mid-flight, and its exception surfaces + # as "Task exception was never retrieved" at GC time instead of + # being logged where it happened. See ``_spawn_background``. + self._background_tasks: set[asyncio.Task[Any]] = set() + def _ensure_static_state(self, name: str) -> StaticServerState: """Get or create the StaticServerState for ``name``. @@ -2712,6 +2719,25 @@ class MCPClientManager: def shutdown(self) -> None: """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: + + 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 + + future = asyncio.run_coroutine_threadsafe(_cancel_background(), self._loop) + try: + future.result(timeout=10) + except Exception: + log.debug("Error cancelling MCP background tasks", exc_info=True) + # Cancel the pool eviction task, then close all pool entries # before tearing down static-path state. Both run on the # mcp-loop so dispatcher coroutines can't race them. @@ -2769,6 +2795,7 @@ class MCPClientManager: self._thread.join(timeout=5) # Clear all state + self._background_tasks.clear() self._static_servers.clear() self._db_managed.clear() self._tools = [] @@ -3327,6 +3354,33 @@ class MCPClientManager: # broken the first failure re-trips the circuit immediately. self._circuit_open_until.pop(server_name, None) + def _spawn_background(self, coro: Coroutine[Any, Any, Any], label: str) -> asyncio.Task[Any]: + """Schedule *coro* as a tracked background task (loop thread only). + + Holds a strong reference until completion and retrieves the task's + outcome in a done-callback: failures are logged once, here, at + warning — never deferred to garbage collection, where they surface + as "Task exception was never retrieved" on whatever stream happens + to be attached at the time. ``shutdown()`` cancels anything still + tracked before stopping the loop. + """ + task = asyncio.create_task(coro) + self._background_tasks.add(task) + + def _done(t: asyncio.Task[Any]) -> None: + try: + if not t.cancelled(): + exc = t.exception() + if exc is not None: + log.warning("MCP background %s failed", label, exc_info=exc) + finally: + # Discard LAST so set-emptiness means "done AND reported" — + # a watcher keying on emptiness must never race the warning. + self._background_tasks.discard(t) + + task.add_done_callback(_done) + return task + def _cb_auto_reconnect(self, server_name: str) -> Any: """Attempt reconnection for a disconnected server during half-open probe. @@ -3355,13 +3409,18 @@ class MCPClientManager: # Schedule catalog refresh on the loop without blocking the caller. # The reconnected session is valid for the imminent dispatch; catalog - # drift will be reconciled on the loop in the background. + # drift will be reconciled on the loop in the background. The task is + # tracked: a refresh FAILURE is logged by the done-callback — this + # except only covers the scheduling itself. def _schedule_refresh() -> None: try: - asyncio.create_task(self._refresh_server(server_name)) + self._spawn_background( + self._refresh_server(server_name), + f"catalog refresh after reconnect for '{server_name}'", + ) except Exception: log.warning( - "Catalog refresh after reconnect failed for '%s'", + "Scheduling catalog refresh after reconnect failed for '%s'", server_name, exc_info=True, )