From dfe9a520b96b1d84ace604ce5bcb6202e4c163ef Mon Sep 17 00:00:00 2001 From: yang rui <142162413+ruirui6946@users.noreply.github.com> Date: Sun, 13 Sep 2026 17:13:45 +0800 Subject: [PATCH] fix(mcp): isolate pooled sessions by owning event loop (#5396) * fix(mcp): isolate pooled sessions by owning event loop * refactor(mcp): remove obsolete eviction cancellation plumbing --- README.md | 7 + backend/docs/MCP_SERVER.md | 14 ++ .../packages/harness/deerflow/mcp/AGENTS.md | 13 ++ .../harness/deerflow/mcp/session_pool.py | 100 ++++----- backend/tests/test_mcp_session_pool.py | 106 +++++----- .../test_mcp_session_pool_loop_isolation.py | 200 ++++++++++++++++++ 6 files changed, 326 insertions(+), 114 deletions(-) create mode 100644 backend/tests/test_mcp_session_pool_loop_isolation.py diff --git a/README.md b/README.md index 8df044a74..ff07c79ce 100644 --- a/README.md +++ b/README.md @@ -1645,6 +1645,13 @@ The response contains normalized `cron`, `timezone`, the effective UTC `start_at `deerflow` is a terminal-native workbench for people who live in the shell. It runs **embedded** over `DeerFlowClient` — no Gateway, frontend, nginx, or Docker required — while honoring the same `config.yaml`, checkpointer, skills, memory, MCP, and sandbox settings as the rest of DeerFlow. +Parallel synchronous stdio MCP calls use independent sessions on their own event +loops. They do not cancel each other's connections, but they do not share +server-side state; session reuse requires the same loop. See the +[MCP session notes](backend/docs/MCP_SERVER.md) for details. +Manually managed event loops must drain pending session owners before closing; +the normal `asyncio.run()` path does this automatically. + ![DeerFlow TUI](docs/tui/tui-preview.svg) ```bash diff --git a/backend/docs/MCP_SERVER.md b/backend/docs/MCP_SERVER.md index f2ecd58fb..31d452753 100644 --- a/backend/docs/MCP_SERVER.md +++ b/backend/docs/MCP_SERVER.md @@ -199,6 +199,20 @@ backward compatibility. Disable it only when every resulting tool name remains unique across the enabled servers. Stdio tools continue to use DeerFlow's persistent per-thread session pool regardless of this setting. +Session reuse also requires the same owning event loop. Parallel synchronous +tool calls from the embedded client use separate loops and separate stdio +sessions, so they can finish independently without cancelling a sibling's +connection. They do not share server-side state. The synchronous wrapper closes +its loop after each call; use the asynchronous path on a shared loop when +session continuity is required. Explicit pool cleanup covers all loops for the +selected server/thread scope. + +If you manage event loops manually, close the pool or cancel and await its owner +tasks before closing their loop. Calling `loop.close()` with pending owners +prevents transport teardown and completion callbacks. Abandoned live-registry +records can be removed by LRU eviction or explicit cleanup, but those operations +cannot finish transport cleanup on a loop that has already closed. + ## Server Timeouts Two independent settings bound stdio MCP servers and durable HTTP/SSE task diff --git a/backend/packages/harness/deerflow/mcp/AGENTS.md b/backend/packages/harness/deerflow/mcp/AGENTS.md index 21d0aedd5..2c32e7464 100644 --- a/backend/packages/harness/deerflow/mcp/AGENTS.md +++ b/backend/packages/harness/deerflow/mcp/AGENTS.md @@ -6,6 +6,19 @@ - **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()` +- **Loop-isolated stdio sessions**: Both the live registry and in-flight creations + are keyed by `(server_name, scope_key, owning_loop)`. Same-loop callers share + initialization and state; live sibling loops must never cancel or replace each + other's sessions merely because their server/scope matches. Explicit server, + scope, pair and global cleanup span all loops; disconnect eviction matches the + exact session identity. Owner completion drops only its own records, including + during `asyncio.run()` shutdown. Sync wrappers use fresh loops per invocation, + so parallel sync calls have independent subprocess/server state. + Manual loop owners must drain pending tasks before `loop.close()`. Closing a + loop with a pending session owner prevents its completion callback and resource + teardown; an abandoned live-registry record remains bounded by `MAX_SESSIONS` + 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. - **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 diff --git a/backend/packages/harness/deerflow/mcp/session_pool.py b/backend/packages/harness/deerflow/mcp/session_pool.py index 9b1e3e257..04aa1facf 100644 --- a/backend/packages/harness/deerflow/mcp/session_pool.py +++ b/backend/packages/harness/deerflow/mcp/session_pool.py @@ -5,8 +5,9 @@ each tool call creates a new MCP session. For stateful servers like Playwright, this means browser state (opened pages, filled forms) is lost between calls. This module provides a session pool that maintains persistent MCP sessions, -scoped by ``(server_name, scope_key)`` — typically scope_key is the thread_id — -so that consecutive tool calls share the same session and server-side state. +scoped by ``(server_name, scope_key, owning_loop)``. Consecutive calls on +the same loop share server-side state; independent loops use separate sessions. +The sync wrapper uses a fresh loop per call, so it does not preserve that state. Sessions are evicted in LRU order when the pool reaches capacity. Lifecycle model (owner task) @@ -120,7 +121,7 @@ async def call_pooled_session_tool( class MCPSessionPool: - """Manages persistent MCP sessions scoped by ``(server_name, scope_key)``.""" + """Manages persistent MCP sessions scoped by ``(server_name, scope_key, owning_loop)``.""" MAX_SESSIONS = 256 SESSION_CLOSE_TIMEOUT = 5.0 # seconds to wait when closing a session on a foreign loop @@ -128,7 +129,7 @@ class MCPSessionPool: def __init__(self) -> None: # Each entry: (session, owning_loop, owner_task, close_event). self._entries: OrderedDict[ - tuple[str, str], + tuple[str, str, asyncio.AbstractEventLoop], tuple[ ClientSession, asyncio.AbstractEventLoop, @@ -136,7 +137,7 @@ class MCPSessionPool: asyncio.Event, ], ] = OrderedDict() - # In-flight creations, keyed by (server, scope). Lets concurrent callers + # In-flight creations, keyed by (server, scope, owning_loop). 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 @@ -144,7 +145,7 @@ class MCPSessionPool: # ``_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[str, str, asyncio.AbstractEventLoop], tuple[ asyncio.AbstractEventLoop, asyncio.Future[ClientSession], @@ -166,9 +167,19 @@ class MCPSessionPool: # Session owner task # ------------------------------------------------------------------ + def _discard_owner(self, key: tuple[str, str, asyncio.AbstractEventLoop], owner: asyncio.Task[Any]) -> None: + """Retire only this owner, including after asyncio.run shuts its loop down.""" + with self._lock: + entry = self._entries.get(key) + if entry is not None and entry[2] is owner: + self._entries.pop(key) + inflight = self._inflight.get(key) + if inflight is not None and inflight[2] is owner: + self._inflight.pop(key) + async def _run_session( self, - key: tuple[str, str], + key: tuple[str, str, asyncio.AbstractEventLoop], connection: dict[str, Any], ready: asyncio.Future[ClientSession], close_evt: asyncio.Event, @@ -261,9 +272,8 @@ class MCPSessionPool: ) -> ClientSession: """Get or create a persistent MCP session. - If an existing session was created in a different (or closed) event - loop, it is evicted and replaced with a fresh one owned by a task on - the current loop. + 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. @@ -273,79 +283,57 @@ class MCPSessionPool: Returns: An initialized ``ClientSession``. """ - key = (server_name, scope_key) current_loop = asyncio.get_running_loop() + key = (server_name, scope_key, current_loop) # 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, 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]] = [] + # 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 key in self._entries: - session, loop, ent_task, ent_close = self._entries[key] - if loop is current_loop and not loop.is_closed(): + session, _loop, ent_task, _ent_close = self._entries[key] + if not ent_task.done(): self._entries.move_to_end(key) return session - # Session belongs to a different/closed event loop – evict it. self._entries.pop(key) - 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(): - # Another caller on this loop is already creating the session; - # wait for the same result instead of building a duplicate. + if inflight is not None: + # Only callers on this same loop share the creation future. join = inflight[1] else: - if inflight is not None: - # Stale in-flight creation owned by a different/closed loop. - # Drop the record and tear its owner down; because that owner - # may be blocked inside initialize() (where close_evt cannot - # 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, 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(key, connection, ready, close_evt)) 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, False, None)) + evicted.append((loop, ent_task, ent_close)) - # Phase 2: shut down evicted sessions/creations. Signal EVERY removed + # 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 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: + # 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) - if cancel: - self._cancel_owner(loop, ent_task, ent_ready) try: - for loop, ent_task, ent_close, cancel, ent_ready in evicted: + for loop, ent_task, ent_close 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. + 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 @@ -634,12 +622,13 @@ class MCPSessionPool: 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.""" - key = (server_name, scope_key) + """Close every session for this server/scope across all owning loops.""" with self._lock: - entry = self._entries.pop(key, None) - inflight = self._inflight.pop(key, None) - await self._close_owners([entry] if entry is not None else [], [inflight] if inflight is not None else []) + keys = [k for k in self._entries if k[:2] == (server_name, scope_key)] + entries = [self._entries.pop(k) for k in keys] + keys = [k for k in self._inflight if k[:2] == (server_name, scope_key)] + inflight = [self._inflight.pop(k) for k in keys] + await self._close_owners(entries, inflight) async def close_session_if_current( self, @@ -648,12 +637,11 @@ class MCPSessionPool: session: ClientSession, ) -> bool: """Close *session* only if it is still the registered entry for the key.""" - key = (server_name, scope_key) with self._lock: - entry = self._entries.get(key) - if entry is None or entry[0] is not session: + key = next((k for k, entry in self._entries.items() if k[:2] == (server_name, scope_key) and entry[0] is session), None) + if key is None: return False - self._entries.pop(key) + entry = self._entries.pop(key) _session, loop, task, close_evt = entry await self._shutdown_entry(loop, task, close_evt) return True diff --git a/backend/tests/test_mcp_session_pool.py b/backend/tests/test_mcp_session_pool.py index 055ae8200..44d743dfd 100644 --- a/backend/tests/test_mcp_session_pool.py +++ b/backend/tests/test_mcp_session_pool.py @@ -292,7 +292,7 @@ async def test_close_scope(): assert cms[1].closed is False # t2 session still exists. - assert ("s", "t2") in pool._entries + assert ("s", "t2") in {k[:2] for k in pool._entries} @pytest.mark.asyncio @@ -327,7 +327,7 @@ async def test_close_session_only_evicts_the_exact_server_scope_pair(): assert cms[0].closed is True assert cms[1].closed is False assert cms[2].closed is False - assert set(pool._entries) == {("s2", "t1"), ("s1", "t2")} + assert {k[:2] for k in pool._entries} == {("s2", "t1"), ("s1", "t2")} @pytest.mark.asyncio @@ -477,7 +477,7 @@ mcp.run(transport="stdio") await wrapped.coroutine(runtime=runtime) assert exc_info.value.error.code == CONNECTION_CLOSED - assert ("crash", "user:thread") not in get_session_pool()._entries + assert ("crash", "user:thread") not in {k[:2] for k in get_session_pool()._entries} content, _artifact = await wrapped.coroutine(runtime=runtime) @@ -735,7 +735,7 @@ async def test_late_disconnect_from_old_session_does_not_evict_replacement(tmp_p with pytest.raises(anyio.ClosedResourceError): await late_call - assert pool._entries[("srv", "test-user-autouse:default")][0] is replacement + assert pool._entries[("srv", "test-user-autouse:default", asyncio.get_running_loop())][0] is replacement await pool.close_all() @@ -1296,7 +1296,7 @@ async def test_session_pool_tool_extracts_thread_id(): # The scope key is "{user_id}:{thread_id}"; the autouse fixture sets # the effective user to "test-user-autouse". pool = get_session_pool() - assert ("server", "test-user-autouse:from-config") in pool._entries + assert ("server", "test-user-autouse:from-config") in {k[:2] for k in pool._entries} @pytest.mark.asyncio @@ -1331,7 +1331,7 @@ async def test_session_pool_tool_default_scope(): await wrapped.coroutine(runtime=None, x=1) pool = get_session_pool() - assert ("server", "test-user-autouse:default") in pool._entries + assert ("server", "test-user-autouse:default") in {k[:2] for k in pool._entries} @pytest.mark.asyncio @@ -1371,7 +1371,7 @@ async def test_session_pool_tool_get_config_fallback(): await wrapped.coroutine(runtime=None, x=1) pool = get_session_pool() - assert ("server", "test-user-autouse:from-langgraph-config") in pool._entries + assert ("server", "test-user-autouse:from-langgraph-config") in {k[:2] for k in pool._entries} def test_session_pool_tool_sync_wrapper_path_is_safe(): @@ -1710,7 +1710,7 @@ async def test_close_scope_does_not_cross_tasks(): assert cms[0].closed is True assert cms[1].closed is False - assert ("s", "t2") in pool._entries + assert ("s", "t2") in {k[:2] for k in pool._entries} @pytest.mark.asyncio @@ -1761,8 +1761,8 @@ def test_close_all_sync_across_loops_does_not_cross_tasks(): assert len(pool._entries) == 0 -def test_get_session_replaces_session_from_closed_loop(): - """A pooled session whose owning loop has closed is evicted and recreated.""" +def test_owner_loop_shutdown_retires_session_records(): + """Loop shutdown closes the owner and removes records before the next call.""" pool = MCPSessionPool() cms: list[_CancelScopeCm] = [] @@ -1776,15 +1776,14 @@ def test_get_session_replaces_session_from_closed_loop(): # asyncio.run (mirrors the sync-tool path). asyncio.run cancels the # pending owner task and runs its __aexit__ on the same loop. asyncio.run(pool.get_session("s", "t1", {"transport": "stdio", "command": "x", "args": []})) - assert ("s", "t1") in pool._entries + assert not pool._entries - # Now request the same key from a fresh loop: the stale entry (closed - # loop) must be evicted and replaced with a fresh session. + # A fresh loop creates a session without retaining its predecessor. session = asyncio.run(pool.get_session("s", "t1", {"transport": "stdio", "command": "x", "args": []})) assert session is not None assert len(cms) == 2 - assert pool._entries[("s", "t1")][0] is session + assert not pool._entries class _BlockingInitCm: @@ -1881,7 +1880,7 @@ async def test_get_session_cancelled_during_eviction_teardown_does_not_leak(): loop = asyncio.get_running_loop() victim_task = asyncio.create_task(victim_owner()) - pool._entries[("victim", "scope")] = (MagicMock(), loop, victim_task, asyncio.Event()) + pool._entries[("victim", "scope", asyncio.get_running_loop())] = (MagicMock(), loop, victim_task, asyncio.Event()) gate = asyncio.Event() cms: list[_BlockingInitCm] = [] @@ -1898,10 +1897,10 @@ async def test_get_session_cancelled_during_eviction_teardown_does_not_leak(): # 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: + if ("srv", "scope-2") in {k[:2] for k in pool._inflight}: break await asyncio.sleep(0.01) - assert ("srv", "scope-2") in pool._inflight + assert ("srv", "scope-2") in {k[:2] for k in pool._inflight} await asyncio.sleep(0.02) call.cancel() @@ -1919,7 +1918,7 @@ async def test_get_session_cancelled_during_eviction_teardown_does_not_leak(): 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 + assert ("srv", "scope-2") not in {k[:2] for k 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())] @@ -1960,7 +1959,7 @@ async def test_cancelled_creator_does_not_close_session_held_by_joiner(): loop = asyncio.get_running_loop() victim_task = asyncio.create_task(victim_owner()) - pool._entries[("victim", "scope")] = (MagicMock(), loop, victim_task, asyncio.Event()) + pool._entries[("victim", "scope", asyncio.get_running_loop())] = (MagicMock(), loop, victim_task, asyncio.Event()) gate = asyncio.Event() cms: list[_BlockingInitCm] = [] @@ -1978,10 +1977,10 @@ async def test_cancelled_creator_does_not_close_session_held_by_joiner(): # 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: + if ("srv", "scope-2") in {k[:2] for k in pool._inflight}: break await asyncio.sleep(0.01) - assert ("srv", "scope-2") in pool._inflight + assert ("srv", "scope-2") in {k[:2] for k in pool._inflight} # The owner finishes initialize() and commits while its creator is # still parked on the victim's teardown. The second caller then @@ -2000,9 +1999,9 @@ async def test_cancelled_creator_does_not_close_session_held_by_joiner(): 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 + assert ("srv", "scope-2") in {k[:2] for k in pool._entries}, "committed session must stay registered after creator cancellation" + assert pool._entries[("srv", "scope-2", asyncio.get_running_loop())][0] is session, "the joiner's session must be the registered one" + assert ("srv", "scope-2") not in {k[:2] for k in pool._inflight} # Cleanup: close the pool so the parked owner finishes deterministically. await pool.close_all() @@ -2048,10 +2047,10 @@ async def test_joiner_follows_creation_outcome_when_creator_is_cancelled(): 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: + if ("s", "same") in {k[:2] for k in pool._inflight}: break await asyncio.sleep(0.01) - assert ("s", "same") in pool._inflight + assert ("s", "same") in {k[:2] for k in pool._inflight} # Second caller joins the in-flight creation instead of duplicating it. joiner = asyncio.create_task(pool.get_session("s", "same", conn)) @@ -2113,8 +2112,8 @@ async def test_eviction_teardown_completes_normally_then_session_is_returned(): 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 + assert ("s", "t2") in {k[:2] for k in pool._entries} + assert ("s", "t1") not in {k[:2] for k in pool._entries} @pytest.mark.asyncio @@ -2170,8 +2169,8 @@ async def test_cancelling_close_scope_does_not_strand_other_removed_owners(): 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) + pool._entries[("s1", "scope", asyncio.get_running_loop())] = (MagicMock(), loop, first_task, first_close) + pool._entries[("s2", "scope", asyncio.get_running_loop())] = (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). @@ -2364,7 +2363,7 @@ async def test_queued_cancel_rechecks_failure_on_owning_loop(): gate2 = holder["gate"] proxy = _DelayedLoop(loop2) - pool._inflight[("s", "t1")] = (proxy, holder["ready"], holder["task"], holder["close_evt"]) + pool._inflight[("s", "t1", proxy)] = (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 @@ -2505,7 +2504,7 @@ async def test_close_all_during_in_flight_creation_does_not_resurrect_session(): call = asyncio.create_task(pool.get_session("s", "t1", conn)) # Let the owner task enter the CM and reach the blocking initialize(). await asyncio.sleep(0.01) - assert ("s", "t1") in pool._inflight + assert ("s", "t1") in {k[:2] for k in pool._inflight} # Close everything while the creation is still in-flight. await pool.close_all() @@ -2528,14 +2527,8 @@ async def test_close_all_during_in_flight_creation_does_not_resurrect_session(): assert not leaked, "in-flight owner task must not leak after close_all" -def test_get_session_cross_loop_in_flight_does_not_raise_assertion(): - """A same-key request from another loop must not hit the in-flight assertion (#3379 CR P1). - - Loop A starts (and leaves running) an in-flight creation, then loop B - requests the same key. The stale in-flight record (owned by loop A) must be - dropped and loop B must become a fresh creator — never fall through to an - AssertionError. - """ +def test_sequential_thread_loops_create_independent_sessions(): + """Sequential sync callers each create and retire their own session.""" pool = MCPSessionPool() cms: list[_CancelScopeCm] = [] @@ -2555,14 +2548,12 @@ def test_get_session_cross_loop_in_flight_does_not_raise_assertion(): errors.append(e) with patch("langchain_mcp_adapters.sessions.create_session", side_effect=make_cm): - # First loop creates and registers an entry, then its loop is torn down - # by asyncio.run, leaving a stale (closed-loop) record behind. + # First loop creates a session, then asyncio.run tears its owner down. t1 = threading.Thread(target=run_in_own_loop) t1.start() t1.join() - # Second loop requests the same key. It must evict the stale record and - # create a fresh session instead of raising AssertionError. + # The second loop requests the same server/scope independently. t2 = threading.Thread(target=run_in_own_loop) t2.start() t2.join() @@ -2572,14 +2563,8 @@ def test_get_session_cross_loop_in_flight_does_not_raise_assertion(): assert all(r is not None for r in results) -def test_cross_loop_preempting_blocked_in_flight_does_not_hang_owner(): - """A foreign-loop request must not leave a still-initializing owner hung (#3379 CR P1). - - Loop A starts a creation that blocks inside initialize() (the in-flight - record stays live). Loop B then requests the same key. B must tear A's owner - down — cancelling it, because close_evt alone cannot wake a task blocked in - initialize() — so that A's get_session unwinds instead of hanging forever. - """ +def test_cross_loop_caller_does_not_cancel_live_in_flight_owner(): + """A sibling loop can finish while the first loop's handshake is blocked.""" pool = MCPSessionPool() conn = {"transport": "stdio", "command": "x", "args": []} first_gate = threading.Event() @@ -2639,13 +2624,18 @@ def test_cross_loop_preempting_blocked_in_flight_does_not_hang_owner(): # B must complete without depending on A's blocked initialize(). assert not tb.is_alive(), "foreign-loop request B must not hang" - # A must already be unwound (cancelled), not waiting on the dead gate. - ta.join(3) - assert not ta.is_alive(), "preempted owner A must not hang forever" + try: + assert ta.is_alive(), "B must not cancel A's live handshake" + assert not errors + finally: + first_gate.set() + ta.join(3) + assert not ta.is_alive() - assert [n for n, _ in results] == ["B"], "only B produces a usable session" - assert any(isinstance(e, asyncio.CancelledError) for _, e in errors), "preempted A must unwind via CancelledError" - assert "blocking" in closed, "preempted owner's __aexit__ must run on teardown" + assert sorted(n for n, _ in results) == ["A", "B"] + assert not errors + assert "blocking" in closed + assert not pool._entries and not pool._inflight @pytest.mark.asyncio diff --git a/backend/tests/test_mcp_session_pool_loop_isolation.py b/backend/tests/test_mcp_session_pool_loop_isolation.py new file mode 100644 index 000000000..420db8e74 --- /dev/null +++ b/backend/tests/test_mcp_session_pool_loop_isolation.py @@ -0,0 +1,200 @@ +"""Persistent-loop regressions for concurrent sync callers (#5256).""" + +import asyncio +import concurrent.futures +import sys +import threading +from contextlib import asynccontextmanager +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from deerflow.mcp.session_pool import MCPSessionPool +from deerflow.tools.sync import make_sync_tool_wrapper + + +@pytest.fixture +def loop_pool(monkeypatch): + pool = MCPSessionPool() + loops = [asyncio.new_event_loop(), asyncio.new_event_loop()] + threads = [threading.Thread(target=loop.run_forever) for loop in loops] + closed = [] + + @asynccontextmanager + async def create_session(connection): + owner = asyncio.current_task() + session = AsyncMock() + try: + yield session + finally: + assert asyncio.current_task() is owner + closed.append(session) + + monkeypatch.setattr("langchain_mcp_adapters.sessions.create_session", create_session) + for thread in threads: + thread.start() + + def run(index, coroutine): + return asyncio.run_coroutine_threadsafe(coroutine, loops[index]).result(timeout=5) + + try: + yield pool, run, closed + finally: + run(0, pool.close_all()) + for loop in loops: + loop.call_soon_threadsafe(loop.stop) + for thread in threads: + thread.join(timeout=5) + assert not thread.is_alive() + for loop in loops: + loop.close() + + +def test_live_loops_reuse_only_their_own_sessions_and_disconnect_is_isolated(loop_pool): + pool, run, closed = loop_pool + first = run(0, pool.get_session("s", "u:t", {})) + sibling = run(1, pool.get_session("s", "u:t", {})) + assert first is not sibling + assert not closed + assert run(0, pool.get_session("s", "u:t", {})) is first + assert run(1, pool.get_session("s", "u:t", {})) is sibling + assert run(0, pool.close_session_if_current("s", "u:t", first)) + assert closed == [first] + replacement = run(0, pool.get_session("s", "u:t", {})) + assert replacement is not first + assert not run(0, pool.close_session_if_current("s", "u:t", first)) + assert run(1, pool.get_session("s", "u:t", {})) is sibling + assert run(0, pool.get_session("s", "u:t", {})) is replacement + + +@pytest.mark.parametrize("operation", ["scope", "server", "session", "all"]) +def test_explicit_cleanup_closes_both_loops(loop_pool, operation): + pool, run, closed = loop_pool + sessions = [run(i, pool.get_session("s", "u:t", {})) for i in range(2)] + other = run(1, pool.get_session("other", "u:other", {})) + cleanup = { + "scope": lambda: pool.close_scope("u:t"), + "server": lambda: pool.close_server("s"), + "session": lambda: pool.close_session("s", "u:t"), + "all": pool.close_all, + }[operation] + run(0, cleanup()) + assert all(session in closed for session in sessions) + if operation == "all": + assert other in closed + else: + assert other not in closed + assert run(1, pool.get_session("other", "u:other", {})) is other + + +def test_owner_completion_does_not_remove_a_replacement(loop_pool): + pool, run, _closed = loop_pool + + async def replace(): + await pool.get_session("s", "u:t", {}) + key = ("s", "u:t", asyncio.get_running_loop()) + old_owner = pool._entries[key][2] + await pool.close_session("s", "u:t") + replacement = await pool.get_session("s", "u:t", {}) + pool._discard_owner(key, old_owner) + assert await pool.get_session("s", "u:t", {}) is replacement + + run(0, replace()) + + +@pytest.mark.parametrize("retirement", ["lru", "explicit"]) +def test_abandoned_closed_loop_entry_can_be_retired(loop_pool, retirement): + """Model the registry left by unsupported loop.close() with pending owners. + + Use a synthetic pending owner so the test itself does not leak a real task + or transport on a closed loop. Normal owner shutdown is tested separately. + """ + pool, run, _closed = loop_pool + closed_loop = asyncio.new_event_loop() + closed_loop.close() + owner = MagicMock(spec=asyncio.Task) + owner.done.return_value = False + key = ("s", "u:t", closed_loop) + pool._entries[key] = (MagicMock(), closed_loop, owner, asyncio.Event()) + pool.MAX_SESSIONS = 1 + if retirement == "explicit": + run(0, pool.close_scope("u:t")) + assert not pool._entries + else: + replacement = run(0, pool.get_session("s", "u:t", {})) + assert key not in pool._entries + assert len(pool._entries) == 1 + assert run(0, pool.get_session("s", "u:t", {})) is replacement + owner.cancel.assert_not_called() + + +@pytest.mark.parametrize("operation", ["scope", "server", "session", "all"]) +def test_cleanup_cancels_inflight_owners_on_both_loops(loop_pool, monkeypatch, operation): + pool, run, _closed = loop_pool + started = [threading.Event(), threading.Event()] + exited = [threading.Event(), threading.Event()] + + @asynccontextmanager + async def create_session(connection): + index = connection["index"] + owner = asyncio.current_task() + + async def initialize(): + started[index].set() + await asyncio.Future() + + try: + yield AsyncMock(initialize=initialize) + finally: + assert asyncio.current_task() is owner + exited[index].set() + + monkeypatch.setattr("langchain_mcp_adapters.sessions.create_session", create_session) + + async def start(index): + return asyncio.create_task(pool.get_session("s", "u:t", {"index": index})) + + calls = [run(i, start(i)) for i in range(2)] + assert all(event.wait(5) for event in started) + cleanup = { + "scope": lambda: pool.close_scope("u:t"), + "server": lambda: pool.close_server("s"), + "session": lambda: pool.close_session("s", "u:t"), + "all": pool.close_all, + }[operation] + run(0, cleanup()) + assert all(event.wait(5) for event in exited) + + async def cancelled(call): + with pytest.raises(asyncio.CancelledError): + await call + + for i, call in enumerate(calls): + run(i, cancelled(call)) + assert not pool._inflight and not pool._entries + + +def test_parallel_sync_wrappers_complete_real_stdio_calls(tmp_path): + server = tmp_path / "echo_server.py" + server.write_text( + 'from mcp.server.fastmcp import FastMCP\nmcp = FastMCP("echo")\n@mcp.tool()\ndef echo(text: str) -> str:\n return text\nmcp.run(transport="stdio")\n', + encoding="utf-8", + ) + pool = MCPSessionPool() + barrier = threading.Barrier(2) + connection = {"transport": "stdio", "command": sys.executable, "args": [str(server)]} + + async def echo(text): + await asyncio.to_thread(barrier.wait, 5) + session = await pool.get_session("echo", "user:thread", connection) + result = await session.call_tool("echo", {"text": text}) + return result.content[0].text + + wrapper = make_sync_tool_wrapper(echo, "echo") + try: + with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor: + futures = [executor.submit(wrapper, text) for text in ("A", "B")] + assert [future.result(timeout=30) for future in futures] == ["A", "B"] + assert not pool._entries and not pool._inflight + finally: + pool.close_all_sync()