mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-26 13:16:24 +00:00
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:
parent
8a3a309d1c
commit
f8b69bd84b
@ -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
|
||||
|
||||
|
||||
|
||||
@ -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`
|
||||
|
||||
@ -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:
|
||||
|
||||
@ -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:
|
||||
|
||||
@ -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`
|
||||
|
||||
@ -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)
|
||||
|
||||
98
backend/tests/test_mcp_scope.py
Normal file
98
backend/tests/test_mcp_scope.py
Normal 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
|
||||
@ -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()
|
||||
|
||||
@ -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()
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user