From 683d146a30411a743bf314866678f1bdae3944e6 Mon Sep 17 00:00:00 2001 From: goloisme Date: Fri, 4 Sep 2026 11:30:36 +0800 Subject: [PATCH] fix(mcp): MCP cache re-initialization broken by cross-loop asyncio.Lock (#5062) * Fixes #5060: P1 snapshot config before loading, P2 RLock for sync path P1: Config changes during initialization can permanently cache stale tools. - Snapshot _config_path and _config_signature BEFORE await get_mcp_tools() - Compare AFTER get_mcp_tools() completes using _current_config_state() - If config changed during loading, discard stale result and retry - Prevents publishing old tools with new signature, which would make _is_cache_stale() permanently return False P2: Module-level asyncio.Lock still fails across event loops after real contention. - _init_lock = threading.RLock() for sync path (reentrant, prevents races) - _async_init_lock = asyncio.Lock() for async init serialization - reset_mcp_tools_cache() now acquires _init_lock for serialization Also fixes test P3: removed duplicated test bodies that leaked state between tests. * fix(mcp): make cache initialization cross-loop safe - Replace the module-level asyncio.Lock with thread-safe generation claiming - Snapshot config state before/after MCP loading and discard stale results - Keep reset state changes short and non-blocking for async endpoints - Add regression coverage for contended cross-loop init, config rewrites during load, and reset while init is in flight * fix: release MCP init claim on cancellation Release the in-flight generation claim from a cancellation-safe finally block so cancelling the task that owns initialization does not strand future callers. Add regression coverage for cancelling the owner and then reinitializing successfully. * fix(mcp): retire session pool before cache reset release Prevent a concurrent MCP cache initializer from publishing tool wrappers bound to the session-pool singleton that reset_mcp_tools_cache() is already retiring. Add regression coverage for that interleaving. * fix(mcp): retire session pool on stale cache invalidation * fix(mcp): retire pool on init discard --------- Co-authored-by: Willem Jiang --- .../packages/harness/deerflow/mcp/AGENTS.md | 2 +- .../packages/harness/deerflow/mcp/cache.py | 157 +++++-- .../harness/deerflow/mcp/session_pool.py | 6 +- backend/tests/test_mcp_cache.py | 412 +++++++++++++++++- 4 files changed, 536 insertions(+), 41 deletions(-) diff --git a/backend/packages/harness/deerflow/mcp/AGENTS.md b/backend/packages/harness/deerflow/mcp/AGENTS.md index 2e8fd793fd..b458b29aa6 100644 --- a/backend/packages/harness/deerflow/mcp/AGENTS.md +++ b/backend/packages/harness/deerflow/mcp/AGENTS.md @@ -6,7 +6,7 @@ - **Long-running ordinary task driver**: `extensions_config.json -> mcpServers..task_toolsets` binds exact raw submit/status/cancel names; one raw tool may occupy only one role across that server's groups. `mcp/tools.py` hides status/cancel and replaces submit with a wrapper that returns only the local task ID after persistence. `ordinary.py` reads only MCP `structuredContent`, maps remote `running` to `working`, and treats `error_code=task_not_found` or malformed structured output as permanent failure. A status call with `isError=true` is a retryable call failure: the first text content block is retained as a bounded diagnostic, while a permanent remote-task outcome must arrive in a normal result with structured `status=failed`. `task_tool_caller.py` restores the same `(server_name, user_id:thread_id)` stdio session scope; HTTP/SSE calls remain ephemeral, apply `session_init_timeout` to initialization and `tool_call_timeout` to task calls, and support server-level OAuth refresh outside an Agent run. `McpTaskService` exponentially backs off transient status/cancel errors without a maximum attempt count, derives API `tracking_degraded` from the consecutive-error threshold, keeps `input_required` on a slower poll, and caps finite positive remote poll hints at 24 hours. Task-enabled server runtime/binding configuration and `mcpInterceptors` are frozen to the Gateway startup snapshot; hot drift fails clearly before tool discovery can diverge from background calls, while presentation-only fields and non-task servers remain reloadable. Configured task toolsets fail startup when the runtime is disabled or persistence is memory. Users still cannot submit an answer back to an `input_required` remote task. - **Durable task payload bounds**: persisted task errors are capped at 4,000 characters. `input_required` and `result_artifact` must each serialize as valid JSON within 64 KiB; an invalid or oversized payload becomes a permanent protocol failure rather than being truncated and changing its semantics. Remote task IDs/task names are limited to 255 characters and task-enabled server names to 128, matching the SQL schema; an oversized submitted remote ID is rejected only after the Service has the handle so compensation cancellation still runs. Oversized results retain the existing bounded preview/truncation/artifact behavior. - **Lazy initialization**: Tools loaded on first use via `get_cached_mcp_tools()` -- **Cache invalidation**: Detects extensions-config changes by comparing the resolved config path and a `(mtime, size, sha256)` content signature against the values recorded at initialization, 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. `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 +- **Cache invalidation**: Detects extensions-config changes by comparing the resolved config path and a `(mtime, size, sha256)` content signature against the values recorded at initialization, 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. `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. If `initialize_mcp_tools()` itself observes a config-signature change between the pre-load and post-load snapshots, that discard branch must reset the tool cache through the same session-pool retirement path as normal stale invalidation before waiters retry; otherwise a stale load can leave `(server_name, scope_key)` sessions from the abandoned connection available to the next wrapper build. - **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. - **OAuth (HTTP/SSE)**: Supports token endpoint flows (`client_credentials`, `refresh_token`) with automatic token refresh + Authorization header injection. The rendered ` ` passes `mcp/headers.py::illegal_header_value_reason` inside `OAuthTokenManager.get_authorization_header` — the one boundary the tool interceptor, the initial discovery headers and the durable task path all read their value from — so a token endpoint returning something the transport would refuse fails closed instead of letting h11 echo the token into a model-visible tool error. The rendered value is what gets checked, not the two fields separately, because that is what the transport sees. The operator's static `headers` get the same check in `mcp/client.py::build_server_params`, where `build_servers_config` already drops just that server and logs the reason. diff --git a/backend/packages/harness/deerflow/mcp/cache.py b/backend/packages/harness/deerflow/mcp/cache.py index cff7d7d722..b02b385b7b 100644 --- a/backend/packages/harness/deerflow/mcp/cache.py +++ b/backend/packages/harness/deerflow/mcp/cache.py @@ -2,6 +2,7 @@ import asyncio import logging +import threading from pathlib import Path from langchain_core.tools import BaseTool @@ -13,7 +14,10 @@ logger = logging.getLogger(__name__) _mcp_tools_cache: list[BaseTool] | None = None _cache_initialized = False -_initialization_lock = asyncio.Lock() +_init_lock = threading.RLock() # Guards cache state transitions. +_init_condition = threading.Condition(_init_lock) +_initializing_generation: int | None = None +_cache_generation = 0 # Cache-invalidation key for the resolved extensions config file. We track the # resolved path *and* a ``(mtime, size, sha256)`` content signature — via the @@ -113,6 +117,12 @@ def _is_cache_stale() -> bool: return False +def _wait_for_initialization(generation: int | None) -> None: + """Wait for an in-flight initialization without binding to any event loop.""" + with _init_condition: + _init_condition.wait_for(lambda: _cache_initialized or _initializing_generation != generation) + + async def initialize_mcp_tools() -> list[BaseTool]: """Initialize and cache MCP tools. @@ -122,21 +132,66 @@ async def initialize_mcp_tools() -> list[BaseTool]: List of LangChain tools from all enabled MCP servers. """ global _mcp_tools_cache, _cache_initialized, _config_path, _config_signature + global _initializing_generation, _cache_generation - async with _initialization_lock: - if _cache_initialized: - logger.info("MCP tools already initialized") - return _mcp_tools_cache or [] + while True: + with _init_condition: + if _cache_initialized: + logger.info("MCP tools already initialized") + return _mcp_tools_cache or [] - from deerflow.mcp.tools import get_mcp_tools + if _initializing_generation is None: + pre_path, pre_sig = _current_config_state() + claim_generation = _cache_generation + _initializing_generation = claim_generation + break + waiting_generation = _initializing_generation + + await asyncio.to_thread(_wait_for_initialization, waiting_generation) + + from deerflow.mcp.tools import get_mcp_tools + + loaded_tools = None + post_path = None + post_sig = None + init_succeeded = False + try: logger.info("Initializing MCP tools...") - _mcp_tools_cache = await get_mcp_tools() - _cache_initialized = True - _config_path, _config_signature = _current_config_state() # Record config path + content signature - logger.info("MCP tools initialized: %d tool(s) loaded (config path: %s)", len(_mcp_tools_cache), _config_path) + loaded_tools = await get_mcp_tools() + post_path, post_sig = _current_config_state() + init_succeeded = True + finally: + if not init_succeeded: + with _init_condition: + if _initializing_generation == claim_generation: + _initializing_generation = None + _init_condition.notify_all() - return _mcp_tools_cache + retired_pool = None + with _init_condition: + try: + if _cache_generation != claim_generation: + logger.info("MCP cache was reset during initialization; discarding stale result") + return [] + + if (pre_path, pre_sig) != (post_path, post_sig): + logger.warning("MCP config changed during initialization; discarding stale result") + retired_pool = _reset_mcp_tools_cache_state_and_retire_pool_locked() + else: + _mcp_tools_cache = loaded_tools + _cache_initialized = True + _config_path, _config_signature = post_path, post_sig + logger.info("MCP tools initialized: %d tool(s) loaded (config path: %s)", len(_mcp_tools_cache), _config_path) + return _mcp_tools_cache + finally: + if _initializing_generation == claim_generation: + _initializing_generation = None + _init_condition.notify_all() + + if retired_pool is not None: + retired_pool.close_all_sync() + return [] def get_cached_mcp_tools() -> list[BaseTool]: @@ -152,31 +207,35 @@ def get_cached_mcp_tools() -> list[BaseTool]: Returns: List of cached MCP tools. """ - global _cache_initialized + while True: + retired_pool = None + with _init_lock: + if _is_cache_stale(): + logger.info("MCP cache is stale, resetting for re-initialization...") + retired_pool = _reset_mcp_tools_cache_state_and_retire_pool_locked() - # Check if cache is stale due to config file changes - if _is_cache_stale(): - logger.info("MCP cache is stale, resetting for re-initialization...") - reset_mcp_tools_cache() + if _cache_initialized: + return _mcp_tools_cache or [] + + if _initializing_generation is not None: + _init_condition.wait_for(lambda: _initializing_generation is None or _cache_initialized) + continue + + if retired_pool is not None: + retired_pool.close_all_sync() - if not _cache_initialized: logger.info("MCP tools not initialized, performing lazy initialization...") try: - # Try to initialize in the current event loop loop = asyncio.get_event_loop() if loop.is_running(): - # If loop is already running (e.g., in LangGraph Studio), - # we need to create a new loop in a thread import concurrent.futures with concurrent.futures.ThreadPoolExecutor() as executor: future = executor.submit(asyncio.run, initialize_mcp_tools()) future.result() else: - # If no loop is running, we can use the current loop loop.run_until_complete(initialize_mcp_tools()) except RuntimeError: - # No event loop exists, create one try: asyncio.run(initialize_mcp_tools()) except Exception: @@ -186,7 +245,37 @@ def get_cached_mcp_tools() -> list[BaseTool]: logger.exception("Failed to lazy-initialize MCP tools") return [] - return _mcp_tools_cache or [] + with _init_lock: + if _cache_initialized: + return _mcp_tools_cache or [] + + +def _reset_mcp_tools_cache_state() -> None: + """Reset cache state under ``_init_condition`` / ``_init_lock``.""" + global _mcp_tools_cache, _cache_initialized, _config_path, _config_signature + global _cache_generation + + _mcp_tools_cache = None + _cache_initialized = False + _config_path = None + _config_signature = None + _cache_generation += 1 + _init_condition.notify_all() + + +def _reset_mcp_tools_cache_state_and_retire_pool_locked(): + """Retire the MCP session pool and reset cache state under one lock. + + Tool wrappers close over the module-level session-pool singleton when they + are built. Any path that invalidates the tool cache must therefore swap the + singleton before waiters/fresh initializers can rebuild wrappers, including + automatic config-signature invalidation in ``get_cached_mcp_tools()``. + """ + from deerflow.mcp.session_pool import reset_session_pool + + retired_pool = reset_session_pool() + _reset_mcp_tools_cache_state() + return retired_pool def reset_mcp_tools_cache() -> None: @@ -196,12 +285,6 @@ def reset_mcp_tools_cache() -> None: Also closes all persistent MCP sessions so they are recreated on the next tool load. """ - global _mcp_tools_cache, _cache_initialized, _config_path, _config_signature - _mcp_tools_cache = None - _cache_initialized = False - _config_path = None - _config_signature = None - # Close persistent sessions – they will be recreated by the next # get_mcp_tools() call with the (possibly updated) connection config. # @@ -215,13 +298,19 @@ def reset_mcp_tools_cache() -> None: # loop to finish teardown here: that is a self-deadlock (the loop can only # run the teardown after this synchronous call returns control to it). try: - from deerflow.mcp.session_pool import get_session_pool + from deerflow.mcp.session_pool import reset_session_pool - get_session_pool().close_all_sync() + with _init_condition: + # Retire the session-pool singleton before cache waiters can start a + # fresh initialization. Otherwise a concurrent initializer can build + # tool wrappers against the soon-to-be-detached pool and publish + # them after this reset replaces the singleton. + retired_pool = reset_session_pool() + _reset_mcp_tools_cache_state() + + if retired_pool is not None: + retired_pool.close_all_sync() except Exception: logger.debug("Could not close MCP session pool on cache reset", exc_info=True) - from deerflow.mcp.session_pool import reset_session_pool - - reset_session_pool() logger.info("MCP tools cache reset") diff --git a/backend/packages/harness/deerflow/mcp/session_pool.py b/backend/packages/harness/deerflow/mcp/session_pool.py index f03dc7aee9..fceb065d40 100644 --- a/backend/packages/harness/deerflow/mcp/session_pool.py +++ b/backend/packages/harness/deerflow/mcp/session_pool.py @@ -738,8 +738,10 @@ def get_session_pool() -> MCPSessionPool: return _pool -def reset_session_pool() -> None: - """Reset the singleton (used in tests and the MCP cache reset path).""" +def reset_session_pool() -> MCPSessionPool | None: + """Reset the singleton and return the retired pool, if any.""" global _pool with _pool_lock: + retired = _pool _pool = None + return retired diff --git a/backend/tests/test_mcp_cache.py b/backend/tests/test_mcp_cache.py index f145e1adb5..63d25080f1 100644 --- a/backend/tests/test_mcp_cache.py +++ b/backend/tests/test_mcp_cache.py @@ -23,6 +23,7 @@ from __future__ import annotations import asyncio import json import os +import threading from pathlib import Path import pytest @@ -42,7 +43,10 @@ _TRACKED_GLOBALS = ( "_config_path", "_config_signature", "_config_mtime", - "_initialization_lock", + "_init_lock", + "_init_condition", + "_initializing_generation", + "_cache_generation", ) @@ -64,9 +68,12 @@ def cache_globals(): for name in ("_config_path", "_config_signature", "_config_mtime"): if hasattr(cache_module, name): setattr(cache_module, name, None) - # asyncio.Lock binds to the first event loop it is awaited on, so each test - # (which drives initialize_mcp_tools via a fresh asyncio.run) needs its own. - cache_module._initialization_lock = asyncio.Lock() + # threading.Lock is safe across threads and does not bind to event loops, + # so each test gets fresh coordination state for isolation. + cache_module._init_lock = threading.RLock() + cache_module._init_condition = threading.Condition(cache_module._init_lock) + cache_module._initializing_generation = None + cache_module._cache_generation = 0 try: yield @@ -286,3 +293,400 @@ def test_config_deleted_after_init_via_real_env_resolution_does_not_raise(cache_ # last-known-good MCP tools), matching the deliberate contract in # test_config_deleted_after_init_is_not_stale above. assert cache_module._is_cache_stale() is False + + +class TestCrossLoopReinitialization: + """Regression for #5060: cache initialization must be cross-loop safe.""" + + def test_cross_loop_reinit_does_not_raise(self, monkeypatch, tmp_path): + """Two successive ``asyncio.run()`` calls (each with its own loop) must not crash.""" + saved = {name: getattr(cache_module, name, _MISSING) for name in _TRACKED_GLOBALS} + try: + cache_module._mcp_tools_cache = None + cache_module._cache_initialized = False + for name in ("_config_path", "_config_signature", "_config_mtime"): + if hasattr(cache_module, name): + setattr(cache_module, name, None) + cache_module._init_lock = threading.RLock() + cache_module._init_condition = threading.Condition(cache_module._init_lock) + cache_module._initializing_generation = None + cache_module._cache_generation = 0 + + cfg = tmp_path / "extensions_config.json" + _write_extensions_config(cfg, {"srv1": _server()}) + monkeypatch.setenv("DEER_FLOW_EXTENSIONS_CONFIG_PATH", str(cfg)) + + async def _fake_tools(): + return [] + + monkeypatch.setattr("deerflow.mcp.tools.get_mcp_tools", _fake_tools) + + asyncio.run(cache_module.initialize_mcp_tools()) + assert cache_module._cache_initialized is True + + cache_module.reset_mcp_tools_cache() + assert cache_module._cache_initialized is False + + asyncio.run(cache_module.initialize_mcp_tools()) + assert cache_module._cache_initialized is True + finally: + for name, value in saved.items(): + if value is _MISSING: + if hasattr(cache_module, name): + delattr(cache_module, name) + else: + setattr(cache_module, name, value) + + def test_contended_cross_loop_reinit_does_not_raise(self, monkeypatch, tmp_path): + """Contended initializers in two event loops must not reuse a loop-bound lock.""" + saved = {name: getattr(cache_module, name, _MISSING) for name in _TRACKED_GLOBALS} + try: + cache_module._mcp_tools_cache = None + cache_module._cache_initialized = False + for name in ("_config_path", "_config_signature", "_config_mtime"): + if hasattr(cache_module, name): + setattr(cache_module, name, None) + cache_module._init_lock = threading.RLock() + cache_module._init_condition = threading.Condition(cache_module._init_lock) + cache_module._initializing_generation = None + cache_module._cache_generation = 0 + + cfg = tmp_path / "extensions_config.json" + _write_extensions_config(cfg, {"srv1": _server()}) + monkeypatch.setenv("DEER_FLOW_EXTENSIONS_CONFIG_PATH", str(cfg)) + + calls = 0 + + async def _fake_tools(): + nonlocal calls + calls += 1 + await asyncio.sleep(0.01) + return [] + + monkeypatch.setattr("deerflow.mcp.tools.get_mcp_tools", _fake_tools) + + async def _contended_init(): + await asyncio.gather(cache_module.initialize_mcp_tools(), cache_module.initialize_mcp_tools()) + + asyncio.run(_contended_init()) + assert cache_module._cache_initialized is True + assert calls == 1 + + cache_module.reset_mcp_tools_cache() + assert cache_module._cache_initialized is False + + asyncio.run(_contended_init()) + assert cache_module._cache_initialized is True + assert calls == 2 + finally: + for name, value in saved.items(): + if value is _MISSING: + if hasattr(cache_module, name): + delattr(cache_module, name) + else: + setattr(cache_module, name, value) + + def test_get_cached_mcp_tools_reinit_after_invalidation(self, monkeypatch, tmp_path): + """Gateway path without ``cache_globals`` resetting coordination state — production scenario.""" + saved = {name: getattr(cache_module, name, _MISSING) for name in _TRACKED_GLOBALS} + try: + cache_module._mcp_tools_cache = None + cache_module._cache_initialized = False + for name in ("_config_path", "_config_signature", "_config_mtime"): + if hasattr(cache_module, name): + setattr(cache_module, name, None) + cache_module._init_lock = threading.RLock() + cache_module._init_condition = threading.Condition(cache_module._init_lock) + cache_module._initializing_generation = None + cache_module._cache_generation = 0 + + cfg = tmp_path / "extensions_config.json" + _write_extensions_config(cfg, {"srv1": _server()}) + monkeypatch.setenv("DEER_FLOW_EXTENSIONS_CONFIG_PATH", str(cfg)) + + async def _fake_tools(): + return [] + + monkeypatch.setattr("deerflow.mcp.tools.get_mcp_tools", _fake_tools) + + result1 = cache_module.get_cached_mcp_tools() + assert result1 == [] + assert cache_module._cache_initialized is True + + _write_extensions_config(cfg, {"srv1": _server(), "srv2": _server("uvx")}) + cache_module.reset_mcp_tools_cache() + + result2 = cache_module.get_cached_mcp_tools() + assert result2 == [] + assert cache_module._cache_initialized is True + finally: + for name, value in saved.items(): + if value is _MISSING: + if hasattr(cache_module, name): + delattr(cache_module, name) + else: + setattr(cache_module, name, value) + + +def test_config_change_during_initialization_discards_stale_tools(cache_globals, monkeypatch, tmp_path): + """A config rewrite during load must not publish tools loaded from the old config.""" + cfg = tmp_path / "extensions_config.json" + _write_extensions_config(cfg, {"old": _server()}) + monkeypatch.setenv("DEER_FLOW_EXTENSIONS_CONFIG_PATH", str(cfg)) + + calls = 0 + + async def _fake_tools(): + nonlocal calls + calls += 1 + if calls == 1: + _write_extensions_config(cfg, {"new": _server("uvx")}) + return ["old-tools"] + return ["new-tools"] + + monkeypatch.setattr("deerflow.mcp.tools.get_mcp_tools", _fake_tools) + + first = asyncio.run(cache_module.initialize_mcp_tools()) + assert first == [] + assert cache_module._cache_initialized is False + assert cache_module._mcp_tools_cache is None + + second = cache_module.get_cached_mcp_tools() + assert second == ["new-tools"] + assert cache_module._cache_initialized is True + assert cache_module._mcp_tools_cache == ["new-tools"] + assert cache_module._is_cache_stale() is False + assert calls == 2 + + +def test_config_change_during_initialization_retires_pool_for_same_server_connection_change(cache_globals, monkeypatch, tmp_path): + """Discarding a mid-load config change must also retire pooled sessions. + + A stale load can create a pooled session before ``initialize_mcp_tools`` + notices that the config changed and discards the loaded tools. If the server + name and scope stay the same while the connection changes, the next load can + otherwise reuse that old session because ``MCPSessionPool`` keys only by + ``(server_name, scope_key)``. + """ + from deerflow.mcp import session_pool as session_pool_module + + class FakeSession: + def __init__(self, command: str) -> None: + self.command = command + + class FakeSessionPool: + def __init__(self) -> None: + self.closed = False + self.sessions = {} + + async def get_session(self, server_name, scope_key, connection): + key = (server_name, scope_key) + if key not in self.sessions: + self.sessions[key] = FakeSession(connection["command"]) + return self.sessions[key] + + def close_all_sync(self) -> None: + self.closed = True + + real_reset_session_pool = session_pool_module.reset_session_pool + monkeypatch.setattr(session_pool_module, "MCPSessionPool", FakeSessionPool) + real_reset_session_pool() + old_pool = session_pool_module.get_session_pool() + + cfg = tmp_path / "extensions_config.json" + _write_extensions_config(cfg, {"same": _server("npx")}) + monkeypatch.setenv("DEER_FLOW_EXTENSIONS_CONFIG_PATH", str(cfg)) + + calls = 0 + loaded_pools = [] + loaded_sessions = [] + + async def _fake_tools(): + nonlocal calls + calls += 1 + server = json.loads(cfg.read_text())["mcpServers"]["same"] + pool = session_pool_module.get_session_pool() + session = await pool.get_session("same", "thread-1", server) + loaded_pools.append(pool) + loaded_sessions.append(session) + if calls == 1: + _write_extensions_config(cfg, {"same": _server("uvx")}) + return [f"session-{session.command}"] + + monkeypatch.setattr("deerflow.mcp.tools.get_mcp_tools", _fake_tools) + + try: + first = asyncio.run(cache_module.initialize_mcp_tools()) + assert first == [] + assert cache_module._cache_initialized is False + + second = cache_module.get_cached_mcp_tools() + + assert second == ["session-uvx"] + assert loaded_pools[0] is old_pool + assert loaded_pools[1] is not old_pool + assert loaded_sessions[0] is not loaded_sessions[1] + assert old_pool.closed is True + assert cache_module._cache_initialized is True + finally: + real_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" + _write_extensions_config(cfg, {"srv1": _server()}) + monkeypatch.setenv("DEER_FLOW_EXTENSIONS_CONFIG_PATH", str(cfg)) + + started = threading.Event() + finish = threading.Event() + + async def _fake_tools(): + started.set() + await asyncio.to_thread(finish.wait) + return [] + + monkeypatch.setattr("deerflow.mcp.tools.get_mcp_tools", _fake_tools) + + worker = threading.Thread(target=lambda: asyncio.run(cache_module.initialize_mcp_tools())) + worker.start() + assert started.wait(timeout=1) + + reset_done = threading.Event() + + def _reset(): + cache_module.reset_mcp_tools_cache() + reset_done.set() + + reset_thread = threading.Thread(target=_reset) + reset_thread.start() + assert reset_done.wait(timeout=1) + + finish.set() + worker.join(timeout=1) + reset_thread.join(timeout=1) + assert not worker.is_alive() + assert not reset_thread.is_alive() + + +def test_automatic_stale_invalidation_retires_session_pool_before_reinitializing(cache_globals, monkeypatch, tmp_path): + """Automatic config-signature invalidation must retire the old session pool. + + ``get_cached_mcp_tools()`` detects runtime edits through ``_is_cache_stale`` + without going through the explicit admin reset endpoint. That automatic path + must still swap the session-pool singleton before rebuilding tool wrappers; + otherwise the fresh wrappers can keep reusing sessions created from the old + connection config. + """ + from deerflow.mcp import session_pool as session_pool_module + + real_reset_session_pool = session_pool_module.reset_session_pool + real_reset_session_pool() + old_pool = session_pool_module.get_session_pool() + + cfg = tmp_path / "extensions_config.json" + _write_extensions_config(cfg, {"old": _server()}) + _initialize_against(monkeypatch, cfg) + assert cache_module._cache_initialized is True + assert session_pool_module.get_session_pool() is old_pool + + _write_extensions_config(cfg, {"new": _server("uvx")}) + loaded_pools = [] + + async def _fake_tools(): + loaded_pools.append(session_pool_module.get_session_pool()) + return ["new-tools"] + + monkeypatch.setattr("deerflow.mcp.tools.get_mcp_tools", _fake_tools) + + try: + result = cache_module.get_cached_mcp_tools() + + assert result == ["new-tools"] + assert loaded_pools == [session_pool_module.get_session_pool()] + assert loaded_pools[0] is not old_pool + assert cache_module._cache_initialized is True + finally: + real_reset_session_pool() + + +def test_reset_mcp_tools_cache_retires_session_pool_before_releasing_initializers(cache_globals, monkeypatch): + """A reset must not let fresh tool wrappers publish with the retiring pool.""" + from deerflow.mcp import session_pool as session_pool_module + + real_reset_session_pool = session_pool_module.reset_session_pool + real_reset_session_pool() + old_pool = session_pool_module.get_session_pool() + + cache_module._mcp_tools_cache = ["cached-tools"] + cache_module._cache_initialized = True + race_results = [] + loaded_pools = [] + + async def _fake_tools(): + loaded_pools.append(session_pool_module.get_session_pool()) + return ["race-tools"] + + def _reset_with_concurrent_initializer(): + # This simulates the old interleaving: a cache waiter starts exactly + # while reset_mcp_tools_cache() is retiring the session-pool singleton. + race_results.append(asyncio.run(cache_module.initialize_mcp_tools())) + return real_reset_session_pool() + + monkeypatch.setattr("deerflow.mcp.tools.get_mcp_tools", _fake_tools) + monkeypatch.setattr(session_pool_module, "reset_session_pool", _reset_with_concurrent_initializer) + + try: + cache_module.reset_mcp_tools_cache() + + # The racing initializer should see the still-valid old cache and avoid + # rebuilding tools until the pool has been swapped. Pre-fix, cache state + # was cleared first, so this race loaded and published wrappers against + # ``old_pool`` just before the singleton was replaced. + assert race_results == [["cached-tools"]] + assert loaded_pools == [] + assert cache_module._cache_initialized is False + assert cache_module._mcp_tools_cache is None + assert session_pool_module.get_session_pool() is not old_pool + finally: + real_reset_session_pool() + + +def test_cancelled_initializer_releases_generation_claim(cache_globals, monkeypatch, tmp_path): + """Cancelling the owner task must not strand waiters on its generation claim.""" + cfg = tmp_path / "extensions_config.json" + _write_extensions_config(cfg, {"srv1": _server()}) + monkeypatch.setenv("DEER_FLOW_EXTENSIONS_CONFIG_PATH", str(cfg)) + + started = asyncio.Event() + release = asyncio.Event() + calls = 0 + + async def _fake_tools(): + nonlocal calls + calls += 1 + started.set() + await release.wait() + return [] + + monkeypatch.setattr("deerflow.mcp.tools.get_mcp_tools", _fake_tools) + + async def _cancel_and_retry(): + owner = asyncio.create_task(cache_module.initialize_mcp_tools()) + await asyncio.wait_for(started.wait(), timeout=1) + owner.cancel() + + with pytest.raises(asyncio.CancelledError): + await owner + + assert cache_module._initializing_generation is None + assert cache_module._cache_initialized is False + + release.set() + result = await asyncio.wait_for(cache_module.initialize_mcp_tools(), timeout=1) + assert result == [] + + asyncio.run(_cancel_and_retry()) + assert cache_module._cache_initialized is True + assert cache_module._initializing_generation is None + assert calls == 2