fix(mcp): track fire-and-forget background tasks; harden loop teardown

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
This commit is contained in:
Patrick Buckley
2026-06-11 17:02:15 -07:00
parent 5ff726dd7a
commit 6c48af1900
2 changed files with 121 additions and 11 deletions
+58 -7
View File
@@ -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)
+63 -4
View File
@@ -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,
)