diff --git a/backend/packages/harness/deerflow/mcp/session_pool.py b/backend/packages/harness/deerflow/mcp/session_pool.py index 64fa091ce..f03dc7aee 100644 --- a/backend/packages/harness/deerflow/mcp/session_pool.py +++ b/backend/packages/harness/deerflow/mcp/session_pool.py @@ -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) diff --git a/backend/tests/test_mcp_session_pool.py b/backend/tests/test_mcp_session_pool.py index b7337c7a5..16132fd5f 100644 --- a/backend/tests/test_mcp_session_pool.py +++ b/backend/tests/test_mcp_session_pool.py @@ -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__``.