github/deer-flow
1
0
Fork 0
mirror of https://github.com/bytedance/deer-flow.git synced 2026-09-14 16:08:41 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(mcp): isolate pooled sessions by owning event loop (#5396)

Browse Source 
* fix(mcp): isolate pooled sessions by owning event loop

* refactor(mcp): remove obsolete eviction cancellation plumbing
This commit is contained in:
yang rui 2026-09-13 17:13:45 +08:00 committed by GitHub
parent 56540fab01
commit dfe9a520b9
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 326 additions and 114 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
README.md
View File

@ -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

14
backend/docs/MCP_SERVER.md
View File

@ -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

13
backend/packages/harness/deerflow/mcp/AGENTS.md
View File

@ -6,6 +6,19 @@
- **Long-running ordinary task driver**: `extensions_config.json -> mcpServers.<server>.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

100
backend/packages/harness/deerflow/mcp/session_pool.py
View File

@ -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

106
backend/tests/test_mcp_session_pool.py
View File

@ -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

200
backend/tests/test_mcp_session_pool_loop_isolation.py Normal file
View File

@ -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()