mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-26 13:16:24 +00:00
fix(mcp): tear down the in-flight owner when get_session is cancelled mid-eviction (#5008)
* fix(mcp): keep session owner teardown safe across cancellation paths * fix(mcp): gate pooled-session publication on the commit inside the owner task The owner resolved `ready` as soon as initialize() finished, but the session only became pool property when the creator promoted it into _entries in Phase 4. A concurrent get_session() could join the in-flight creation and receive that session from Phase 2b while the creator was still parked in the Phase-2 eviction teardown; cancelling the creator then ran the Phase-2 unwind, which unconditionally shut the owner down — closing the session underneath the joiner (#5008 review). Move the commit into the owner task: initialize() success now pops the in-flight record, registers the session in _entries, and resolves ready with the session in one atomic critical section, so 'ready resolved with a result' is exactly 'session registered and pool-owned'. Joiners can therefore only ever receive a committed session, and both creator unwind paths (Phase 2 and Phase 3) skip teardown when the creation already committed, leaving the pooled session to LRU eviction / close_*. When the record was removed before the commit (close_* or creator unwind), the owner aborts and ready carries the same cancellation the old Phase-4 not-still-ours path raised, so joiners fail with the creation's outcome instead of hanging or holding an unmanaged session. test_cancelled_creator_does_not_close_session_held_by_joiner reproduces the review's scenario deterministically (MAX_SESSIONS=1, hung LRU victim, gated initialize, second caller receives the session, creator cancelled): red on the previous commit, green now. test_joiner_follows_creation_outcome_when_creator_is_cancelled pins the joiner outcome-gating semantics as a drift guard.
This commit is contained in:
hataa
GitHub
parent
0dd233afc4
commit
567a06783c
@ -29,6 +29,13 @@ To make this impossible, every pooled session is owned by a dedicated
|
||||
session back to the caller, and then *waits* on a close event. All shutdown
|
||||
paths only ever **signal** that event; the owner task performs ``__aexit__``
|
||||
itself, guaranteeing enter and exit always happen in the same task.
|
||||
|
||||
The owner task is also the only writer that promotes a creation into the pool:
|
||||
once ``initialize()`` succeeds it registers the session in ``_entries`` and
|
||||
resolves the creation's future in one atomic critical section, so callers can
|
||||
only ever receive a session the pool already owns (and will retire via LRU
|
||||
eviction or the close_* paths) — never one whose lifetime is still tied to a
|
||||
single caller that might get cancelled.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@ -132,6 +139,10 @@ class MCPSessionPool:
|
||||
# In-flight creations, keyed by (server, scope). Lets concurrent callers
|
||||
# on the same loop share a single creation instead of each spawning a
|
||||
# duplicate session. Value: (loop, ready_future, owner_task, close_event).
|
||||
# The owner task promotes the record into ``_entries`` and resolves
|
||||
# ``ready`` with the session in one atomic critical section (see
|
||||
# ``_run_session``), so ``ready`` resolving with a *result* always means
|
||||
# the session is registered — never merely handed to one caller.
|
||||
self._inflight: dict[
|
||||
tuple[str, str],
|
||||
tuple[
|
||||
@ -144,6 +155,12 @@ class MCPSessionPool:
|
||||
# threading.Lock is not bound to any event loop, so it is safe to
|
||||
# acquire from both async paths and sync/worker-thread paths.
|
||||
self._lock = threading.Lock()
|
||||
# Strong references to detached teardown reaper tasks. The event loop
|
||||
# only keeps weak references to tasks, so an unheld reaper — and with
|
||||
# it the owner it is awaiting, mid-__aexit__ — could be
|
||||
# garbage-collected before teardown completes; the done callback keeps
|
||||
# the set from growing without bound.
|
||||
self._teardown_tasks: set[asyncio.Task[Any]] = set()
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Session owner task
|
||||
@ -151,16 +168,19 @@ class MCPSessionPool:
|
||||
|
||||
async def _run_session(
|
||||
self,
|
||||
key: tuple[str, str],
|
||||
connection: dict[str, Any],
|
||||
ready: asyncio.Future[ClientSession],
|
||||
close_evt: asyncio.Event,
|
||||
) -> None:
|
||||
"""Own a single MCP session for its entire lifetime.
|
||||
|
||||
Enters the session context manager, initializes it, publishes the live
|
||||
session via ``ready``, then blocks until ``close_evt`` is set. The
|
||||
context manager is *always* exited from this task, satisfying anyio's
|
||||
cancel-scope same-task requirement.
|
||||
Enters the session context manager, initializes it, and *commits* the
|
||||
session: promotes the in-flight record into ``_entries`` and resolves
|
||||
``ready`` with the live session in one atomic critical section, then
|
||||
blocks until ``close_evt`` is set. The context manager is *always*
|
||||
exited from this task, satisfying anyio's cancel-scope same-task
|
||||
requirement.
|
||||
"""
|
||||
from langchain_mcp_adapters.sessions import create_session
|
||||
|
||||
@ -179,8 +199,29 @@ class MCPSessionPool:
|
||||
# leaking the session/subprocess.
|
||||
try:
|
||||
await session.initialize()
|
||||
if not ready.done():
|
||||
ready.set_result(session)
|
||||
# Commit point. ``ready`` resolves with a *result* only inside the
|
||||
# critical section that also registers the session in ``_entries``:
|
||||
# a resolved-with-result future therefore means the session is
|
||||
# pool-owned (visible to LRU eviction and the close_* paths), never
|
||||
# the private property of one caller. Joiners waiting on ``ready``
|
||||
# can only ever receive a committed session, and an unwind path can
|
||||
# check the outcome race-free. If the record was already removed
|
||||
# (the creator unwound or a close_* ran), the creation is aborted:
|
||||
# ``ready`` carries the cancellation instead, so no caller can end
|
||||
# up holding an unmanaged session.
|
||||
loop = asyncio.get_running_loop()
|
||||
task = asyncio.current_task()
|
||||
with self._lock:
|
||||
still_ours = self._inflight.get(key) == (loop, ready, task, close_evt)
|
||||
if still_ours:
|
||||
self._inflight.pop(key)
|
||||
self._entries[key] = (session, loop, task, close_evt)
|
||||
if not ready.done():
|
||||
ready.set_result(session)
|
||||
if still_ours:
|
||||
logger.info("Created persistent MCP session for %s/%s", key[0], key[1])
|
||||
elif not ready.done():
|
||||
ready.set_exception(asyncio.CancelledError("MCP session pool was closed while the session was being created"))
|
||||
await close_evt.wait()
|
||||
except BaseException as e:
|
||||
if not ready.done():
|
||||
@ -217,10 +258,13 @@ class MCPSessionPool:
|
||||
# Phase 1: inspect/mutate the registry under the thread lock (no awaits).
|
||||
# Decide one of three outcomes atomically: return an existing session,
|
||||
# join an in-flight creation, or become the creator for this key.
|
||||
# Each item: (loop, owner_task, close_event, cancel). ``cancel`` is True
|
||||
# for in-flight creations, whose owner may be blocked inside
|
||||
# ``initialize()`` where close_evt cannot wake it — it must be cancelled.
|
||||
evicted: list[tuple[asyncio.AbstractEventLoop, asyncio.Task[Any], asyncio.Event, bool]] = []
|
||||
# Each item: (loop, owner_task, close_event, cancel, ready_future).
|
||||
# ``cancel`` is True for in-flight creations, whose owner may be blocked
|
||||
# inside ``initialize()`` where close_evt cannot wake it — it must be
|
||||
# cancelled (guarded: never when the owner already failed and is
|
||||
# unwinding in __aexit__). ``ready`` is the creation's future, used only
|
||||
# for that guard.
|
||||
evicted: list[tuple[asyncio.AbstractEventLoop, asyncio.Task[Any], asyncio.Event, bool, asyncio.Future[ClientSession] | None]] = []
|
||||
join: asyncio.Future[ClientSession] | None = None
|
||||
ready: asyncio.Future[ClientSession] | None = None
|
||||
close_evt: asyncio.Event | None = None
|
||||
@ -233,7 +277,7 @@ class MCPSessionPool:
|
||||
return session
|
||||
# Session belongs to a different/closed event loop – evict it.
|
||||
self._entries.pop(key)
|
||||
evicted.append((loop, ent_task, ent_close, False))
|
||||
evicted.append((loop, ent_task, ent_close, False, None))
|
||||
|
||||
inflight = self._inflight.get(key)
|
||||
if inflight is not None and inflight[0] is current_loop and not inflight[0].is_closed():
|
||||
@ -248,86 +292,118 @@ class MCPSessionPool:
|
||||
# wake it), it must be cancelled. We then create a fresh
|
||||
# session here.
|
||||
self._inflight.pop(key)
|
||||
evicted.append((inflight[0], inflight[2], inflight[3], True))
|
||||
evicted.append((inflight[0], inflight[2], inflight[3], True, inflight[1]))
|
||||
# Become the creator: publish an in-flight record before any
|
||||
# await so concurrent callers join us instead of racing.
|
||||
ready = current_loop.create_future()
|
||||
close_evt = asyncio.Event()
|
||||
task = current_loop.create_task(self._run_session(connection, ready, close_evt))
|
||||
task = current_loop.create_task(self._run_session(key, connection, ready, close_evt))
|
||||
self._inflight[key] = (current_loop, ready, task, close_evt)
|
||||
|
||||
# Evict LRU entries when at capacity.
|
||||
while len(self._entries) >= self.MAX_SESSIONS:
|
||||
oldest_key, (_, loop, ent_task, ent_close) = next(iter(self._entries.items()))
|
||||
self._entries.pop(oldest_key)
|
||||
evicted.append((loop, ent_task, ent_close, False))
|
||||
evicted.append((loop, ent_task, ent_close, False, None))
|
||||
|
||||
# Phase 2: shut down evicted sessions/creations. Same-loop owners are
|
||||
# awaited so they finish deterministically; foreign-loop owners are
|
||||
# routed to their own loop. In every case the owner task — never this
|
||||
# one — runs __aexit__. In-flight owners are cancelled (cancel=True) so a
|
||||
# blocking initialize() cannot leave them hung.
|
||||
for loop, ent_task, ent_close, cancel in evicted:
|
||||
if loop is current_loop and not loop.is_closed():
|
||||
await self._shutdown(ent_close, ent_task, cancel)
|
||||
elif cancel:
|
||||
await self._shutdown_entry(loop, ent_task, ent_close, cancel=True)
|
||||
else:
|
||||
self._signal_close(loop, ent_close)
|
||||
# Phase 2: shut down evicted sessions/creations. Signal EVERY removed
|
||||
# owner first — the signal loop contains no awaits, so it completes
|
||||
# atomically and a cancellation during the teardown awaits below can
|
||||
# never strand an owner that was already removed from the registries.
|
||||
# Then await teardowns (same-loop deterministically; foreign-loop
|
||||
# in-flight creations routed to their loop). In every case the owner
|
||||
# task — never this one — runs __aexit__.
|
||||
for loop, ent_task, ent_close, cancel, ent_ready in evicted:
|
||||
self._signal_close(loop, ent_close)
|
||||
if cancel:
|
||||
self._cancel_owner(loop, ent_task, ent_ready)
|
||||
try:
|
||||
for loop, ent_task, ent_close, cancel, ent_ready in evicted:
|
||||
if loop is current_loop and not loop.is_closed():
|
||||
await self._shutdown(ent_close, ent_task, cancel=False, ready=ent_ready)
|
||||
elif cancel:
|
||||
await self._shutdown_entry(loop, ent_task, ent_close, cancel=False, ready=ent_ready)
|
||||
# else: foreign-loop registered entry — already signalled above;
|
||||
# its teardown completes on its own loop.
|
||||
except BaseException:
|
||||
# We may already be the creator for ``key``: the in-flight record
|
||||
# and owner task were published under the lock *before* these
|
||||
# awaits. A cancellation here (routine — both production call
|
||||
# sites wrap get_session in asyncio.wait_for(session_init_timeout))
|
||||
# would otherwise orphan that owner: it finishes initialize(),
|
||||
# commits, and blocks on close_evt forever — invisible to LRU
|
||||
# eviction, which only scans _entries. Unwind exactly like the
|
||||
# Phase-3 path — unless the owner already committed while we were
|
||||
# parked here: then the session is registered in _entries and may
|
||||
# already be in a joiner's hands, so tearing it down would close it
|
||||
# underneath them. A committed session is pool property; LRU
|
||||
# eviction and the close_* paths own it from here on. Every evicted
|
||||
# owner was already signalled in the atomic loop above, so skipping
|
||||
# the remaining teardown *awaits* (not signals) on unwind is safe:
|
||||
# each victim still tears down in its own task.
|
||||
if task is not None and not self._creation_committed(ready):
|
||||
assert ready is not None and close_evt is not None
|
||||
try:
|
||||
await self._shutdown(close_evt, task, cancel=True, ready=ready)
|
||||
except BaseException:
|
||||
logger.debug("Owner teardown interrupted during eviction unwind", exc_info=True)
|
||||
with self._lock:
|
||||
if self._inflight.get(key) == (current_loop, ready, task, close_evt):
|
||||
self._inflight.pop(key)
|
||||
raise
|
||||
|
||||
# Phase 2b: a concurrent creation for this key is already in progress on
|
||||
# this loop — share its result rather than create a duplicate session.
|
||||
# The creator's owner task resolves ``ready`` with a result only at the
|
||||
# commit critical section that also registers the session, so a value
|
||||
# returned here is always pool-owned; if the creation is aborted (the
|
||||
# creator unwound or a close_* ran), ``ready`` carries the exception and
|
||||
# this joiner fails with it instead of holding an unmanaged session.
|
||||
if join is not None:
|
||||
return await asyncio.shield(join)
|
||||
|
||||
assert ready is not None and close_evt is not None and task is not None
|
||||
|
||||
# Phase 3: wait for our owner task to publish the initialized session.
|
||||
# Phase 3: wait for our owner task to commit the initialized session.
|
||||
# A successful result means the session is already registered in
|
||||
# _entries — the commit critical section did both — so from that moment
|
||||
# the pool, not this call, owns its lifetime.
|
||||
try:
|
||||
session = await asyncio.shield(ready)
|
||||
except BaseException:
|
||||
# Two distinct cases reach here:
|
||||
# Three distinct cases reach here:
|
||||
#
|
||||
# 1. The owner task failed (e.g. connect/initialize error) and
|
||||
# reported it via ready.set_exception(). It is *already* in its
|
||||
# finally block running cm.__aexit__ in its own task, so we must
|
||||
# NOT cancel it — doing so would interrupt that cleanup. We only
|
||||
# wait for it to finish unwinding.
|
||||
# 2. This call itself was cancelled (CancelledError). Because of the
|
||||
# shield, `ready` is still pending and the owner task is alive and
|
||||
# blocked. We signal close and cancel it so it exits the cancel
|
||||
# scope in its own task, then wait for it to finish.
|
||||
#
|
||||
# The session is never registered yet, so nobody else can close it;
|
||||
# waiting here guarantees we never leak a session or owner task.
|
||||
owner_already_failed = ready.done() and not ready.cancelled() and ready.exception() is not None
|
||||
if not owner_already_failed:
|
||||
close_evt.set()
|
||||
task.cancel()
|
||||
try:
|
||||
await asyncio.shield(task)
|
||||
except BaseException:
|
||||
logger.debug("Owner task ended during get_session unwind", exc_info=True)
|
||||
with self._lock:
|
||||
if self._inflight.get(key) == (current_loop, ready, task, close_evt):
|
||||
self._inflight.pop(key)
|
||||
# 2. This call itself was cancelled (CancelledError) while the
|
||||
# creation was still in flight. Because of the shield, `ready`
|
||||
# is still pending and the owner task is alive. We signal close
|
||||
# and cancel it so it exits the cancel scope in its own task,
|
||||
# then wait for it to finish.
|
||||
# 3. This call was cancelled at the very moment the owner
|
||||
# committed: `ready` resolved with a result and the session is
|
||||
# registered in _entries. Joiners may already hold it, so we
|
||||
# must NOT tear it down — just propagate the cancellation and
|
||||
# leave the pooled session to LRU eviction / close_*.
|
||||
if not self._creation_committed(ready):
|
||||
owner_already_failed = ready.done() and not ready.cancelled() and ready.exception() is not None
|
||||
if not owner_already_failed:
|
||||
close_evt.set()
|
||||
task.cancel()
|
||||
try:
|
||||
await asyncio.shield(task)
|
||||
except BaseException:
|
||||
logger.debug("Owner task ended during get_session unwind", exc_info=True)
|
||||
with self._lock:
|
||||
if self._inflight.get(key) == (current_loop, ready, task, close_evt):
|
||||
self._inflight.pop(key)
|
||||
raise
|
||||
|
||||
# Phase 4: promote the in-flight creation to a registered entry — but
|
||||
# only if our in-flight record is still the live one. A concurrent
|
||||
# close_* / close_all may have removed it while we were initializing; in
|
||||
# that case we must NOT resurrect the session into _entries. Instead we
|
||||
# own the teardown: signal our owner task and wait for it to run
|
||||
# __aexit__ in its own task, then surface the cancellation.
|
||||
with self._lock:
|
||||
still_ours = self._inflight.get(key) == (current_loop, ready, task, close_evt)
|
||||
if still_ours:
|
||||
self._inflight.pop(key)
|
||||
self._entries[key] = (session, current_loop, task, close_evt)
|
||||
if not still_ours:
|
||||
await self._shutdown(close_evt, task)
|
||||
raise asyncio.CancelledError("MCP session pool was closed while the session was being created")
|
||||
logger.info("Created persistent MCP session for %s/%s", server_name, scope_key)
|
||||
# Phase 4: the commit inside the owner task already promoted the
|
||||
# creation into a registered entry; nothing left to decide here.
|
||||
return session
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
@ -349,25 +425,123 @@ class MCPSessionPool:
|
||||
# Loop was closed between the is_closed() check and now.
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
def _owner_unwinding_after_failure(ready: asyncio.Future[ClientSession] | None) -> bool:
|
||||
"""True when the owner already failed and is running ``__aexit__``.
|
||||
|
||||
Cancelling such an owner would interrupt that in-task cleanup — the
|
||||
same-task exit anyio requires — so callers must skip the cancel.
|
||||
"""
|
||||
return ready is not None and ready.done() and not ready.cancelled() and ready.exception() is not None
|
||||
|
||||
@staticmethod
|
||||
def _creation_committed(ready: asyncio.Future[ClientSession] | None) -> bool:
|
||||
"""True once the creation's session is registered in ``_entries``.
|
||||
|
||||
``_run_session`` resolves ``ready`` with a *result* only inside the
|
||||
commit critical section that also moves the record into ``_entries``,
|
||||
so this is exactly the state in which the session is pool-owned —
|
||||
visible to LRU eviction and the close_* paths, and possibly already
|
||||
handed to a joiner. Unwind paths must not tear such a session down:
|
||||
closing it here would yank a live session from a joiner's hands
|
||||
(#5008 review).
|
||||
"""
|
||||
return ready is not None and ready.done() and not ready.cancelled() and ready.exception() is None
|
||||
|
||||
@classmethod
|
||||
def _cancel_owner(cls, loop: asyncio.AbstractEventLoop, task: asyncio.Task[Any], ready: asyncio.Future[ClientSession] | None = None) -> None:
|
||||
"""Thread-safe guarded cancel of an owner task on its owning loop.
|
||||
|
||||
The failure-state recheck runs INSIDE the callback that executes on the
|
||||
owning loop, immediately before ``task.cancel()``: evaluating it on the
|
||||
caller's loop first and queueing the cancel after opens a
|
||||
time-of-check/time-of-use window in which the owner can fail, publish
|
||||
the exception to ``ready``, and enter ``__aexit__`` — the already-queued
|
||||
cancel would then interrupt that in-task cleanup. On the owning loop
|
||||
the check and the cancel are atomic with respect to owner-task
|
||||
progress (the owner cannot advance between two non-awaiting
|
||||
statements), so an owner already unwinding after failure is never
|
||||
cancelled.
|
||||
"""
|
||||
|
||||
def _guarded_cancel() -> None:
|
||||
if cls._owner_unwinding_after_failure(ready):
|
||||
return
|
||||
task.cancel()
|
||||
|
||||
if loop.is_closed():
|
||||
return
|
||||
try:
|
||||
loop.call_soon_threadsafe(_guarded_cancel)
|
||||
except RuntimeError:
|
||||
# Loop was closed between the is_closed() check and now.
|
||||
pass
|
||||
|
||||
def _track_owner_teardown(self, task: asyncio.Task[Any]) -> None:
|
||||
"""Keep an owner's teardown observable after its awaiter was cancelled.
|
||||
|
||||
The awaiter unwinds, but the owner still has to finish ``__aexit__`` in
|
||||
its own task; the reaper awaits that completion so exceptions are
|
||||
retrieved and the teardown stays observable instead of dangling. The
|
||||
reaper is strongly retained in ``_teardown_tasks`` until it finishes —
|
||||
the loop only keeps weak task references, so an unheld reaper (and
|
||||
transitively the owner mid-``__aexit__``) could be garbage-collected
|
||||
before the exit completes (#5008 review).
|
||||
"""
|
||||
|
||||
async def _reap() -> None:
|
||||
try:
|
||||
await task
|
||||
except BaseException:
|
||||
logger.debug("Owner task ended after caller cancellation", exc_info=True)
|
||||
|
||||
try:
|
||||
reaper = asyncio.get_running_loop().create_task(_reap(), name=f"mcp-session-owner-reap:{task.get_name()}")
|
||||
except RuntimeError:
|
||||
return
|
||||
self._teardown_tasks.add(reaper)
|
||||
reaper.add_done_callback(self._teardown_tasks.discard)
|
||||
|
||||
async def _shutdown(
|
||||
self,
|
||||
close_evt: asyncio.Event,
|
||||
task: asyncio.Task[Any],
|
||||
cancel: bool = False,
|
||||
ready: asyncio.Future[ClientSession] | None = None,
|
||||
) -> None:
|
||||
"""Signal an owner task and wait for it to finish (runs on its loop).
|
||||
|
||||
``cancel=True`` is used for in-flight creations: the owner task may be
|
||||
blocked inside ``initialize()`` where ``close_evt`` cannot wake it, so it
|
||||
must be cancelled. Its ``finally`` block still runs ``__aexit__`` in its
|
||||
own task, satisfying anyio's same-task cancel-scope requirement.
|
||||
must be cancelled — unless it already failed and is unwinding in its
|
||||
``finally`` block (``ready`` carries an exception), where a cancel would
|
||||
interrupt the in-task ``__aexit__``. The exit always runs in the owner
|
||||
task itself, satisfying anyio's same-task cancel-scope requirement.
|
||||
|
||||
The await is shielded: a cancellation of *this* awaiting task
|
||||
propagates immediately without cancelling the owner, whose teardown
|
||||
keeps running under a tracked reaper task. The victim's own
|
||||
cancellation or exception is swallowed and logged as before.
|
||||
"""
|
||||
close_evt.set()
|
||||
if cancel:
|
||||
if cancel and not self._owner_unwinding_after_failure(ready):
|
||||
task.cancel()
|
||||
caller = asyncio.current_task()
|
||||
caller_cancels = caller.cancelling() if caller is not None else 0
|
||||
try:
|
||||
await task
|
||||
except (Exception, asyncio.CancelledError):
|
||||
await asyncio.shield(task)
|
||||
except asyncio.CancelledError:
|
||||
# ``shield`` surfaces the victim's cancellation (we cancelled it
|
||||
# above, or someone else did). But if the count rose, the
|
||||
# cancellation belongs to US — the awaiting task was cancelled
|
||||
# while waiting (e.g. a get_session parked in eviction teardown
|
||||
# under asyncio.wait_for). Propagate it while keeping the owner's
|
||||
# teardown tracked so __aexit__ still completes.
|
||||
if caller is not None and caller.cancelling() > caller_cancels:
|
||||
self._track_owner_teardown(task)
|
||||
raise
|
||||
logger.debug("Owner task cancelled during shutdown")
|
||||
except Exception:
|
||||
logger.debug("Owner task ended during shutdown", exc_info=True)
|
||||
|
||||
async def _shutdown_entry(
|
||||
@ -376,15 +550,16 @@ class MCPSessionPool:
|
||||
task: asyncio.Task[Any],
|
||||
close_evt: asyncio.Event,
|
||||
cancel: bool = False,
|
||||
ready: asyncio.Future[ClientSession] | None = None,
|
||||
) -> None:
|
||||
"""Shut down one entry, routing the close to its owning loop."""
|
||||
if loop.is_closed():
|
||||
return
|
||||
current_loop = asyncio.get_running_loop()
|
||||
if loop is current_loop:
|
||||
await self._shutdown(close_evt, task, cancel)
|
||||
await self._shutdown(close_evt, task, cancel, ready=ready)
|
||||
elif loop.is_running():
|
||||
future = asyncio.run_coroutine_threadsafe(self._shutdown(close_evt, task, cancel), loop)
|
||||
future = asyncio.run_coroutine_threadsafe(self._shutdown(close_evt, task, cancel, ready=ready), loop)
|
||||
try:
|
||||
await asyncio.wrap_future(future)
|
||||
except Exception:
|
||||
@ -402,10 +577,31 @@ class MCPSessionPool:
|
||||
logger.warning("Owning loop for MCP session is idle; signalling close best-effort. Session may leak until the loop runs again.")
|
||||
self._signal_close(loop, close_evt)
|
||||
if cancel:
|
||||
try:
|
||||
loop.call_soon_threadsafe(task.cancel)
|
||||
except RuntimeError:
|
||||
pass
|
||||
self._cancel_owner(loop, task, ready)
|
||||
|
||||
async def _close_owners(
|
||||
self,
|
||||
entries: list[tuple[ClientSession, asyncio.AbstractEventLoop, asyncio.Task[Any], asyncio.Event]],
|
||||
inflight: list[tuple[asyncio.AbstractEventLoop, asyncio.Future[ClientSession], asyncio.Task[Any], asyncio.Event]],
|
||||
) -> None:
|
||||
"""Shut down already-removed owners: signal all first, then await.
|
||||
|
||||
Signalling every removed owner BEFORE awaiting any teardown guarantees
|
||||
a cancellation of the close call can never strand an owner that is no
|
||||
longer reachable through the registries: each has its close event (and,
|
||||
for in-flight creations, its guarded cancel) in hand and tears down in
|
||||
its own task regardless. The awaiting phase adds best-effort
|
||||
determinism on top; skipping the remaining awaits on unwind is safe.
|
||||
"""
|
||||
for _session, loop, ent_task, ent_close in entries:
|
||||
self._signal_close(loop, ent_close)
|
||||
for loop, ent_ready, ent_task, ent_close in inflight:
|
||||
self._signal_close(loop, ent_close)
|
||||
self._cancel_owner(loop, ent_task, ent_ready)
|
||||
for _session, loop, ent_task, ent_close in entries:
|
||||
await self._shutdown_entry(loop, ent_task, ent_close)
|
||||
for loop, ent_ready, ent_task, ent_close in inflight:
|
||||
await self._shutdown_entry(loop, ent_task, ent_close, ready=ent_ready)
|
||||
|
||||
async def close_scope(self, scope_key: str) -> None:
|
||||
"""Close all sessions for a given scope (e.g. thread_id)."""
|
||||
@ -414,10 +610,7 @@ class MCPSessionPool:
|
||||
entries = [(self._entries.pop(k)) for k in keys]
|
||||
inflight_keys = [k for k in self._inflight if k[1] == scope_key]
|
||||
inflight = [self._inflight.pop(k) for k in inflight_keys]
|
||||
for _session, loop, task, close_evt in entries:
|
||||
await self._shutdown_entry(loop, task, close_evt)
|
||||
for loop, _ready, task, close_evt in inflight:
|
||||
await self._shutdown_entry(loop, task, close_evt, cancel=True)
|
||||
await self._close_owners(entries, inflight)
|
||||
|
||||
async def close_session(self, server_name: str, scope_key: str) -> None:
|
||||
"""Close one exact server/scope session so a retry reconnects cleanly."""
|
||||
@ -425,12 +618,7 @@ class MCPSessionPool:
|
||||
with self._lock:
|
||||
entry = self._entries.pop(key, None)
|
||||
inflight = self._inflight.pop(key, None)
|
||||
if entry is not None:
|
||||
_session, loop, task, close_evt = entry
|
||||
await self._shutdown_entry(loop, task, close_evt)
|
||||
if inflight is not None:
|
||||
loop, _ready, task, close_evt = inflight
|
||||
await self._shutdown_entry(loop, task, close_evt, cancel=True)
|
||||
await self._close_owners([entry] if entry is not None else [], [inflight] if inflight is not None else [])
|
||||
|
||||
async def close_session_if_current(
|
||||
self,
|
||||
@ -456,10 +644,7 @@ class MCPSessionPool:
|
||||
entries = [(self._entries.pop(k)) for k in keys]
|
||||
inflight_keys = [k for k in self._inflight if k[0] == server_name]
|
||||
inflight = [self._inflight.pop(k) for k in inflight_keys]
|
||||
for _session, loop, task, close_evt in entries:
|
||||
await self._shutdown_entry(loop, task, close_evt)
|
||||
for loop, _ready, task, close_evt in inflight:
|
||||
await self._shutdown_entry(loop, task, close_evt, cancel=True)
|
||||
await self._close_owners(entries, inflight)
|
||||
|
||||
async def close_all(self) -> None:
|
||||
"""Close every managed session."""
|
||||
@ -468,10 +653,7 @@ class MCPSessionPool:
|
||||
self._entries.clear()
|
||||
inflight = list(self._inflight.values())
|
||||
self._inflight.clear()
|
||||
for _session, loop, task, close_evt in entries:
|
||||
await self._shutdown_entry(loop, task, close_evt)
|
||||
for loop, _ready, task, close_evt in inflight:
|
||||
await self._shutdown_entry(loop, task, close_evt, cancel=True)
|
||||
await self._close_owners(entries, inflight)
|
||||
|
||||
def close_all_sync(self) -> None:
|
||||
"""Close all sessions on their owning event loops (synchronous).
|
||||
@ -499,14 +681,17 @@ class MCPSessionPool:
|
||||
self._inflight.clear()
|
||||
|
||||
# Entries are initialized (gentle close_evt path). In-flight creations
|
||||
# may be blocked mid-init, so they are cancelled to unblock teardown.
|
||||
owners = [(loop, task, close_evt, False) for _s, loop, task, close_evt in entries]
|
||||
owners += [(loop, task, close_evt, True) for loop, _r, task, close_evt in inflight]
|
||||
# may be blocked mid-init, so they are cancelled to unblock teardown —
|
||||
# but every guard below re-checks on the owning loop (or atomically on
|
||||
# this thread for the current-loop branch), so an owner that has since
|
||||
# failed and is unwinding in __aexit__ is never cancelled.
|
||||
owners = [(loop, task, close_evt, False, None) for _s, loop, task, close_evt in entries]
|
||||
owners += [(loop, task, ent_close, True, ent_ready) for loop, ent_ready, task, ent_close in inflight]
|
||||
try:
|
||||
current_running_loop = asyncio.get_running_loop()
|
||||
except RuntimeError:
|
||||
current_running_loop = None
|
||||
for loop, task, close_evt, cancel in owners:
|
||||
for loop, task, close_evt, cancel, ent_ready in owners:
|
||||
if loop.is_closed():
|
||||
continue
|
||||
try:
|
||||
@ -515,16 +700,18 @@ class MCPSessionPool:
|
||||
# waiting on run_coroutine_threadsafe(...).result() would
|
||||
# deadlock until timeout. Signal the owner task directly and
|
||||
# let it finish once this synchronous call returns control to
|
||||
# the running loop.
|
||||
# the running loop. Same-thread, no yield between the guard
|
||||
# and the cancel, so the check is atomic here.
|
||||
close_evt.set()
|
||||
if cancel:
|
||||
if cancel and not self._owner_unwinding_after_failure(ent_ready):
|
||||
task.cancel()
|
||||
elif loop.is_running():
|
||||
# Schedule the shutdown on the owning loop from this thread.
|
||||
future = asyncio.run_coroutine_threadsafe(self._shutdown(close_evt, task, cancel), loop)
|
||||
# Schedule the shutdown on the owning loop from this thread;
|
||||
# _shutdown applies the failure guard on that loop.
|
||||
future = asyncio.run_coroutine_threadsafe(self._shutdown(close_evt, task, cancel, ready=ent_ready), loop)
|
||||
future.result(timeout=self.SESSION_CLOSE_TIMEOUT)
|
||||
else:
|
||||
loop.run_until_complete(self._shutdown(close_evt, task, cancel))
|
||||
loop.run_until_complete(self._shutdown(close_evt, task, cancel, ready=ent_ready))
|
||||
except Exception:
|
||||
logger.debug("Error closing MCP session during sync close", exc_info=True)
|
||||
|
||||
|
||||
@ -1,10 +1,12 @@
|
||||
"""Tests for the MCP persistent-session pool."""
|
||||
|
||||
import asyncio
|
||||
import gc
|
||||
import logging
|
||||
import stat
|
||||
import sys
|
||||
import threading
|
||||
import weakref
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import anyio
|
||||
@ -1721,6 +1723,549 @@ async def test_get_session_cancelled_while_initializing_does_not_leak():
|
||||
assert not leaked, "owner task must not be left pending after cancellation"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_session_cancelled_during_eviction_teardown_does_not_leak():
|
||||
"""Cancelling get_session while it awaits an evicted session's teardown
|
||||
must not orphan the just-created owner task.
|
||||
|
||||
The in-flight record and owner task are published before the Phase-2
|
||||
eviction awaits, so a caller cancelled there — routine, since both
|
||||
production call sites wrap get_session in asyncio.wait_for(
|
||||
session_init_timeout) — would otherwise leak the owner: it finishes
|
||||
initialize(), publishes ready, and blocks on close_evt forever, invisible
|
||||
to LRU eviction, which only scans _entries.
|
||||
"""
|
||||
pool = MCPSessionPool()
|
||||
pool.MAX_SESSIONS = 1
|
||||
|
||||
# A pre-registered LRU victim whose teardown hangs (wedged server): the
|
||||
# creator parks in Phase 2 awaiting it, so the cancel below is guaranteed
|
||||
# to land between publishing the in-flight record and Phase 3.
|
||||
victim_hang = asyncio.Event()
|
||||
|
||||
async def victim_owner() -> None:
|
||||
await victim_hang.wait()
|
||||
|
||||
loop = asyncio.get_running_loop()
|
||||
victim_task = asyncio.create_task(victim_owner())
|
||||
pool._entries[("victim", "scope")] = (MagicMock(), loop, victim_task, asyncio.Event())
|
||||
|
||||
gate = asyncio.Event()
|
||||
cms: list[_BlockingInitCm] = []
|
||||
call: asyncio.Task | None = None
|
||||
|
||||
def make_cm(*args, **kwargs):
|
||||
cm = _BlockingInitCm(gate)
|
||||
cms.append(cm)
|
||||
return cm
|
||||
|
||||
try:
|
||||
with patch("langchain_mcp_adapters.sessions.create_session", side_effect=make_cm):
|
||||
call = asyncio.create_task(pool.get_session("srv", "scope-2", {"transport": "stdio", "command": "x", "args": []}))
|
||||
# The creator publishes its in-flight record before awaiting the
|
||||
# hung victim, so this is deterministic.
|
||||
for _ in range(100):
|
||||
if ("srv", "scope-2") in pool._inflight:
|
||||
break
|
||||
await asyncio.sleep(0.01)
|
||||
assert ("srv", "scope-2") in pool._inflight
|
||||
await asyncio.sleep(0.02)
|
||||
|
||||
call.cancel()
|
||||
# On the bug, the caller's cancellation is swallowed by the
|
||||
# eviction teardown and the task never finishes; the shield keeps
|
||||
# the timeout from cancelling it so the hang is observable.
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await asyncio.wait_for(asyncio.shield(call), timeout=5.0)
|
||||
|
||||
# The orphaned owner must have been torn down: its __aexit__ ran even
|
||||
# though initialize() was never released.
|
||||
for _ in range(50):
|
||||
if cms and cms[0].closed:
|
||||
break
|
||||
await asyncio.sleep(0.01)
|
||||
assert cms, "owner task must have been created"
|
||||
assert cms[0].closed, "owner must run __aexit__ after Phase-2 cancellation"
|
||||
assert ("srv", "scope-2") not in pool._inflight
|
||||
|
||||
current = asyncio.current_task()
|
||||
leaked = [t for t in asyncio.all_tasks() if t is not current and not t.done() and "_run_session" in str(t.get_coro())]
|
||||
assert not leaked, "owner task must not be left pending after Phase-2 cancellation"
|
||||
finally:
|
||||
if call is not None and not call.done():
|
||||
call.cancel()
|
||||
victim_hang.set()
|
||||
victim_task.cancel()
|
||||
try:
|
||||
await victim_task
|
||||
except BaseException:
|
||||
pass
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancelled_creator_does_not_close_session_held_by_joiner():
|
||||
"""Cancelling the creator must not close a session a joiner already holds.
|
||||
|
||||
Reproduces the review race deterministically: MAX_SESSIONS=1, the creator
|
||||
parks in Phase 2 awaiting a hung LRU victim while its owner finishes
|
||||
initialize(); a second caller for the same key then receives the session.
|
||||
Cancelling the creator afterwards must leave that session open and
|
||||
registered — once the creation committed, the session is pool property
|
||||
(LRU eviction / close_* own it), never the cancelled caller's (#5008
|
||||
review).
|
||||
"""
|
||||
pool = MCPSessionPool()
|
||||
pool.MAX_SESSIONS = 1
|
||||
|
||||
# A pre-registered LRU victim whose teardown hangs: the creator parks in
|
||||
# Phase 2 awaiting it, so the cancel below lands while the creator is
|
||||
# between publishing the in-flight record and awaiting ready.
|
||||
victim_hang = asyncio.Event()
|
||||
|
||||
async def victim_owner() -> None:
|
||||
await victim_hang.wait()
|
||||
|
||||
loop = asyncio.get_running_loop()
|
||||
victim_task = asyncio.create_task(victim_owner())
|
||||
pool._entries[("victim", "scope")] = (MagicMock(), loop, victim_task, asyncio.Event())
|
||||
|
||||
gate = asyncio.Event()
|
||||
cms: list[_BlockingInitCm] = []
|
||||
creator: asyncio.Task | None = None
|
||||
|
||||
def make_cm(*args, **kwargs):
|
||||
cm = _BlockingInitCm(gate)
|
||||
cms.append(cm)
|
||||
return cm
|
||||
|
||||
try:
|
||||
with patch("langchain_mcp_adapters.sessions.create_session", side_effect=make_cm):
|
||||
conn = {"transport": "stdio", "command": "x", "args": []}
|
||||
creator = asyncio.create_task(pool.get_session("srv", "scope-2", conn))
|
||||
# The creator publishes its in-flight record before awaiting the
|
||||
# hung victim, so this is deterministic.
|
||||
for _ in range(100):
|
||||
if ("srv", "scope-2") in pool._inflight:
|
||||
break
|
||||
await asyncio.sleep(0.01)
|
||||
assert ("srv", "scope-2") in pool._inflight
|
||||
|
||||
# The owner finishes initialize() and commits while its creator is
|
||||
# still parked on the victim's teardown. The second caller then
|
||||
# receives the initialized session (Phase 2b before this fix, the
|
||||
# _entries fast path after) — either way it now holds it.
|
||||
gate.set()
|
||||
joiner = asyncio.create_task(pool.get_session("srv", "scope-2", conn))
|
||||
session = await asyncio.wait_for(asyncio.shield(joiner), timeout=5.0)
|
||||
assert cms, "owner task must have been created"
|
||||
|
||||
# On the bug, this unwind unconditionally shut the owner down even
|
||||
# though a joiner held the session: __aexit__ ran underneath it.
|
||||
creator.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await asyncio.wait_for(asyncio.shield(creator), timeout=5.0)
|
||||
|
||||
await asyncio.sleep(0.02)
|
||||
assert not cms[0].closed, "cancelling the creator must not close a session already handed to a caller"
|
||||
assert ("srv", "scope-2") in pool._entries, "committed session must stay registered after creator cancellation"
|
||||
assert pool._entries[("srv", "scope-2")][0] is session, "the joiner's session must be the registered one"
|
||||
assert ("srv", "scope-2") not in pool._inflight
|
||||
|
||||
# Cleanup: close the pool so the parked owner finishes deterministically.
|
||||
await pool.close_all()
|
||||
for _ in range(100):
|
||||
if cms[0].closed:
|
||||
break
|
||||
await asyncio.sleep(0.01)
|
||||
assert cms[0].closed, "owner must still tear down through the pool close paths"
|
||||
|
||||
current = asyncio.current_task()
|
||||
leaked = [t for t in asyncio.all_tasks() if t is not current and not t.done() and "_run_session" in str(t.get_coro())]
|
||||
assert not leaked, "owner task must not be left pending after cleanup"
|
||||
finally:
|
||||
if creator is not None and not creator.done():
|
||||
creator.cancel()
|
||||
victim_hang.set()
|
||||
victim_task.cancel()
|
||||
try:
|
||||
await victim_task
|
||||
except BaseException:
|
||||
pass
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_joiner_follows_creation_outcome_when_creator_is_cancelled():
|
||||
"""A joiner must follow the creation's outcome, never hold an orphan.
|
||||
|
||||
If the creator is cancelled while the shared owner is still initializing,
|
||||
the creation is aborted before it commits: the joiner fails with the same
|
||||
cancellation instead of receiving (or hanging on) a session whose teardown
|
||||
is already underway. Drift guard for the outcome-gated join semantics.
|
||||
"""
|
||||
pool = MCPSessionPool()
|
||||
gate = asyncio.Event()
|
||||
cms: list[_BlockingInitCm] = []
|
||||
|
||||
def make_cm(*a, **kw):
|
||||
cm = _BlockingInitCm(gate)
|
||||
cms.append(cm)
|
||||
return cm
|
||||
|
||||
with patch("langchain_mcp_adapters.sessions.create_session", side_effect=make_cm):
|
||||
conn = {"transport": "stdio", "command": "x", "args": []}
|
||||
creator = asyncio.create_task(pool.get_session("s", "same", conn))
|
||||
for _ in range(100):
|
||||
if ("s", "same") in pool._inflight:
|
||||
break
|
||||
await asyncio.sleep(0.01)
|
||||
assert ("s", "same") in pool._inflight
|
||||
|
||||
# Second caller joins the in-flight creation instead of duplicating it.
|
||||
joiner = asyncio.create_task(pool.get_session("s", "same", conn))
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
# The creator is cancelled mid-init: the creation aborts, and the
|
||||
# joiner must observe that outcome rather than succeed or hang.
|
||||
creator.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await creator
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await asyncio.wait_for(asyncio.shield(joiner), timeout=5.0)
|
||||
|
||||
assert len(cms) == 1, "the joiner must not have created a duplicate session"
|
||||
for _ in range(100):
|
||||
if cms[0].closed:
|
||||
break
|
||||
await asyncio.sleep(0.01)
|
||||
assert cms[0].closed, "aborted creation must still run its __aexit__"
|
||||
assert len(pool._entries) == 0
|
||||
assert len(pool._inflight) == 0
|
||||
|
||||
current = asyncio.current_task()
|
||||
leaked = [t for t in asyncio.all_tasks() if t is not current and not t.done() and "_run_session" in str(t.get_coro())]
|
||||
assert not leaked, "owner task must not be left pending after the creation aborted"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_eviction_teardown_completes_normally_then_session_is_returned():
|
||||
"""Non-cancelled path: eviction teardown runs to completion and the caller
|
||||
proceeds to Phase 3 unchanged — guards the Phase-2 try/except against
|
||||
happy-path drift."""
|
||||
pool = MCPSessionPool()
|
||||
pool.MAX_SESSIONS = 1
|
||||
|
||||
class CmFactory:
|
||||
def __init__(self):
|
||||
self.closed = False
|
||||
|
||||
async def __aenter__(self):
|
||||
return AsyncMock()
|
||||
|
||||
async def __aexit__(self, *args):
|
||||
self.closed = True
|
||||
return False
|
||||
|
||||
cms: list[CmFactory] = []
|
||||
|
||||
def make_cm(*a, **kw):
|
||||
cm = CmFactory()
|
||||
cms.append(cm)
|
||||
return cm
|
||||
|
||||
with patch("langchain_mcp_adapters.sessions.create_session", side_effect=make_cm):
|
||||
first = await pool.get_session("s", "t1", {"transport": "stdio", "command": "x", "args": []})
|
||||
# Pool is full (1): this call evicts t1 through Phase 2 and then
|
||||
# returns the fresh session through Phase 3/4.
|
||||
second = await pool.get_session("s", "t2", {"transport": "stdio", "command": "x", "args": []})
|
||||
|
||||
assert first is not second
|
||||
assert cms[0].closed, "evicted owner must complete teardown before the caller proceeds"
|
||||
assert ("s", "t2") in pool._entries
|
||||
assert ("s", "t1") not in pool._entries
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_close_scope_does_not_cancel_owner_already_unwinding_in_aexit():
|
||||
"""A failed in-flight owner is already running __aexit__ in its own task;
|
||||
the close paths must not cancel it — the cancel would interrupt that
|
||||
in-task cleanup and the exit would never finish."""
|
||||
pool = MCPSessionPool()
|
||||
cms: list[_InitFailCm] = []
|
||||
|
||||
def make_cm(*args, **kwargs):
|
||||
cm = _InitFailCm()
|
||||
cms.append(cm)
|
||||
return cm
|
||||
|
||||
with patch("langchain_mcp_adapters.sessions.create_session", side_effect=make_cm):
|
||||
call = asyncio.create_task(pool.get_session("s", "t1", {"transport": "stdio", "command": "x", "args": []}))
|
||||
# Wait until the owner failed initialize() and entered its slow __aexit__.
|
||||
for _ in range(100):
|
||||
if cms and cms[0].exit_started:
|
||||
break
|
||||
await asyncio.sleep(0.01)
|
||||
assert cms and cms[0].exit_started, "owner must reach its slow __aexit__ first"
|
||||
|
||||
await pool.close_scope("t1")
|
||||
|
||||
with pytest.raises(RuntimeError):
|
||||
await call
|
||||
|
||||
assert cms[0].closed is True, "close_scope must not cancel an owner already unwinding in __aexit__"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancelling_close_scope_does_not_strand_other_removed_owners():
|
||||
"""Cancelling close_scope mid-teardown must not strand the other owners it
|
||||
already removed from the registry: every removed owner gets its close
|
||||
signal before any teardown is awaited."""
|
||||
pool = MCPSessionPool()
|
||||
loop = asyncio.get_running_loop()
|
||||
|
||||
first_hang = asyncio.Event() # first owner's slow teardown stand-in
|
||||
first_close = asyncio.Event()
|
||||
second_close = asyncio.Event()
|
||||
second_done = asyncio.Event()
|
||||
|
||||
async def first_owner() -> None:
|
||||
await first_close.wait()
|
||||
await first_hang.wait()
|
||||
|
||||
async def second_owner() -> None:
|
||||
await second_close.wait()
|
||||
second_done.set()
|
||||
|
||||
first_task = asyncio.create_task(first_owner())
|
||||
second_task = asyncio.create_task(second_owner())
|
||||
pool._entries[("s1", "scope")] = (MagicMock(), loop, first_task, first_close)
|
||||
pool._entries[("s2", "scope")] = (MagicMock(), loop, second_task, second_close)
|
||||
|
||||
closer = asyncio.create_task(pool.close_scope("scope"))
|
||||
# Let the closer reach its teardown await (both signals are already out).
|
||||
for _ in range(100):
|
||||
if first_close.is_set() and second_close.is_set():
|
||||
break
|
||||
await asyncio.sleep(0.01)
|
||||
await asyncio.sleep(0.02)
|
||||
|
||||
closer.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await closer
|
||||
|
||||
assert second_close.is_set(), "every removed owner must be signalled before teardown awaits"
|
||||
for _ in range(100):
|
||||
if second_done.is_set():
|
||||
break
|
||||
await asyncio.sleep(0.01)
|
||||
assert second_done.is_set(), "the second owner must finish its own teardown"
|
||||
|
||||
first_hang.set()
|
||||
first_task.cancel()
|
||||
try:
|
||||
await first_task
|
||||
except BaseException:
|
||||
pass
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_owner_mid_aexit_survives_gc_after_closer_cancellation():
|
||||
"""The reaper spawned for a cancelled closer must be strongly retained.
|
||||
|
||||
The event loop holds only weak references to tasks, so an unheld reaper —
|
||||
and transitively the owner it awaits, mid-``__aexit__`` — can be
|
||||
garbage-collected before the teardown completes once the registry entry
|
||||
is gone (#5008 review). The closer scenario runs in its own dropped task
|
||||
so no harness frame (including the caught CancelledError's traceback)
|
||||
keeps the teardown alive."""
|
||||
pool = MCPSessionPool()
|
||||
|
||||
class GcProneCm:
|
||||
def __init__(self):
|
||||
self.entered = False
|
||||
self.exit_started = False
|
||||
self.closed = False
|
||||
|
||||
async def __aenter__(self):
|
||||
self.entered = True
|
||||
session = MagicMock()
|
||||
session.initialize = self._initialize
|
||||
return session
|
||||
|
||||
async def _initialize(self):
|
||||
return None
|
||||
|
||||
async def __aexit__(self, *args):
|
||||
self.exit_started = True
|
||||
# An orphaned future: nothing external references it, so the
|
||||
# awaiting owner forms a collectable cycle once no strong root
|
||||
# holds the reaper.
|
||||
await asyncio.get_running_loop().create_future()
|
||||
self.closed = True
|
||||
|
||||
holder: dict[str, GcProneCm] = {}
|
||||
|
||||
def make_cm(*args, **kwargs):
|
||||
cm = GcProneCm()
|
||||
holder["cm"] = cm
|
||||
return cm
|
||||
|
||||
async def scenario() -> None:
|
||||
closer = asyncio.create_task(pool.close_scope("t1"))
|
||||
for _ in range(200):
|
||||
cm = cm_weak()
|
||||
if cm is not None and cm.exit_started:
|
||||
break
|
||||
await asyncio.sleep(0.01)
|
||||
closer.cancel()
|
||||
try:
|
||||
await closer
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
with patch("langchain_mcp_adapters.sessions.create_session", side_effect=make_cm):
|
||||
await pool.get_session("s", "t1", {"transport": "stdio", "command": "x", "args": []})
|
||||
|
||||
cm_weak = weakref.ref(holder["cm"])
|
||||
holder.clear()
|
||||
|
||||
scen = asyncio.create_task(scenario())
|
||||
await scen
|
||||
del scen
|
||||
|
||||
# Drop every remaining strong root and force collection: the registry
|
||||
# entry is gone (close_scope popped it), the scenario task and its
|
||||
# frames are gone, and the only thing that may keep the teardown alive
|
||||
# is the retained reaper.
|
||||
gc.collect()
|
||||
|
||||
assert cm_weak() is not None, "owner mid-__aexit__ must survive GC while teardown is pending"
|
||||
|
||||
# Cleanup: cancel the surviving owner/reaper tasks.
|
||||
for t in [t for t in asyncio.all_tasks() if t is not asyncio.current_task() and ("_run_session" in str(t.get_coro()) or "_reap" in str(t.get_coro()))]:
|
||||
t.cancel()
|
||||
await asyncio.sleep(0)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_queued_cancel_rechecks_failure_on_owning_loop():
|
||||
"""A cancellation queued from another loop must recheck the owner's failure
|
||||
state ON the owning loop, immediately before cancelling.
|
||||
|
||||
Reproduces the time-of-check/time-of-use window: the guard snapshot is
|
||||
taken while the owner has not failed yet, the queued cancellation is
|
||||
delivered (delayed here deterministically) only after the owner failed and
|
||||
parked inside ``__aexit__`` — delivering it must be a no-op, not an
|
||||
interruption of that cleanup (#5008 review)."""
|
||||
pool = MCPSessionPool()
|
||||
loop2 = asyncio.new_event_loop()
|
||||
thread2 = threading.Thread(target=loop2.run_forever, daemon=True)
|
||||
thread2.start()
|
||||
|
||||
class _DelayedLoop:
|
||||
"""Wraps the foreign loop; while holding, callbacks queue locally.
|
||||
|
||||
``is_running()`` reports False so the pool's await phase takes its
|
||||
harmless idle-loop branch instead of needing a real loop handle.
|
||||
"""
|
||||
|
||||
def __init__(self, wrapped: asyncio.AbstractEventLoop) -> None:
|
||||
self._wrapped = wrapped
|
||||
self.hold = False
|
||||
self._held: list[tuple] = []
|
||||
|
||||
def is_closed(self) -> bool:
|
||||
return self._wrapped.is_closed()
|
||||
|
||||
def is_running(self) -> bool:
|
||||
return False
|
||||
|
||||
def call_soon_threadsafe(self, callback, *args) -> None:
|
||||
if self.hold:
|
||||
self._held.append((callback, args))
|
||||
else:
|
||||
self._wrapped.call_soon_threadsafe(callback, *args)
|
||||
|
||||
def flush(self) -> None:
|
||||
held, self._held = self._held, []
|
||||
for callback, args in held:
|
||||
self._wrapped.call_soon_threadsafe(callback, *args)
|
||||
|
||||
class _GatedAexitCm:
|
||||
def __init__(self):
|
||||
self.exit_started = False
|
||||
self.closed = False
|
||||
|
||||
async def __aexit__(self, *args):
|
||||
self.exit_started = True
|
||||
await gate2.wait()
|
||||
self.closed = True
|
||||
|
||||
cm = _GatedAexitCm()
|
||||
holder: dict[str, object] = {}
|
||||
|
||||
def _spawn() -> None:
|
||||
ready2 = loop2.create_future()
|
||||
release2 = asyncio.Event()
|
||||
close_evt2 = asyncio.Event()
|
||||
|
||||
async def owner() -> None:
|
||||
try:
|
||||
await release2.wait()
|
||||
raise RuntimeError("init boom")
|
||||
except BaseException as exc:
|
||||
if not ready2.done():
|
||||
ready2.set_exception(exc)
|
||||
finally:
|
||||
await cm.__aexit__(None, None, None)
|
||||
|
||||
task2 = asyncio.ensure_future(owner())
|
||||
holder.update(ready=ready2, task=task2, close_evt=close_evt2, release=release2, gate=asyncio.Event())
|
||||
|
||||
gate2 = None
|
||||
loop2.call_soon_threadsafe(_spawn)
|
||||
for _ in range(200):
|
||||
if "task" in holder:
|
||||
break
|
||||
await asyncio.sleep(0.005)
|
||||
assert "task" in holder, "foreign owner must be spawned"
|
||||
gate2 = holder["gate"]
|
||||
|
||||
proxy = _DelayedLoop(loop2)
|
||||
pool._inflight[("s", "t1")] = (proxy, holder["ready"], holder["task"], holder["close_evt"])
|
||||
|
||||
# Snapshot moment: the owner has NOT failed yet. close_scope's signal phase
|
||||
# queues the (guarded) cancellation into the holding proxy; its await phase
|
||||
# takes the idle branch and returns.
|
||||
proxy.hold = True
|
||||
await pool.close_scope("t1")
|
||||
assert pool._inflight == {}
|
||||
|
||||
# The owner now fails on its own loop and parks inside the gated __aexit__.
|
||||
loop2.call_soon_threadsafe(holder["release"].set)
|
||||
for _ in range(400):
|
||||
if cm.exit_started:
|
||||
break
|
||||
await asyncio.sleep(0.005)
|
||||
assert cm.exit_started, "owner must enter its gated __aexit__ before the cancel is delivered"
|
||||
|
||||
# Deliver the held cancellation — this is the stale-snapshot catch-up the
|
||||
# on-owning-loop recheck must neutralize.
|
||||
proxy.hold = False
|
||||
proxy.flush()
|
||||
|
||||
# Let the exit finish: the delivered cancel must have been skipped.
|
||||
loop2.call_soon_threadsafe(gate2.set)
|
||||
for _ in range(400):
|
||||
if cm.closed:
|
||||
break
|
||||
await asyncio.sleep(0.005)
|
||||
assert cm.closed, "queued cancel must not interrupt an owner already unwinding in __aexit__"
|
||||
|
||||
loop2.call_soon_threadsafe(loop2.stop)
|
||||
thread2.join(timeout=2)
|
||||
loop2.close()
|
||||
|
||||
|
||||
class _InitFailCm:
|
||||
"""Fake session CM whose ``initialize`` fails, with a slow ``__aexit__``.
|
||||
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user