fix(mcp): close thread-scoped MCP sessions on thread deletion (#5956)

* fix(mcp): close thread-scoped MCP sessions on thread deletion

Deleting a thread removed its filesystem data, checkpoints, runs, events,
feedback, metadata and browser session, but never closed its persistent MCP
sessions. MCPSessionPool.close_scope() had no production caller at all, so a
recreated thread id under the same user could reuse stale stateful MCP server
state, and the leaked owner task plus subprocess kept holding memory and file
descriptors.

Add MCPSessionPool.close_thread_scope(user_id, thread_id) and call it from the
reserved deletion path under the same best-effort policy as the neighbouring
cleanups. It matches every incarnation of the thread -- the legacy user:thread
scope and any v2:[user, thread, incarnation] scope -- because a deletion cannot
reliably read the current incarnation: the record may predate incarnation
tracking, and a new incarnation can be minted concurrently. close_scope() is
unchanged for callers that already hold one exact key.

Add mcp_scope_belongs_to_thread() to own the matching rule, and hoist the v2:
prefix into a shared constant so the encoder and matcher cannot drift.

Refs #5188

* test(mcp): bind owner identity in thread cleanup checks

---------

Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
This commit is contained in:
NanHan Qing 2026-09-29 16:53:18 +08:00 committed by GitHub
parent 8a3a309d1c
commit f8b69bd84b
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
9 changed files with 391 additions and 13 deletions

View File

@ -830,6 +830,22 @@ async def _delete_thread_data_with_reservation(thread_id: str, request: Request)
except Exception:
logger.debug("Could not close browser session for %s (not critical)", sanitize_log_param(thread_id))
# Tear down persistent MCP sessions scoped to this user/thread (best-effort).
# The same invariant as the browser session above applies, and harder: a
# persistent stdio MCP server such as Playwright keeps retained pages and
# cookies, and its leaked owner task plus subprocess keep holding memory and
# file descriptors. Scope keys encode user/thread/incarnation, so the whole
# thread identity is closed across every incarnation -- a legacy
# (incarnation-less) scope and a newer versioned scope are both stale once
# the thread is gone, and reading the current incarnation here would race a
# concurrently minted one. See #5188.
try:
from deerflow.mcp.session_pool import get_session_pool
await get_session_pool().close_thread_scope(user_id=user_id, thread_id=thread_id)
except Exception:
logger.debug("Could not close MCP sessions for %s (not critical)", sanitize_log_param(thread_id))
return response

View File

@ -15,10 +15,17 @@
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.
scope, pair, thread-identity and global cleanup span all loops; disconnect
eviction matches the exact session identity. Thread-identity cleanup
(`close_thread_scope`) deliberately matches *every incarnation* of a
user/thread pair instead of one exact scope key: a thread-deletion caller
cannot rely on reading the current incarnation (the record may predate
incarnation tracking, and a new incarnation can be minted concurrently), and
every generation of a deleted thread is equally stale. Keep `close_scope` for
callers that already hold one exact key. 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`

View File

@ -52,6 +52,8 @@ from mcp import ClientSession
from mcp.shared.exceptions import McpError
from mcp.types import CONNECTION_CLOSED
from deerflow.mcp_scope import mcp_scope_belongs_to_thread
logger = logging.getLogger(__name__)
_MCP_CLOSED_STREAM_ERRORS = (
@ -621,6 +623,32 @@ class MCPSessionPool:
inflight = [self._inflight.pop(k) for k in inflight_keys]
await self._close_owners(entries, inflight)
async def close_thread_scope(self, *, user_id: str, thread_id: str) -> None:
"""Close every session scoped to one user/thread identity.
Differs from :meth:`close_scope` in matching *all incarnations* of the
thread rather than one exact scope key. A thread-deletion path knows the
user and thread but not reliably the incarnation: the record may predate
incarnation tracking (legacy ``user:thread`` scope), and a new
incarnation can be minted between reading the record and tearing down.
Keying cleanup on a read incarnation would therefore both miss legacy
sessions and race a concurrent incarnation. Every generation of a
deleted thread is equally stale, so all of them are closed.
``user_id`` must be the same string used by
``resolve_runtime_user_id(runtime)`` when minting the session scope.
The Gateway delete route passes ``get_effective_user_id()`` and relies
on both resolving to the same identity in its embedded runtime.
Sessions belonging to another user or thread are left untouched.
"""
with self._lock:
keys = [k for k in self._entries if mcp_scope_belongs_to_thread(k[1], user_id=user_id, thread_id=thread_id)]
entries = [self._entries.pop(k) for k in keys]
inflight_keys = [k for k in self._inflight if mcp_scope_belongs_to_thread(k[1], user_id=user_id, thread_id=thread_id)]
inflight = [self._inflight.pop(k) for k in inflight_keys]
await self._close_owners(entries, inflight)
async def close_session(self, server_name: str, scope_key: str) -> None:
"""Close every session for this server/scope across all owning loops."""
with self._lock:

View File

@ -9,6 +9,11 @@ THREAD_INCARNATION_CONTEXT_KEY = "thread_incarnation"
THREAD_INCARNATION_METADATA_GUARD_KEY = "__deerflow_thread_incarnation_metadata_guard"
_MISSING = object()
# Prefix that distinguishes a versioned scope from the legacy ``user:thread``
# encoding. Shared with :func:`mcp_scope_belongs_to_thread` so the two can
# never drift apart.
_V2_SCOPE_PREFIX = "v2:"
def is_valid_thread_incarnation(value: object) -> TypeGuard[str | None]:
"""Return whether *value* is a supported legacy or versioned incarnation."""
@ -33,13 +38,39 @@ def mcp_session_scope_key(
return scope
# A JSON tuple is an unambiguous, versioned encoding even when an opaque
# user/thread id contains the delimiter used by the legacy scope.
return "v2:" + json.dumps(
return _V2_SCOPE_PREFIX + json.dumps(
[user_id, thread_id, thread_incarnation],
ensure_ascii=True,
separators=(",", ":"),
)
def mcp_scope_belongs_to_thread(scope_key: str, *, user_id: str, thread_id: str) -> bool:
"""Return whether *scope_key* was minted for this user/thread identity.
Matches both encodings :func:`mcp_session_scope_key` can produce: the legacy
``user:thread`` scope (NULL incarnation) and the versioned
``v2:[user, thread, incarnation]`` scope. The incarnation is deliberately
*not* compared — a caller tearing down a whole thread has to invalidate every
generation of it, and the current generation cannot be read reliably at
teardown time (see ``MCPSessionPool.close_thread_scope``).
The legacy branch is a prefix-free string compare, so a crafted ``user_id``
containing ``:`` can collide with another owner's legacy key. That can only
ever over-close (tear down an extra session), never leak one, and the
versioned encoding — which every incarnation-aware caller uses — is exact.
"""
if scope_key == f"{user_id}:{thread_id}":
return True
if not scope_key.startswith(_V2_SCOPE_PREFIX):
return False
try:
decoded = json.loads(scope_key[len(_V2_SCOPE_PREFIX) :])
except ValueError:
return False
return isinstance(decoded, list) and len(decoded) == 3 and decoded[0] == user_id and decoded[1] == thread_id
def runtime_thread_incarnation(runtime: Any | None) -> str | None:
"""Read the server-owned incarnation from ``ToolRuntime.context``."""
if runtime is None:

View File

@ -2,6 +2,13 @@
Backend tests must preserve the runtime invariants they exercise without changing production execution topology.
## Router auth fixtures
For owner-scoped route assertions, pass a stable `user_factory` and
`bind_current_user=True` to `make_authed_test_app`. The default stub stamps
request state but leaves the ambient user context unchanged; the opt-in binding
uses the stub user during the request and restores the previous context after it.
## Lark CLI blocking-I/O fixtures
`blocking_io/test_integrations_router.py` uses a real local CLI stub: a `.cmd`

View File

@ -13,7 +13,8 @@ This module provides two surfaces:
1. :func:`make_authed_test_app` — wraps ``FastAPI()`` with a tiny
``BaseHTTPMiddleware`` that stamps a fake user / AuthContext on every
request, plus a permissive ``thread_store`` mock on
``app.state``. Use from TestClient-based router tests.
``app.state``. Owner-scoped tests can opt into binding that user to the
request context. Use from TestClient-based router tests.
2. :func:`call_unwrapped` — invokes the underlying function bypassing
the ``@require_permission`` decorator chain by walking ``__wrapped__``.
@ -38,6 +39,7 @@ from starlette.types import ASGIApp
from app.gateway.auth.models import User
from app.gateway.authz import AuthContext, Permissions
from deerflow.runtime.user_context import reset_current_user, set_current_user
# Default permission set granted to the stub user. Mirrors `_ALL_PERMISSIONS`
# in authz.py — kept inline so the tests don't import a private symbol.
@ -73,21 +75,29 @@ class _StubAuthMiddleware(BaseHTTPMiddleware):
authenticated context and skips its own re-authentication path.
"""
def __init__(self, app: ASGIApp, user_factory: Callable[[], User]) -> None:
def __init__(self, app: ASGIApp, user_factory: Callable[[], User], bind_current_user: bool = False) -> None:
super().__init__(app)
self._user_factory = user_factory
self._bind_current_user = bind_current_user
async def dispatch(self, request: Request, call_next: Callable) -> Response:
user = self._user_factory()
request.state.user = user
request.state.auth = AuthContext(user=user, permissions=list(_STUB_PERMISSIONS))
return await call_next(request)
if not self._bind_current_user:
return await call_next(request)
token = set_current_user(user)
try:
return await call_next(request)
finally:
reset_current_user(token)
def make_authed_test_app(
*,
user_factory: Callable[[], User] | None = None,
owner_check_passes: bool = True,
bind_current_user: bool = False,
) -> FastAPI:
"""Build a FastAPI test app with stub auth + permissive thread_store.
@ -99,6 +109,8 @@ def make_authed_test_app(
returns True for every call so ``@require_permission(owner_check=True)``
never blocks the route under test. Pass False to verify that
permission failures surface correctly.
bind_current_user: Also bind the stub user to the request's user
context when a test needs to assert owner-scoped behavior.
Returns:
A ``FastAPI`` app with the stub middleware installed and
@ -107,7 +119,7 @@ def make_authed_test_app(
"""
factory = user_factory or _make_stub_user
app = FastAPI()
app.add_middleware(_StubAuthMiddleware, user_factory=factory)
app.add_middleware(_StubAuthMiddleware, user_factory=factory, bind_current_user=bind_current_user)
repo = MagicMock()
repo.check_access = AsyncMock(return_value=owner_check_passes)

View File

@ -0,0 +1,98 @@
"""Tests for canonical MCP session scope construction and matching."""
import pytest
from deerflow.mcp_scope import (
mcp_scope_belongs_to_thread,
mcp_session_scope_key,
)
def test_legacy_scope_is_user_colon_thread():
assert mcp_session_scope_key(user_id="u1", thread_id="t1", thread_incarnation=None) == "u1:t1"
def test_versioned_scope_is_prefixed_json_tuple():
scope = mcp_session_scope_key(user_id="u1", thread_id="t1", thread_incarnation="inc-1")
assert scope == 'v2:["u1","t1","inc-1"]'
def test_versioned_scope_disambiguates_ids_containing_the_legacy_delimiter():
"""The v2 encoding must stay unambiguous when an id contains ``:``."""
left = mcp_session_scope_key(user_id="a:b", thread_id="c", thread_incarnation="inc")
right = mcp_session_scope_key(user_id="a", thread_id="b:c", thread_incarnation="inc")
assert left != right
def test_scope_rejects_invalid_incarnation():
with pytest.raises(RuntimeError):
mcp_session_scope_key(user_id="u1", thread_id="t1", thread_incarnation="")
@pytest.mark.parametrize(
"scope_key",
[
"u1:t1",
'v2:["u1","t1","inc-1"]',
'v2:["u1","t1","inc-2"]',
],
)
def test_belongs_to_thread_matches_legacy_and_every_incarnation(scope_key):
"""Teardown matches the legacy scope and any versioned incarnation.
A deleted thread cannot reliably know its current incarnation, and every
generation of it is stale once the thread is gone (#5188).
"""
assert mcp_scope_belongs_to_thread(scope_key, user_id="u1", thread_id="t1") is True
@pytest.mark.parametrize(
"scope_key",
[
"u1:t2",
"u2:t1",
"u1:t10",
"u10:t1",
'v2:["u1","t2","inc-1"]',
'v2:["u2","t1","inc-1"]',
],
)
def test_belongs_to_thread_rejects_other_user_or_thread(scope_key):
assert mcp_scope_belongs_to_thread(scope_key, user_id="u1", thread_id="t1") is False
@pytest.mark.parametrize(
"scope_key",
[
'v2:["u1","t1",null]',
'v2:["u1","t1",1]',
'v2:["u1","t1",""]',
],
)
def test_belongs_to_thread_matches_on_identity_not_incarnation_shape(scope_key):
"""A key carrying this thread's identity is claimed even when its third
element is not a shape the current encoder emits.
Matching is deliberately lenient about the incarnation: teardown must never
*miss* a session belonging to the deleted thread, and an unexpected
incarnation shape cannot make the session belong to somebody else.
"""
assert mcp_scope_belongs_to_thread(scope_key, user_id="u1", thread_id="t1") is True
@pytest.mark.parametrize(
"scope_key",
[
"",
"v2:",
"v2:not-json",
'v2:{"user":"u1"}',
'v2:["u1","t1"]',
'v2:["u1","t1","inc","extra"]',
'v3:["u1","t1","inc"]',
"u1",
],
)
def test_belongs_to_thread_rejects_malformed_scopes(scope_key):
"""Unknown or structurally invalid encodings must not be claimed as this thread's."""
assert mcp_scope_belongs_to_thread(scope_key, user_id="u1", thread_id="t1") is False

View File

@ -18,6 +18,7 @@ from mcp.types import CONNECTION_CLOSED, CallToolResult, ErrorData, TextContent
from deerflow.mcp.session_pool import MCPSessionPool, call_pooled_session_tool, get_session_pool, reset_session_pool
from deerflow.mcp_scope import (
THREAD_INCARNATION_METADATA_GUARD_KEY,
mcp_scope_belongs_to_thread,
mcp_session_scope_key,
runtime_thread_incarnation,
)
@ -370,6 +371,131 @@ async def test_close_scope():
assert ("s", "t2") in {k[:2] for k in pool._entries}
@pytest.mark.asyncio
async def test_close_thread_scope_closes_every_incarnation_of_one_thread():
"""Thread deletion must invalidate all of the thread's scopes, and only its own.
A thread can hold both the legacy ``user:thread`` scope and a versioned
``v2:[user, thread, incarnation]`` scope, so a teardown that closed a single
exact key would leave the other alive. Sessions of another thread or another
user must survive (#5188).
"""
pool = MCPSessionPool()
class CmFactory:
def __init__(self):
self.closed = False
async def __aenter__(self):
return AsyncMock()
async def __aexit__(self, *args):
self.closed = True
return False
cms: list[CmFactory] = []
def make_cm(*a, **kw):
cm = CmFactory()
cms.append(cm)
return cm
legacy = mcp_session_scope_key(user_id="u1", thread_id="t1", thread_incarnation=None)
versioned = mcp_session_scope_key(user_id="u1", thread_id="t1", thread_incarnation="inc-2")
other_thread = mcp_session_scope_key(user_id="u1", thread_id="t2", thread_incarnation="inc-1")
other_user = mcp_session_scope_key(user_id="u2", thread_id="t1", thread_incarnation="inc-1")
with patch("langchain_mcp_adapters.sessions.create_session", side_effect=make_cm):
for scope in (legacy, versioned, other_thread, other_user):
await pool.get_session("s", scope, {"transport": "stdio", "command": "x", "args": []})
await pool.close_thread_scope(user_id="u1", thread_id="t1")
assert cms[0].closed is True, "legacy scope must be closed"
assert cms[1].closed is True, "versioned scope must be closed"
assert cms[2].closed is False, "another thread of the same user must survive"
assert cms[3].closed is False, "the same thread id under another user must survive"
remaining = {k[:2] for k in pool._entries}
assert ("s", other_thread) in remaining
assert ("s", other_user) in remaining
assert not any(mcp_scope_belongs_to_thread(k[1], user_id="u1", thread_id="t1") for k in pool._entries)
@pytest.mark.asyncio
async def test_close_thread_scope_tears_down_inflight_creation_for_same_thread():
"""A session being created when the thread is deleted must not be admitted.
The pool registers a session only after ``initialize()`` succeeds, so an
in-flight creation for the deleted thread has to be cancelled out of
``_inflight``; otherwise it commits afterwards and recreates the very state
the deletion removed (#5188).
"""
pool = MCPSessionPool()
gate = asyncio.Event()
class CmFactory:
def __init__(self):
self.closed = False
async def __aenter__(self):
return AsyncMock()
async def __aexit__(self, *args):
self.closed = True
return False
established_cms: list[CmFactory] = []
def make_established(*a, **kw):
cm = CmFactory()
established_cms.append(cm)
return cm
blocked = _BlockingInitCm(gate)
def make_blocked(*a, **kw):
return blocked
target_scope = mcp_session_scope_key(user_id="u1", thread_id="t1", thread_incarnation="inc-1")
other_scope = mcp_session_scope_key(user_id="u1", thread_id="t2", thread_incarnation="inc-1")
with patch("langchain_mcp_adapters.sessions.create_session", side_effect=make_established):
await pool.get_session("s", other_scope, {"transport": "stdio", "command": "x", "args": []})
with patch("langchain_mcp_adapters.sessions.create_session", side_effect=make_blocked):
conn = {"transport": "stdio", "command": "x", "args": []}
creator = asyncio.create_task(pool.get_session("s", target_scope, conn))
for _ in range(200):
if target_scope in {k[1] for k in pool._inflight}:
break
await asyncio.sleep(0.005)
assert target_scope in {k[1] for k in pool._inflight}, "creation must be in flight"
await pool.close_thread_scope(user_id="u1", thread_id="t1")
assert target_scope not in {k[1] for k in pool._inflight}, "in-flight creation must be removed"
gate.set()
with pytest.raises(asyncio.CancelledError):
await creator
assert blocked.closed is True, "the cancelled creation must still run its __aexit__"
assert established_cms[0].closed is False, "another thread's session must survive"
# The surviving session's owner task is legitimately parked on its close
# event, so retire it before checking that the cancelled creation left no
# owner task behind.
await pool.close_scope(other_scope)
leaked: list[asyncio.Task] = []
for _ in range(200):
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())]
if not leaked:
break
await asyncio.sleep(0.005)
assert not leaked, "owner task must not be left pending after the thread teardown"
@pytest.mark.asyncio
async def test_close_session_only_evicts_the_exact_server_scope_pair():
pool = MCPSessionPool()

View File

@ -16,6 +16,7 @@ from langgraph.store.memory import InMemoryStore
from langgraph.types import Overwrite
from app.gateway import services as gateway_services
from app.gateway.auth.models import User
from app.gateway.routers import thread_runs, threads
from deerflow.config.paths import Paths
from deerflow.persistence.engine import close_engine, get_session_factory, init_engine
@ -413,15 +414,14 @@ def test_delete_thread_data_rejects_invalid_thread_id(tmp_path):
def test_delete_thread_route_cleans_thread_directory(tmp_path):
from deerflow.runtime.user_context import get_effective_user_id
paths = Paths(tmp_path)
user_id = get_effective_user_id()
owner = User(email="thread-owner@example.com", password_hash="x")
user_id = str(owner.id)
thread_dir = paths.thread_dir("thread-route", user_id=user_id)
paths.sandbox_work_dir("thread-route", user_id=user_id).mkdir(parents=True, exist_ok=True)
(paths.sandbox_work_dir("thread-route", user_id=user_id) / "notes.txt").write_text("hello", encoding="utf-8")
app = make_authed_test_app()
app = make_authed_test_app(user_factory=lambda: owner, bind_current_user=True)
app.state.run_manager = _ThreadTestRunManager()
app.include_router(threads.router)
@ -458,6 +458,59 @@ def test_delete_thread_route_closes_browser_session(tmp_path):
manager.close_session.assert_awaited_once_with("thread-browser")
def test_delete_thread_route_closes_mcp_sessions(tmp_path):
"""Deleting a thread tears down its persistent MCP sessions so a later
caller who reuses the id gets fresh MCP server state instead of a retained
session (and its leaked owner task plus subprocess) — the same invariant the
browser-session cleanup above guards (#5188)."""
paths = Paths(tmp_path)
owner = User(email="mcp-owner@example.com", password_hash="x")
app = make_authed_test_app(user_factory=lambda: owner, bind_current_user=True)
app.state.run_manager = _ThreadTestRunManager()
app.include_router(threads.router)
pool = SimpleNamespace(close_thread_scope=AsyncMock(return_value=None))
with (
patch("app.gateway.routers.threads.get_paths", return_value=paths),
patch("deerflow.mcp.session_pool.get_session_pool", return_value=pool),
):
with TestClient(app) as client:
response = client.delete("/api/threads/thread-mcp")
assert response.status_code == 200
pool.close_thread_scope.assert_awaited_once_with(
user_id=str(owner.id),
thread_id="thread-mcp",
)
def test_delete_thread_route_isolates_failing_mcp_cleanup(tmp_path):
"""A failing MCP teardown must stay best-effort like every other cleanup step.
The delete has already removed the thread's filesystem data, checkpoints and
metadata by the time MCP sessions are closed, so an exception here must not
turn a successful deletion into a 500.
"""
paths = Paths(tmp_path)
app = make_authed_test_app()
app.state.run_manager = _ThreadTestRunManager()
app.include_router(threads.router)
pool = SimpleNamespace(close_thread_scope=AsyncMock(side_effect=RuntimeError("simulated MCP teardown failure")))
with (
patch("app.gateway.routers.threads.get_paths", return_value=paths),
patch("deerflow.mcp.session_pool.get_session_pool", return_value=pool),
):
with TestClient(app) as client:
response = client.delete("/api/threads/thread-mcp-fail")
assert response.status_code == 200
assert response.json() == {"success": True, "message": "Deleted local thread data for thread-mcp-fail"}
pool.close_thread_scope.assert_awaited_once()
def _persistence_cleanup_app(tmp_path, *, run_store, event_store, feedback_repo):
app = make_authed_test_app()
app.state.run_manager = _ThreadTestRunManager()