From b58aea31013fe029e6c5ca5e59d5d902dc7ef0e4 Mon Sep 17 00:00:00 2001 From: spud <92900806+jamespud@users.noreply.github.com> Date: Thu, 8 Oct 2026 09:02:47 +0800 Subject: [PATCH] fix(mcp): fence superseded stdio server bindings (#6435) * test(mcp): pin stdio binding fencing * fix(mcp): fence superseded stdio server bindings * docs(mcp): document stdio binding fence * fix(mcp): address binding fencing review * fix(mcp): bind discovery to cache-owned pool --- .../packages/harness/deerflow/mcp/AGENTS.md | 1 + .../packages/harness/deerflow/mcp/cache.py | 31 +- .../harness/deerflow/mcp/session_pool.py | 396 ++++++++++++----- .../harness/deerflow/mcp/task_tool_caller.py | 26 +- .../packages/harness/deerflow/mcp/tools.py | 70 ++- backend/tests/test_mcp_binding_fencing.py | 413 ++++++++++++++++++ backend/tests/test_mcp_cache.py | 93 +++- 7 files changed, 896 insertions(+), 134 deletions(-) create mode 100644 backend/tests/test_mcp_binding_fencing.py diff --git a/backend/packages/harness/deerflow/mcp/AGENTS.md b/backend/packages/harness/deerflow/mcp/AGENTS.md index b8f77362db..fa9912e474 100644 --- a/backend/packages/harness/deerflow/mcp/AGENTS.md +++ b/backend/packages/harness/deerflow/mcp/AGENTS.md @@ -34,6 +34,7 @@ and can be removed by LRU pressure or explicit cleanup. This cannot close the transport after its loop is gone. Normal `asyncio.run()` drains its owners. - **Persistent stdio session capacity**: `MCPSessionPool.MAX_SESSIONS` is a hard cap on the live LRU registry. Capacity is enforced both before creating a session and when an in-flight session is promoted, because different keys can finish initialization concurrently after all observing spare capacity. Promotion-time victims are signalled and drained by separately tracked teardown tasks on their owning loops; the newly promoted owner must never await a victim, so a blocked victim exit cannot prevent its own closure or disconnect recovery. +- **Persistent stdio binding fence**: pooled session registry keys remain `(server_name, scope_key, owning_loop, ownership_domain)`. Separately, each `(ownership_domain, server_name)` has an opaque `ServerBinding` capability derived from a SHA-256 digest of the base stdio process identity (`transport`, `command`, `args`, configured `cwd`, configured `env`; the resolved values themselves are never stored in the binding repr). Discovery captures the exact pool plus binding before its first await, and durable stdio calls bind before per-call workspace `cwd`/`TMPDIR` augmentation. `get_session()` rejects bindings that are already superseded at admission, owner promotion, or the final post-await check. That final check is a check-before-return guard, not a lease across subsequent `call_tool()` use: a reconciliation racing after it detaches and closes the owner, and the caller then fails through the existing transport-disconnect handling. Binding reconciliation consumes already-classified active fingerprints/removals, atomically installs epochs/tombstones and detaches matching owners under the pool lock, and keeps those detached owners in the pool's teardown reaper so their `__aexit__` stays observable even when the caller drops the returned `PreparedRetirement`; configuration diffing remains a higher-layer responsibility. - **Cache invalidation**: Staleness is detected in two stages. First, the resolved config path and a `(mtime, size, sha256)` content signature are compared against the values recorded when the cache was published, not a strict mtime `>` comparison — this catches same-second edits, mtime that stays put or moves backward (`git checkout`, `cp -p` / backup restore, `tar` / `rsync`, object-store / network mounts), and a switch to a different config file with an equal-or-older mtime. The signature helper (`config/file_signature.py::get_config_signature`) is shared with `config/app_config.py::get_app_config()` for the sibling runtime-editable config file, rather than each maintaining its own copy. Second, because the file also carries `skills` and middleware settings, a changed signature alone is not enough: `mcp/cache.py::_effective_mcp_config_snapshot()` serializes the effective MCP slice (enabled servers' parsed config in declaration order plus `mcpInterceptors`; the raw `transport` alias is dropped so equivalent `type`/`transport` spellings compare equal) and the cache is retired only when that slice changed. A skills-only or middleware-only edit therefore adopts the new signature and keeps the cached tools and pooled sessions, while any real MCP change keeps the existing process-wide reset through `_reset_mcp_tools_cache_state_and_retire_pool_locked()`. The snapshot may embed resolved credentials, so it stays in process memory and is never logged or persisted. `_is_cache_stale()` adopts a new signature when the MCP slice is unchanged, so it is not a read-only predicate; production callers run it under `_init_condition`. `ExtensionsConfig.resolve_config_path()` raises `FileNotFoundError` for an explicit `config_path`/`DEER_FLOW_EXTENSIONS_CONFIG_PATH` that points at a missing file — an operator-asserted path going missing is a real misconfiguration, so this is intentionally loud for callers that load the config for actual use (e.g. `from_file()` via `get_mcp_tools()`); only the fallback search mode returns `None`. The MCP cache's own path resolution (`mcp/cache.py::_resolve_config_path`) is narrower: it catches that specific `FileNotFoundError` locally and treats it the same as "unconfigured", so this staleness check degrades to "not stale" instead of propagating an exception when a previously-valid explicit/env-var config disappears mid-run. Two no-config states are distinguished: `_initialized_without_config` makes a cache published while no config existed stale as soon as one appears, while a config deleted after a successful load keeps the last-known-good fail-soft behavior and stays "not stale", while a file that is present but whose content cannot be hashed is treated as stale instead, because the cache cannot prove the revision it published is still the current one. `initialize_mcp_tools()` reads one `ExtensionsConfig` instance, hands that exact object to `get_mcp_tools(extensions_config=...)`, and publishes only when the effective MCP snapshot of the current config still equals the snapshot of the revision that produced the tools; otherwise it resets the tool cache through the same session-pool retirement path as normal stale invalidation before waiters retry. This closes the discovery window — tools built from a superseded MCP revision, or from a resolved path whose signature cannot be re-read, are never published — while a skills-only edit during discovery still publishes because the MCP slice is unchanged. - **Transports**: stdio (command-based), SSE, HTTP - **Per-server tool-name prefixing**: `mcpServers..tool_name_prefix` defaults to `true`, preserving the collision-safe `_` prefix. Servers whose tools already carry a stable namespace may set it to `false`; discovery then calls `langchain_mcp_adapters.tools.load_mcp_tools` with that server's flag. Source routing and stdio session-pool wrapping are based on the producing server and transport, never on whether the visible tool name starts with the server prefix. diff --git a/backend/packages/harness/deerflow/mcp/cache.py b/backend/packages/harness/deerflow/mcp/cache.py index f3a62b3418..921d6d0982 100644 --- a/backend/packages/harness/deerflow/mcp/cache.py +++ b/backend/packages/harness/deerflow/mcp/cache.py @@ -300,6 +300,7 @@ async def initialize_mcp_tools() -> list[BaseTool]: await asyncio.to_thread(_wait_for_initialization, waiting_generation) from deerflow.config.extensions_config import ExtensionsConfig + from deerflow.mcp.session_pool import StaleMCPBindingError, get_session_pool from deerflow.mcp.tools import get_mcp_tools loaded_tools = None @@ -329,7 +330,35 @@ async def initialize_mcp_tools() -> list[BaseTool]: raise RuntimeError("Extensions config could not be loaded for MCP tool discovery") from None loaded_snapshot = _effective_mcp_config_snapshot(loaded_config) loaded_reset_signature = _current_cache_reset_marker_signature(_resolve_config_path()) - loaded_tools = await get_mcp_tools(extensions_config=loaded_config) + # Claim the exact pool this generation owns under the same lock the + # reset path takes, so verify-generation + capture-pool is one atomic + # ownership handoff. A superseded initializer must never resolve the + # singleton after the reset: that would install its stale fingerprint + # into the replacement pool the successor initializer runs on. + with _init_condition: + if _cache_generation != claim_generation: + logger.info("MCP cache was reset before tool discovery; discarding superseded initialization") + return [] + + claimed_pool = get_session_pool() + + try: + loaded_tools = await get_mcp_tools( + extensions_config=loaded_config, + session_pool=claimed_pool, + ) + except StaleMCPBindingError: + # The pool this claim owns was retired while discovery was running. + # That is the same cache race as a generation change, so keep the + # existing "discard stale initialization" semantics; a stale binding + # on a pool this claim still owns stays a real error. + with _init_condition: + superseded = _cache_generation != claim_generation + + if superseded: + logger.info("MCP cache was reset during binding installation; discarding superseded initialization") + return [] + raise post_path, post_sig = _current_config_state() post_reset_signature = _current_cache_reset_marker_signature(post_path) if post_path is not None and post_sig is not None: diff --git a/backend/packages/harness/deerflow/mcp/session_pool.py b/backend/packages/harness/deerflow/mcp/session_pool.py index 7a74f91924..83a8d896ec 100644 --- a/backend/packages/harness/deerflow/mcp/session_pool.py +++ b/backend/packages/harness/deerflow/mcp/session_pool.py @@ -42,9 +42,13 @@ single caller that might get cancelled. from __future__ import annotations import asyncio +import hashlib +import json import logging import threading from collections import OrderedDict +from collections.abc import Collection, Mapping +from dataclasses import dataclass, field from typing import Any, Literal import anyio @@ -58,6 +62,53 @@ logger = logging.getLogger(__name__) MCPPoolDomain = Literal["deployment", "personal"] +_STDIO_IDENTITY_KEYS = ("transport", "command", "args", "cwd", "env") + + +def normalized_connection_fingerprint(connection: Mapping[str, Any]) -> str: + """Return a secret-safe digest of the stdio process identity.""" + identity = {key: connection[key] for key in _STDIO_IDENTITY_KEYS if key in connection} + canonical = json.dumps(identity, sort_keys=True, separators=(",", ":"), ensure_ascii=False, default=str) + return hashlib.sha256(canonical.encode("utf-8")).hexdigest() + + +@dataclass(frozen=True, slots=True, eq=False) +class ServerBinding: + """Opaque capability naming one configuration epoch for a pooled server.""" + + domain: MCPPoolDomain + server_name: str + epoch: int + fingerprint: str | None = field(repr=False) + + +@dataclass(frozen=True, slots=True) +class PreparedRetirement: + """Owners atomically detached and signalled by a binding transition.""" + + entries: tuple[ + tuple[ClientSession, asyncio.AbstractEventLoop, asyncio.Task[Any], asyncio.Event], + ..., + ] = () + inflight: tuple[ + tuple[ + asyncio.AbstractEventLoop, + asyncio.Future[ClientSession], + asyncio.Task[Any], + asyncio.Event, + ], + ..., + ] = () + + +class StaleMCPBindingError(RuntimeError): + """Raised when a caller holds a superseded stdio server binding.""" + + def __init__(self, server_name: str) -> None: + super().__init__(f"MCP server '{server_name}' configuration changed; rebuild the MCP tools before calling it again") + self.server_name = server_name + + _MCP_CLOSED_STREAM_ERRORS = ( anyio.ClosedResourceError, anyio.BrokenResourceError, @@ -168,6 +219,118 @@ class MCPSessionPool: # garbage-collected before teardown completes; the done callback keeps # the set from growing without bound. self._teardown_tasks: set[asyncio.Task[Any]] = set() + # Per-(ownership-domain, server-name) configuration capability. + # Session registry keys remain the four-tuple introduced by PR A. + self._bindings: dict[tuple[MCPPoolDomain, str], ServerBinding] = {} + self._binding_lifecycle_servers: set[tuple[MCPPoolDomain, str]] = set() + self._next_binding_epoch = 1 + self._retired = False + + # ------------------------------------------------------------------ + # Server binding identity + # ------------------------------------------------------------------ + + def _install_binding_locked( + self, + binding_key: tuple[MCPPoolDomain, str], + fingerprint: str | None, + ) -> ServerBinding: + domain, server_name = binding_key + binding = ServerBinding( + domain=domain, + server_name=server_name, + epoch=self._next_binding_epoch, + fingerprint=fingerprint, + ) + self._next_binding_epoch += 1 + self._bindings[binding_key] = binding + return binding + + def _binding_is_current(self, expected: ServerBinding) -> bool: + with self._lock: + return not self._retired and self._bindings.get((expected.domain, expected.server_name)) is expected + + def ensure_binding( + self, + server_name: str, + fingerprint: str, + *, + domain: MCPPoolDomain = "deployment", + ) -> ServerBinding: + """Resolve the discovery/task binding atomically without overwriting newer state.""" + binding_key = (domain, server_name) + with self._lock: + self._binding_lifecycle_servers.add(binding_key) + if self._retired: + raise StaleMCPBindingError(server_name) + current = self._bindings.get(binding_key) + if current is None: + return self._install_binding_locked(binding_key, fingerprint) + if current.fingerprint != fingerprint: + raise StaleMCPBindingError(server_name) + return current + + def reconcile_bindings( + self, + active: Mapping[str, str], + removed: Collection[str], + *, + domain: MCPPoolDomain = "deployment", + ) -> PreparedRetirement: + """Apply already-classified binding changes and detach only changed owners. + + This method deliberately does not inspect configuration. A higher layer + supplies current fingerprints and removed names. Epoch installation, + registry detach and owner signalling are one non-awaiting critical + section, so a stale creator cannot resurrect after the transition. + """ + entries: list[tuple[ClientSession, asyncio.AbstractEventLoop, asyncio.Task[Any], asyncio.Event]] = [] + inflight: list[tuple[asyncio.AbstractEventLoop, asyncio.Future[ClientSession], asyncio.Task[Any], asyncio.Event]] = [] + changed: set[str] = set() + + with self._lock: + for server_name, fingerprint in active.items(): + binding_key = (domain, server_name) + self._binding_lifecycle_servers.add(binding_key) + current = self._bindings.get(binding_key) + if current is not None and current.fingerprint == fingerprint: + continue + self._install_binding_locked(binding_key, fingerprint) + changed.add(server_name) + + for server_name in removed: + binding_key = (domain, server_name) + self._binding_lifecycle_servers.add(binding_key) + current = self._bindings.get(binding_key) + if current is not None and current.fingerprint is None: + continue + self._install_binding_locked(binding_key, None) + changed.add(server_name) + + for key in [k for k in self._entries if k[3] == domain and k[0] in changed]: + entries.append(self._entries.pop(key)) + for key in [k for k in self._inflight if k[3] == domain and k[0] in changed]: + inflight.append(self._inflight.pop(key)) + + for _session, loop, _task, close_evt in entries: + self._signal_close(loop, close_evt) + for loop, ready, task, close_evt in inflight: + self._signal_close(loop, close_evt) + self._cancel_owner(loop, task, ready) + + # Owners are already out of the registry, so retain their teardown even + # if the caller never awaits the returned PreparedRetirement. + for _session, loop, task, _close_evt in entries: + self._track_detached_owner(loop, task) + for loop, _ready, task, _close_evt in inflight: + self._track_detached_owner(loop, task) + + return PreparedRetirement(entries=tuple(entries), inflight=tuple(inflight)) + + def retire_all(self) -> None: + """Fence this pool so stale wrappers cannot create sessions after reset.""" + with self._lock: + self._retired = True # ------------------------------------------------------------------ # Session owner task @@ -193,54 +356,38 @@ class MCPSessionPool: connection: dict[str, Any], ready: asyncio.Future[ClientSession], close_evt: asyncio.Event, + expected_binding: ServerBinding, ) -> None: - """Own a single MCP session for its entire lifetime. - - 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. - """ + """Own one MCP session and fence promotion against binding replacement.""" from langchain_mcp_adapters.sessions import create_session cm = create_session(connection) try: session = await cm.__aenter__() - except BaseException as e: - # Never entered the cancel scope, so there is nothing to exit. + except BaseException as error: if not ready.done(): - ready.set_exception(e) + if self._binding_is_current(expected_binding): + ready.set_exception(error) + else: + logger.debug( + "Suppressing MCP session error for superseded binding '%s'", + expected_binding.server_name, + exc_info=(type(error), error, error.__traceback__), + ) + ready.set_exception(StaleMCPBindingError(expected_binding.server_name)) return - # The context manager is now entered. From here on __aexit__ MUST run in - # this task — on init failure, on cancellation, or on the close signal — - # to satisfy anyio's same-task cancel-scope requirement and to avoid - # leaking the session/subprocess. try: await session.initialize() - # 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() promoted_evicted: list[tuple[asyncio.AbstractEventLoop, asyncio.Task[Any], asyncio.Event]] = [] + binding_key = (expected_binding.domain, expected_binding.server_name) with self._lock: + binding_ok = not self._retired and self._bindings.get(binding_key) is expected_binding still_ours = self._inflight.get(key) == (loop, ready, task, close_evt) - if still_ours: + if still_ours and binding_ok: self._inflight.pop(key) - # Different keys can finish initialization concurrently. - # They may all pass the earlier capacity check while no - # session is registered, so enforce the cap again at the - # single owner-controlled commit point. while len(self._entries) >= self.MAX_SESSIONS: oldest_key, (_, ent_loop, ent_task, ent_close) = next(iter(self._entries.items())) self._entries.pop(oldest_key) @@ -248,26 +395,40 @@ class MCPSessionPool: self._entries[key] = (session, loop, task, close_evt) if not ready.done(): ready.set_result(session) - if still_ours: - # Drain victims independently: a blocked victim's __aexit__ - # must not prevent this owner from handling its own close. + elif still_ours: + self._inflight.pop(key) + + if still_ours and binding_ok: for ent_loop, _ent_task, ent_close in promoted_evicted: self._signal_close(ent_loop, ent_close) - for ent_loop, ent_task, ent_close in promoted_evicted: + for ent_loop, ent_task, _ent_close in promoted_evicted: if ent_loop is loop: self._track_owner_teardown(ent_task) elif not ent_loop.is_closed(): try: ent_loop.call_soon_threadsafe(self._track_owner_teardown, ent_task) except RuntimeError: - pass # The owning loop closed before scheduling. + pass 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: + await close_evt.wait() + elif not binding_ok: + if not ready.done(): + ready.set_exception(StaleMCPBindingError(expected_binding.server_name)) + else: + if 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 error: if not ready.done(): - ready.set_exception(e) + if self._binding_is_current(expected_binding): + ready.set_exception(error) + else: + logger.debug( + "Suppressing MCP session error for superseded binding '%s'", + expected_binding.server_name, + exc_info=(type(error), error, error.__traceback__), + ) + ready.set_exception(StaleMCPBindingError(expected_binding.server_name)) finally: try: await cm.__aexit__(None, None, None) @@ -280,36 +441,46 @@ class MCPSessionPool: scope_key: str, connection: dict[str, Any], *, + binding: ServerBinding | None = None, domain: MCPPoolDomain = "deployment", ) -> ClientSession: - """Get or create a persistent MCP session. + """Get/create a pooled session while proving the caller owns the current binding.""" + effective_domain = binding.domain if binding is not None else domain + if binding is not None and binding.server_name != server_name: + raise StaleMCPBindingError(server_name) - Reuse sessions and in-flight creations only on the current loop. - Other live loops retain their own independent sessions. - - Args: - server_name: MCP server name. - scope_key: Isolation key (typically thread_id). - connection: Connection configuration for ``create_session``. - domain: Ownership namespace for the connection. Deployment and - personal resources with the same server/scope remain isolated. - - Returns: - An initialized ``ClientSession``. - """ current_loop = asyncio.get_running_loop() - key = (server_name, scope_key, current_loop, domain) + key = (server_name, scope_key, current_loop, effective_domain) + binding_key = (effective_domain, server_name) - # 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. - # LRU victims are established sessions: (loop, owner_task, close_event). evicted: list[tuple[asyncio.AbstractEventLoop, asyncio.Task[Any], asyncio.Event]] = [] join: asyncio.Future[ClientSession] | None = None ready: asyncio.Future[ClientSession] | None = None close_evt: asyncio.Event | None = None task: asyncio.Task[Any] | None = None + with self._lock: + if self._retired: + raise StaleMCPBindingError(server_name) + + if binding is not None: + if self._bindings.get(binding_key) is not binding: + raise StaleMCPBindingError(server_name) + expected_binding = binding + else: + # Compatibility for callers not yet binding-aware. Once a + # server participates in explicit binding lifecycle, callers + # must present the opaque capability; fingerprint equality + # alone cannot distinguish remove/re-add ABA. + if binding_key in self._binding_lifecycle_servers: + raise StaleMCPBindingError(server_name) + expected_binding = self._bindings.get(binding_key) + fingerprint = normalized_connection_fingerprint(connection) + if expected_binding is None: + expected_binding = self._install_binding_locked(binding_key, fingerprint) + elif expected_binding.fingerprint != fingerprint: + raise StaleMCPBindingError(server_name) + if key in self._entries: session, _loop, ent_task, _ent_close = self._entries[key] if not ent_task.done(): @@ -319,29 +490,19 @@ class MCPSessionPool: inflight = self._inflight.get(key) if inflight is not None: - # Only callers on this same loop share the creation future. join = inflight[1] else: - # 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(key, connection, ready, close_evt)) + task = current_loop.create_task(self._run_session(key, connection, ready, close_evt, expected_binding)) self._inflight[key] = (current_loop, ready, task, close_evt) task.add_done_callback(lambda owner: self._discard_owner(key, owner)) - # 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)) - # Phase 2: shut down evicted sessions. 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 same-loop teardowns; foreign-loop owners finish on their - # own loops. The owner task — never this one — runs __aexit__. for loop, _ent_task, ent_close in evicted: self._signal_close(loop, ent_close) try: @@ -349,21 +510,6 @@ class MCPSessionPool: if loop is current_loop and not loop.is_closed(): await self._shutdown(ent_close, ent_task) 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: @@ -375,42 +521,20 @@ class MCPSessionPool: 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) + session = await asyncio.shield(join) + # Check-before-return fence, not a lease across subsequent caller + # use: a reconcile racing after this check detaches and closes the + # owner, and the caller then fails through the existing transport + # disconnect path. + if not self._binding_is_current(expected_binding): + raise StaleMCPBindingError(server_name) + return session assert ready is not None and close_evt is not None and task is not None - - # 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: - # 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) 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: @@ -425,8 +549,9 @@ class MCPSessionPool: self._inflight.pop(key) raise - # Phase 4: the owner task already promoted the initialized session and - # enforced capacity in the same commit critical section. + # Same check-before-return fence as the joiner path above; not a lease. + if not self._binding_is_current(expected_binding): + raise StaleMCPBindingError(server_name) return session # ------------------------------------------------------------------ @@ -500,6 +625,33 @@ class MCPSessionPool: # Loop was closed between the is_closed() check and now. pass + def _track_detached_owner( + self, + loop: asyncio.AbstractEventLoop, + task: asyncio.Task[Any], + ) -> None: + """Keep an owner detached by binding reconciliation strongly retained. + + ``reconcile_bindings()`` removes owners from the registry without + awaiting their teardown, so the pool itself must keep that teardown + observable even when the caller drops the returned + ``PreparedRetirement``. The reaper has to run on the owner's loop, which + is where ``__aexit__`` executes. + """ + try: + current_loop: asyncio.AbstractEventLoop | None = asyncio.get_running_loop() + except RuntimeError: + current_loop = None + + if loop is current_loop: + self._track_owner_teardown(task) + elif not loop.is_closed(): + try: + loop.call_soon_threadsafe(self._track_owner_teardown, task) + except RuntimeError: + # Loop was closed between the is_closed() check and now. + pass + def _track_owner_teardown(self, task: asyncio.Task[Any]) -> None: """Keep detached eviction or cancelled-caller teardown observable. @@ -799,9 +951,15 @@ def get_session_pool() -> MCPSessionPool: def reset_session_pool() -> MCPSessionPool | None: - """Reset the singleton and return the retired pool, if any.""" + """Reset the singleton and return the retired pool, if any. + + The old pool is fenced before publication of the replacement so wrappers + that captured it cannot create a post-reset session. + """ global _pool with _pool_lock: retired = _pool + if retired is not None: + retired.retire_all() _pool = None return retired diff --git a/backend/packages/harness/deerflow/mcp/task_tool_caller.py b/backend/packages/harness/deerflow/mcp/task_tool_caller.py index 8fb5166157..d21401ed92 100644 --- a/backend/packages/harness/deerflow/mcp/task_tool_caller.py +++ b/backend/packages/harness/deerflow/mcp/task_tool_caller.py @@ -25,6 +25,7 @@ from deerflow.mcp.session_pool import ( MCPSessionPool, call_pooled_session_tool, get_session_pool, + normalized_connection_fingerprint, ) from deerflow.mcp.tasks.runtime import get_mcp_task_oauth_token_manager from deerflow.mcp.user_config import PersonalMcpConfigSnapshot, load_user_mcp_config_if_changed @@ -195,15 +196,36 @@ class McpTaskToolCaller: ) if transport == "stdio": + pool = get_session_pool() + binding = None + if isinstance(pool, MCPSessionPool): + # Resolve the opaque binding from the base connection before + # per-call workspace cwd/TMP variables are added. Direct unit + # tests may inject a lightweight fake pool; the production + # singleton always takes the binding-aware path. + binding = pool.ensure_binding( + server_name, + normalized_connection_fingerprint(connection), + domain=connection_scope, + ) connection = await asyncio.to_thread( _prepare_stdio_connection, connection, user_id=user_id, thread_id=thread_id, ) - pool = get_session_pool() session_init_timeout = server_config.session_init_timeout - session_request = pool.get_session(server_name, scope_key, connection) if connection_scope == "deployment" else pool.get_session(server_name, scope_key, connection, domain=connection_scope) + if binding is not None: + session_request = pool.get_session( + server_name, + scope_key, + connection, + binding=binding, + ) + elif connection_scope == "deployment": + session_request = pool.get_session(server_name, scope_key, connection) + else: + session_request = pool.get_session(server_name, scope_key, connection, domain=connection_scope) if session_init_timeout is not None: try: session = await asyncio.wait_for( diff --git a/backend/packages/harness/deerflow/mcp/tools.py b/backend/packages/harness/deerflow/mcp/tools.py index 64863e3c9c..ee04ea2474 100644 --- a/backend/packages/harness/deerflow/mcp/tools.py +++ b/backend/packages/harness/deerflow/mcp/tools.py @@ -24,8 +24,11 @@ from deerflow.mcp.interceptors import build_mcp_tool_interceptors, compose_tool_ from deerflow.mcp.oauth import OAuthTokenManager, build_oauth_tool_interceptor, get_initial_oauth_headers from deerflow.mcp.session_pool import ( MCPPoolDomain, + MCPSessionPool, + ServerBinding, call_pooled_session_tool, get_session_pool, + normalized_connection_fingerprint, ) from deerflow.mcp.tasks import ORDINARY_MCP_TASK_DRIVER, TaskSubmitRequest from deerflow.mcp.tasks.runtime import ( @@ -545,6 +548,9 @@ def _make_session_pool_tool( session_init_timeout: float | None = None, tool_name_prefix: bool = True, ownership_domain: MCPPoolDomain = "deployment", + *, + pool: MCPSessionPool | None = None, + binding: ServerBinding | None = None, ) -> BaseTool: """Wrap an MCP tool so it reuses a persistent session from the pool. @@ -563,7 +569,15 @@ def _make_session_pool_tool( if tool_name_prefix and original_name.startswith(prefix): original_name = original_name[len(prefix) :] - pool = get_session_pool() + if pool is None: + pool = get_session_pool() + if binding is None and isinstance(pool, MCPSessionPool): + binding = pool.ensure_binding( + server_name, + normalized_connection_fingerprint(connection), + domain=ownership_domain, + ) + effective_domain = binding.domain if binding is not None else ownership_domain async def call_with_persistent_session( runtime: Runtime | None = None, @@ -613,16 +627,25 @@ def _make_session_pool_tool( session_env.setdefault("TMP", str(tmp_dir)) session_env.setdefault("TEMP", str(tmp_dir)) session_connection["env"] = session_env - session_request = ( - pool.get_session(server_name, scope_key, session_connection) - if ownership_domain == "deployment" - else pool.get_session( + if binding is not None: + session_request = pool.get_session( server_name, scope_key, session_connection, - domain=ownership_domain, + binding=binding, + ) + elif effective_domain == "deployment": + # Compatibility for direct unit-test fakes that construct this + # wrapper without going through get_mcp_tools(). Production + # discovery always supplies the explicit binding above. + session_request = pool.get_session(server_name, scope_key, session_connection) + else: + session_request = pool.get_session( + server_name, + scope_key, + session_connection, + domain=effective_domain, ) - ) if session_init_timeout is not None: # Cancellation here is safe: MCPSessionPool.get_session owns the # teardown of a session stuck mid-creation (it signals close and @@ -647,7 +670,7 @@ def _make_session_pool_tool( else: session = await session_request - domain_kwargs = {"domain": ownership_domain} if ownership_domain != "deployment" else {} + domain_kwargs = {"domain": effective_domain} if effective_domain != "deployment" else {} # Build common call_tool kwargs once — only add keys when needed so # existing call-sites that assert on exact arguments are not affected. @@ -858,7 +881,12 @@ def _configure_task_tools_for_server( return configured -async def get_mcp_tools(extensions_config: ExtensionsConfig | None = None, *, personal_user_id: str | None = None) -> list[BaseTool]: +async def get_mcp_tools( + extensions_config: ExtensionsConfig | None = None, + *, + personal_user_id: str | None = None, + session_pool: MCPSessionPool | None = None, +) -> list[BaseTool]: """Get all tools from enabled MCP servers. Tools using stdio transport are wrapped with persistent-session logic so @@ -870,6 +898,10 @@ async def get_mcp_tools(extensions_config: ExtensionsConfig | None = None, *, pe extensions_config: Optional pre-loaded extensions config. Callers that must prove which config revision produced these tools pass the exact instance they snapshotted; ``None`` loads the latest config from disk. + session_pool: Optional session pool captured by the caller. Cache + initializers pass the exact pool their generation owns so a + superseded claim cannot install bindings into a replacement pool; + ``None`` resolves the current singleton. Returns: List of LangChain tools from all enabled MCP servers. @@ -907,6 +939,22 @@ async def get_mcp_tools(extensions_config: ExtensionsConfig | None = None, *, pe logger.info("No enabled MCP servers configured") return [] + # Capture the exact pool and stdio binding before the first discovery await. + # A later reset/reconcile can supersede this capability, but a stale wrapper + # can never silently pair its old connection with a replacement pool/epoch. + pool = session_pool + if pool is None: + pool = get_session_pool() + ownership_domain: MCPPoolDomain = "personal" if personal_user_id is not None else "deployment" + server_bindings: dict[str, ServerBinding] = {} + for server_name, server_connection in servers_config.items(): + if server_connection.get("transport", "stdio") == "stdio": + server_bindings[server_name] = pool.ensure_binding( + server_name, + normalized_connection_fingerprint(server_connection), + domain=ownership_domain, + ) + try: # Create the multi-server MCP client logger.info(f"Initializing MCP client with {len(servers_config)} server(s)") @@ -1042,7 +1090,9 @@ async def get_mcp_tools(extensions_config: ExtensionsConfig | None = None, *, pe tool_call_timeout=_timeout, session_init_timeout=_init_timeout, tool_name_prefix=tool_name_prefix, - ownership_domain=("personal" if personal_user_id is not None else "deployment"), + ownership_domain=ownership_domain, + pool=pool, + binding=server_bindings[source_name], ) ) else: diff --git a/backend/tests/test_mcp_binding_fencing.py b/backend/tests/test_mcp_binding_fencing.py new file mode 100644 index 0000000000..63ca2f7679 --- /dev/null +++ b/backend/tests/test_mcp_binding_fencing.py @@ -0,0 +1,413 @@ +"""Regression tests for stdio MCP binding ownership/fencing.""" + +import asyncio +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from langchain_core.tools import StructuredTool +from pydantic import BaseModel + +from deerflow.config.extensions_config import ExtensionsConfig +from deerflow.mcp import session_pool as session_pool_module +from deerflow.mcp.session_pool import ( + MCPSessionPool, + StaleMCPBindingError, + normalized_connection_fingerprint, +) + +_CONNECTION = {"transport": "stdio", "command": "x", "args": []} + + +class _GatedSessionCm: + def __init__(self, gate: asyncio.Event | None = None) -> None: + self.gate = gate + self.initialize_started = asyncio.Event() + self.closed = False + self.enter_task = None + self.exit_task = None + self.session = MagicMock() + self.session.initialize = self.initialize + + async def __aenter__(self): + self.enter_task = asyncio.current_task() + return self.session + + async def initialize(self): + self.initialize_started.set() + if self.gate is not None: + await self.gate.wait() + + async def __aexit__(self, *_args): + self.exit_task = asyncio.current_task() + self.closed = True + + +def test_binding_digest_tracks_stdio_identity_without_leaking_env_value(): + a = normalized_connection_fingerprint({**_CONNECTION, "cwd": "/a", "env": {"TOKEN": "secret-a"}}) + b = normalized_connection_fingerprint({**_CONNECTION, "cwd": "/a", "env": {"TOKEN": "secret-b"}}) + c = normalized_connection_fingerprint({**_CONNECTION, "cwd": "/b", "env": {"TOKEN": "secret-a"}}) + + assert a != b != c + assert "secret-a" not in a + assert len(a) == 64 + + +def test_same_binding_is_idempotent_and_same_name_domains_are_independent(): + pool = MCPSessionPool() + deployment = pool.ensure_binding("A", "fp-a", domain="deployment") + + assert pool.ensure_binding("A", "fp-a", domain="deployment") is deployment + personal = pool.ensure_binding("A", "fp-a", domain="personal") + assert personal is not deployment + + +@pytest.mark.asyncio +async def test_remove_readd_same_fingerprint_never_reauthorizes_old_binding(): + pool = MCPSessionPool() + old = pool.ensure_binding("A", "same-fp") + + pool.reconcile_bindings({}, {"A"}) + tombstone = pool._bindings[("deployment", "A")] + pool.reconcile_bindings({"A": "same-fp"}, ()) + current = pool._bindings[("deployment", "A")] + + assert tombstone.fingerprint is None + assert current is not old + assert current.epoch > tombstone.epoch > old.epoch + + with pytest.raises(StaleMCPBindingError): + await pool.get_session("A", "u:t", _CONNECTION, binding=old) + + +@pytest.mark.asyncio +async def test_reconcile_detaches_only_changed_server_and_preserves_registry_key_shape(): + pool = MCPSessionPool() + loop = asyncio.get_running_loop() + a = pool.ensure_binding("A", "a1") + b = pool.ensure_binding("B", "b1") + cm_a = _GatedSessionCm() + cm_b = _GatedSessionCm() + + with patch("langchain_mcp_adapters.sessions.create_session", side_effect=[cm_a, cm_b]): + session_a = await pool.get_session("A", "u:t", _CONNECTION, binding=a) + session_b = await pool.get_session("B", "u:t", _CONNECTION, binding=b) + + prepared = pool.reconcile_bindings({"A": "a2", "B": "b1"}, ()) + assert [entry[0] for entry in prepared.entries] == [session_a] + assert prepared.inflight == () + assert pool._entries[("B", "u:t", loop, "deployment")][0] is session_b + assert cm_b.closed is False + + await asyncio.wait_for(prepared.entries[0][2], timeout=1) + assert cm_a.closed is True + assert cm_a.exit_task is cm_a.enter_task + + +@pytest.mark.asyncio +async def test_stale_binding_fails_before_session_creation(): + pool = MCPSessionPool() + old = pool.ensure_binding("A", "a1") + pool.reconcile_bindings({"A": "a2"}, ()) + + with patch("langchain_mcp_adapters.sessions.create_session") as create_session: + with pytest.raises(StaleMCPBindingError): + await pool.get_session("A", "u:t", _CONNECTION, binding=old) + + create_session.assert_not_called() + + +@pytest.mark.asyncio +async def test_reconcile_fences_creator_blocked_in_initialize(): + pool = MCPSessionPool() + old = pool.ensure_binding("A", "a1") + gate = asyncio.Event() + cm = _GatedSessionCm(gate) + + with patch("langchain_mcp_adapters.sessions.create_session", return_value=cm): + creator = asyncio.create_task(pool.get_session("A", "u:t", _CONNECTION, binding=old)) + await asyncio.wait_for(cm.initialize_started.wait(), 1) + + prepared = pool.reconcile_bindings({"A": "a2"}, ()) + assert len(prepared.inflight) == 1 + gate.set() + + with pytest.raises(StaleMCPBindingError): + await asyncio.wait_for(creator, 1) + + assert not pool._entries + assert not pool._inflight + assert cm.closed is True + assert cm.exit_task is cm.enter_task + + +@pytest.mark.asyncio +async def test_creator_and_joiner_both_observe_superseded_binding(): + pool = MCPSessionPool() + old = pool.ensure_binding("A", "a1") + gate = asyncio.Event() + cm = _GatedSessionCm(gate) + + with patch("langchain_mcp_adapters.sessions.create_session", return_value=cm): + creator = asyncio.create_task(pool.get_session("A", "u:t", _CONNECTION, binding=old)) + await asyncio.wait_for(cm.initialize_started.wait(), 1) + joiner = asyncio.create_task(pool.get_session("A", "u:t", _CONNECTION, binding=old)) + await asyncio.sleep(0) + + pool.reconcile_bindings({"A": "a2"}, ()) + gate.set() + + with pytest.raises(StaleMCPBindingError): + await asyncio.wait_for(creator, 1) + with pytest.raises(StaleMCPBindingError): + await asyncio.wait_for(joiner, 1) + + +@pytest.mark.asyncio +async def test_reset_fences_old_pool_and_old_binding(): + pool = MCPSessionPool() + binding = pool.ensure_binding("A", "a1") + session_pool_module._pool = pool + + retired = session_pool_module.reset_session_pool() + assert retired is pool + + with patch("langchain_mcp_adapters.sessions.create_session") as create_session: + with pytest.raises(StaleMCPBindingError): + await pool.get_session("A", "u:t", _CONNECTION, binding=binding) + create_session.assert_not_called() + + +class _EmptyArgs(BaseModel): + pass + + +@pytest.mark.asyncio +async def test_discovery_captures_pool_and_binding_before_first_await(): + from deerflow.mcp.tools import get_mcp_tools + + config = ExtensionsConfig.model_validate( + { + "mcpServers": { + "A": { + "type": "stdio", + "command": "x", + "args": [], + } + } + } + ) + servers = {"A": {"transport": "stdio", "command": "x", "args": []}} + pool = MagicMock() + binding = SimpleNamespace(domain="deployment", server_name="A") + order: list[str] = [] + + def ensure_binding(*_args, **_kwargs): + order.append("bind") + return binding + + pool.ensure_binding.side_effect = ensure_binding + + async def raw_tool() -> str: + return "ok" + + tool = StructuredTool( + name="A_echo", + description="echo", + args_schema=_EmptyArgs, + coroutine=raw_tool, + ) + + class FakeClient: + def __init__(self, *_args, **_kwargs): + self.callbacks = None + self.tool_interceptors = [] + + async def get_tools(self, *, server_name=None): + assert server_name == "A" + order.append("discover") + return [tool] + + with ( + patch("deerflow.mcp.tools.get_session_pool", return_value=pool), + patch("deerflow.mcp.tools.build_servers_config", return_value=servers), + patch("deerflow.mcp.tools.get_initial_oauth_headers", new_callable=AsyncMock, return_value={}), + patch("deerflow.mcp.tools.build_mcp_tool_interceptors", return_value=[]), + patch("langchain_mcp_adapters.client.MultiServerMCPClient", FakeClient), + patch("deerflow.mcp.tools._make_session_pool_tool", side_effect=lambda wrapped, *_args, **_kwargs: wrapped) as wrap, + ): + tools = await get_mcp_tools(config) + + assert order[:2] == ["bind", "discover"] + assert tools == [tool] + assert wrap.call_args.kwargs["pool"] is pool + assert wrap.call_args.kwargs["binding"] is binding + + +@pytest.mark.asyncio +async def test_durable_stdio_call_binds_base_connection_before_workspace_augmentation(): + from deerflow.mcp.task_tool_caller import McpTaskToolCaller + + config = ExtensionsConfig.model_validate( + { + "mcpServers": { + "reports": { + "type": "stdio", + "command": "report-mcp", + "args": [], + } + } + } + ) + result = SimpleNamespace(structuredContent={"status": "running"}, isError=False) + session = SimpleNamespace(call_tool=AsyncMock(return_value=result)) + pool = MCPSessionPool() + pool.get_session = AsyncMock(return_value=session) + + prepared = { + "transport": "stdio", + "command": "report-mcp", + "args": [], + "cwd": "/workspace", + "env": {"TMPDIR": "/workspace/tmp"}, + } + + caller = McpTaskToolCaller(config) + with ( + patch("deerflow.mcp.task_tool_caller.get_session_pool", return_value=pool), + patch("deerflow.mcp.task_tool_caller._prepare_stdio_connection", return_value=prepared), + ): + actual = await caller.call_tool( + server_name="reports", + tool_name="status_report", + arguments={"task_id": "1"}, + user_id="u1", + thread_id="t1", + ) + + assert actual is result + binding = pool._bindings[("deployment", "reports")] + assert binding.fingerprint == normalized_connection_fingerprint({"transport": "stdio", "command": "report-mcp", "args": []}) + pool.get_session.assert_awaited_once_with( + "reports", + "u1:t1", + prepared, + binding=binding, + ) + + +def test_ensure_binding_rejects_changed_fingerprint(): + pool = MCPSessionPool() + pool.ensure_binding("A", "fp-a") + + with pytest.raises(StaleMCPBindingError): + pool.ensure_binding("A", "fp-b") + + +def test_ensure_binding_rejects_tombstoned_server(): + pool = MCPSessionPool() + pool.ensure_binding("A", "same-fp") + pool.reconcile_bindings({}, ("A",)) + + # remove -> identical re-add must not re-authorize from fingerprint equality + with pytest.raises(StaleMCPBindingError): + pool.ensure_binding("A", "same-fp") + + +def test_ensure_binding_rejects_retired_pool(): + pool = MCPSessionPool() + pool.retire_all() + + with pytest.raises(StaleMCPBindingError): + pool.ensure_binding("A", "fp-a") + + +@pytest.mark.asyncio +async def test_reconcile_retains_detached_owners_when_retirement_is_dropped(): + """Detached owners keep a tracked teardown even without PreparedRetirement.""" + pool = MCPSessionPool() + binding = pool.ensure_binding("A", "a1") + gate = asyncio.Event() + cm = _GatedSessionCm(gate) + + with patch("langchain_mcp_adapters.sessions.create_session", return_value=cm): + creator = asyncio.create_task(pool.get_session("A", "u:t", _CONNECTION, binding=binding)) + await asyncio.wait_for(cm.initialize_started.wait(), 1) + + # Return value deliberately dropped: the pool itself must retain the + # detached owner's teardown instead of leaving a dangling task. + pool.reconcile_bindings({"A": "a2"}, ()) + reapers = set(pool._teardown_tasks) + assert reapers, "detached owner teardown was not retained" + + gate.set() + with pytest.raises(StaleMCPBindingError): + await asyncio.wait_for(creator, 1) + + await asyncio.wait_for(asyncio.gather(*reapers, return_exceptions=True), 1) + deadline = asyncio.get_running_loop().time() + 1 + while pool._teardown_tasks and asyncio.get_running_loop().time() < deadline: + await asyncio.sleep(0) + + assert not pool._teardown_tasks + assert cm.closed is True + assert cm.exit_task is cm.enter_task + + +@pytest.mark.asyncio +async def test_joiner_rechecks_binding_after_ready_resolves(): + """A joiner cannot return a session superseded after ready resolved. + + The joiner path must be pinned to the post-await fence: the test waits until + the joiner has passed admission and is parked inside ``shield(ready)``, lets + the shared future resolve, and only then supersedes the binding. A joiner + that merely raced into admission would fail for the wrong reason. + """ + pool = MCPSessionPool() + old = pool.ensure_binding("A", "a1") + loop = asyncio.get_running_loop() + key = ("A", "u:t", loop, "deployment") + ready = loop.create_future() + close_evt = asyncio.Event() + + async def owner_waiter(): + await close_evt.wait() + + owner = asyncio.create_task(owner_waiter()) + pool._inflight[key] = (loop, ready, owner, close_evt) + + real_shield = asyncio.shield + joined = asyncio.Event() + resolved = asyncio.Event() + release = asyncio.Event() + + async def gated_shield(awaitable): + if awaitable is not ready: + return await real_shield(awaitable) + # Past admission, parked on the shared future: report progress so the + # test can sequence the supersede strictly after the result is consumed. + joined.set() + result = await real_shield(awaitable) + resolved.set() + await release.wait() + return result + + with patch("deerflow.mcp.session_pool.asyncio.shield", new=gated_shield): + joiner = asyncio.create_task(pool.get_session("A", "u:t", _CONNECTION, binding=old)) + + # Proves the joiner cleared admission, saw the in-flight entry and + # captured the join future instead of failing on the admission fence. + await asyncio.wait_for(joined.wait(), 1) + + ready.set_result(MagicMock()) + # Proves shield already consumed the result while still parked, so the + # supersede below lands before the post-await fence is evaluated. + await asyncio.wait_for(resolved.wait(), 1) + + pool.reconcile_bindings({"A": "a2"}, ()) + release.set() + + with pytest.raises(StaleMCPBindingError): + await asyncio.wait_for(joiner, 1) + + await asyncio.gather(owner, return_exceptions=True) diff --git a/backend/tests/test_mcp_cache.py b/backend/tests/test_mcp_cache.py index 60b9d9c6f9..28969fc4c1 100644 --- a/backend/tests/test_mcp_cache.py +++ b/backend/tests/test_mcp_cache.py @@ -507,6 +507,9 @@ def test_config_change_during_initialization_retires_pool_for_same_server_connec self.sessions[key] = FakeSession(connection["command"]) return self.sessions[key] + def retire_all(self) -> None: + pass + def close_all_sync(self) -> None: self.closed = True @@ -554,6 +557,85 @@ def test_config_change_during_initialization_retires_pool_for_same_server_connec real_reset_session_pool() +def test_superseded_initializer_cannot_seed_replacement_pool(cache_globals, monkeypatch, tmp_path): + """A superseded initializer must not install its binding into the new pool. + + Generation invalidation already discards the stale result, but + ``ensure_binding()`` is a persistent side effect on the session pool. If the + stale initializer re-resolves the singleton *after* the reset, it seeds the + replacement pool with the fingerprint it read from the superseded config; + the successor initializer then fails closed on its own clean pool and the + cache can never publish again until an explicit reset. + """ + from unittest.mock import AsyncMock + + from deerflow.mcp import session_pool as session_pool_module + from deerflow.mcp import tools as tools_module + + cfg = tmp_path / "extensions_config.json" + _write_extensions_config(cfg, {"A": _server("old-cmd")}) + monkeypatch.setenv("DEER_FLOW_EXTENSIONS_CONFIG_PATH", str(cfg)) + + real_get_mcp_tools = tools_module.get_mcp_tools + + class FakeClient: + def __init__(self, connections, **kwargs) -> None: + self.connections = connections + self.callbacks = kwargs.get("callbacks") or [] + self.tool_interceptors = kwargs.get("tool_interceptors") or [] + self.tool_name_prefix = kwargs.get("tool_name_prefix", True) + + async def get_tools(self, *, server_name=None): + return [] + + async def _scenario() -> None: + session_pool_module.reset_session_pool() + old_pool = session_pool_module.get_session_pool() + + entered = asyncio.Event() + release = asyncio.Event() + + async def _gated_get_mcp_tools(*args, **kwargs): + # Gate after the cache prologue so the test can supersede the claim + # while discovery still holds the config revision it read for G0. + entered.set() + await release.wait() + return await real_get_mcp_tools(*args, **kwargs) + + monkeypatch.setattr(tools_module, "get_mcp_tools", _gated_get_mcp_tools) + monkeypatch.setattr("deerflow.mcp.tools.get_initial_oauth_headers", AsyncMock(return_value={})) + monkeypatch.setattr("deerflow.mcp.tools.build_mcp_tool_interceptors", lambda *a, **k: []) + monkeypatch.setattr("langchain_mcp_adapters.client.MultiServerMCPClient", FakeClient) + + stale = asyncio.create_task(cache_module.initialize_mcp_tools()) + await asyncio.wait_for(entered.wait(), 1) + + _write_extensions_config(cfg, {"A": _server("new-cmd")}) + cache_module.reset_mcp_tools_cache() + assert old_pool._retired is True + assert session_pool_module._pool is None + + release.set() + assert await asyncio.wait_for(stale, 1) == [] + assert cache_module._cache_initialized is False + + # The successor must initialize against a pool the stale claim never + # touched; poisoning it here previously raised StaleMCPBindingError and + # left the cache permanently uninitialized. + assert await asyncio.wait_for(cache_module.initialize_mcp_tools(), 5) == [] + assert cache_module._cache_initialized is True + + new_pool = session_pool_module.get_session_pool() + assert new_pool is not old_pool + assert old_pool._bindings == {} + assert ("deployment", "A") in new_pool._bindings + + try: + asyncio.run(_scenario()) + finally: + session_pool_module.reset_session_pool() + + def test_reset_mcp_tools_cache_does_not_wait_for_in_flight_initialization(cache_globals, monkeypatch, tmp_path): """Event-loop callers can reset cache state without waiting for a slow tool load.""" cfg = tmp_path / "extensions_config.json" @@ -1072,10 +1154,16 @@ def test_initialization_hands_the_snapshotted_config_to_discovery(cache_globals, _write_extensions_config(cfg, {"srv1": _server("npx")}) monkeypatch.setenv("DEER_FLOW_EXTENSIONS_CONFIG_PATH", str(cfg)) - seen: dict[str, str] = {} + from deerflow.mcp import session_pool as session_pool_module - async def _fake_tools(*, extensions_config): + seen: dict[str, object] = {} + claimed_pool = session_pool_module.get_session_pool() + + async def _fake_tools(*, extensions_config, session_pool=None): seen["command"] = extensions_config.mcp_servers["srv1"].command + # Discovery must run against the pool this generation captured, never a + # singleton re-resolved after the claim. + seen["pool"] = session_pool # A transient flip to Y and back to X must not invalidate X-built tools. _write_extensions_config(cfg, {"srv1": _server("uvx")}) _write_extensions_config(cfg, {"srv1": _server("npx")}) @@ -1086,6 +1174,7 @@ def test_initialization_hands_the_snapshotted_config_to_discovery(cache_globals, result = asyncio.run(cache_module.initialize_mcp_tools()) assert seen["command"] == "npx" + assert seen["pool"] is claimed_pool assert result == ["loaded-tools"] assert cache_module._cache_initialized is True assert cache_module._mcp_config_snapshot is not None