mirror of
https://github.com/turnstonelabs/turnstone.git
synced 2026-08-27 22:34:51 -06:00
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:
@@ -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)
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user