github/deer-flow
1
0
Fork 0
mirror of https://github.com/bytedance/deer-flow.git synced 2026-10-26 21:26:39 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(gateway): close pooled MCP sessions on shutdown (#5961)

Browse Source 
* fix(gateway): close pooled MCP sessions on shutdown

The Gateway lifespan closes the browser session manager on the way out but
never closed MCPSessionPool, so every pooled MCP session was abandoned when
the worker exited: the owner task's transport __aexit__ never ran and a stdio
MCP server or SSE/HTTP connection had no orderly teardown.

MCPSessionPool.close_all() already existed for exactly this and had no
production caller -- only the config-retirement path used close_all_sync().
The async form is the right one here: close_all_sync()'s own docstring directs
callers running inside a live event loop to await close_all() instead of
signalling and returning.

Wired into lifespan shutdown with the same bounded, best-effort shape as the
neighbouring browser-session hook.

* fix(gateway): close MCP sessions after run drain

---------

Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
This commit is contained in:
NanHan Qing 2026-09-28 09:14:27 +08:00 committed by GitHub
parent bc0167636f
commit 2377a49217
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 219 additions and 22 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

49
backend/app/gateway/AGENTS.md
View File

@ -4,12 +4,12 @@ Reject external run/state writes with `sandbox`, `thread_data`, or `viewed_image
Studio retains sanitized creation metadata.
Capability Center's `business` adapter validates bundled-provider credentials
and creates an MCP connection. The MCP API accepts only the exact isolated
interpreter/module/provider launcher and credential environment keys generated
by `deerflow.capabilities.business`. Reject arbitrary Python commands; manifest
metadata cannot bypass execution policy, and credentials stay out of tool
schemas. Preserve admin checks, masked edits, atomic writes and MCP cache reloads.
Capability Center's `business` adapter validates bundled credentials and
creates an MCP connection. The MCP API permits only the isolated launcher and
credential environment keys generated by `deerflow.capabilities.business`;
reject arbitrary Python commands and manifest policy bypasses. Keep credentials
out of tool schemas; preserve admin checks, masked edits, atomic writes, and
MCP cache reloads.
Memory shutdown resolves hot-reloaded config and the backend, flushes, then
closes as one `await_drained` operation. Keep config resolution inside the
@ -23,21 +23,30 @@ is drained across host cancellation before later runtime resources unwind; each
service `stop()` remains bounded by its 30-second `asyncio.timeout`, so include
that bound in pod grace-period budgeting.
`conversation_access.py` binds an opt-in read-only tool to a run request's
explicit `conversation_references` and effective `runs:read` permission. Never
derive grants from message contents or checkpoints. The callback travels through
`RunContext`, not serialized config; the worker owns its context injection and
terminal cleanup. `conversation_reader.py` shares the existing HTTP transcript
visibility/pagination logic without a Request dependency; ownership remains a
caller responsibility. The tool additionally excludes non-text/internal content.
`RunCreateRequest` also accepts the list as `context.conversation_references`
(LangGraph SDK clients cannot send top-level extras): a before-validator lifts
it into the top-level field, so bounds and error locations are shared, and drops
it from `context`, so `merge_run_context_overrides` never sees it; sending both
is a 422. `conversation_references_enabled()` is the one "tool is configured"
predicate, used by admission and by `/api/features`.
Gateway lifespan closes the MCP session pool after `langgraph_runtime` exits:
its run drain and extension-service teardown can still make MCP calls during
earlier shutdown hooks. Keep the pool close bounded and failure-isolated, but
do not move it ahead of those session producers.
FastAPI listens on port 8001; health: `GET /health` (liveness) and `GET /health/ready` (readiness; concurrently probes the ORM engine behind `database:` plus the effective LangGraph checkpointer/Store backend - the legacy `checkpointer:` section, otherwise derived from `database:`, resolved from the startup config snapshot recorded on `app.state` - beneath a single bounded deadline, with connection-opening probes serialized behind a strict per-process gate, 503 while either is unreachable or the startup backend cannot be resolved, `not_configured` for process-local backends such as `backend=memory`). Set `GATEWAY_ENABLE_DOCS=false` to disable the default `/docs`, `/redoc`, and `/openapi.json` endpoints.
`conversation_access.py` grants the opt-in read-only tool only for explicit
`conversation_references` with effective `runs:read`; never infer grants from
messages or checkpoints. Pass the callback through `RunContext`, not serialized
config, and let the worker inject and clear it. `conversation_reader.py` shares
HTTP transcript visibility/pagination without a Request; callers own authorization.
Exclude non-text/internal content. `RunCreateRequest` lifts SDK-compatible
`context.conversation_references` into its top-level field before validation,
then removes it from context; supplying both is 422. This shares bounds/error
locations and keeps `merge_run_context_overrides` clear. Admission and
`/api/features` use `conversation_references_enabled()` as the sole configured
predicate.
FastAPI listens on port 8001. `GET /health` is liveness; `GET /health/ready`
concurrently probes the ORM engine and effective LangGraph checkpointer/Store
under one deadline and a strict connection-opening gate. It reads the startup
config snapshot: legacy `checkpointer:` if set, else `database:`. Unreachable or
unresolved backends return 503; process-local backends such as `memory` report
`not_configured`. `GATEWAY_ENABLE_DOCS=false` disables `/docs`, `/redoc`, and
`/openapi.json`.
`build_run_config()` resolves the default LangGraph super-step budget from the
hot-reloaded top-level `recursion_limit` setting. A valid request-level value

31
backend/app/gateway/app.py
View File

@ -392,6 +392,33 @@ async def _shutdown_startup_trash_sweep(app: FastAPI) -> None:
logger.exception("Startup trash sweep failed during shutdown")
@asynccontextmanager
async def _runtime_with_mcp_pool_shutdown(app: FastAPI, startup_config: AppConfig) -> AsyncGenerator[None, None]:
"""Close pooled MCP transports after runtime producers have stopped."""
try:
async with langgraph_runtime(app, startup_config):
yield
finally:
# RunManager drains active graph tasks when langgraph_runtime exits.
# Those tasks can still acquire new MCP sessions during earlier
# shutdown hooks, so closing the pool inside the runtime leaves fresh
# owner tasks and transports alive after the only close pass.
try:
from deerflow.mcp.session_pool import get_session_pool
await asyncio.wait_for(
get_session_pool().close_all(),
timeout=_SHUTDOWN_HOOK_TIMEOUT_SECONDS,
)
except TimeoutError:
logger.warning(
"MCP session pool shutdown exceeded %.1fs; proceeding with worker exit.",
_SHUTDOWN_HOOK_TIMEOUT_SECONDS,
)
except Exception:
logger.exception("Failed to close MCP sessions")
@asynccontextmanager
async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
"""Application lifespan handler."""
@ -499,7 +526,7 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
logger.warning("Upload staging file cleanup skipped", exc_info=True)
# Initialize LangGraph runtime components (StreamBridge, RunManager, checkpointer, store)
async with langgraph_runtime(app, startup_config):
async with _runtime_with_mcp_pool_shutdown(app, startup_config):
logger.info("LangGraph runtime initialised")
# Check admin bootstrap state and migrate orphan threads after admin exists.
@ -702,6 +729,8 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
set_subagent_batch_submitter(None)
# Browser sessions have their own bounded teardown. MCP sessions close
# after the runtime drains runs, since those runs may still call tools.
try:
from deerflow.community.browser_automation import get_browser_session_manager

161
backend/tests/test_gateway_lifespan_shutdown.py
View File

@ -12,7 +12,7 @@ from __future__ import annotations
import asyncio
import logging
import threading
from contextlib import asynccontextmanager
from contextlib import ExitStack, asynccontextmanager
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock, patch
@ -568,3 +568,162 @@ async def test_lifespan_pins_batch_service_to_app_extensions(monkeypatch):
):
async with lifespan(app):
assert app.state.subagent_batch_service._extensions is snapshot
def _gateway_lifespan_patches(startup_config, *, pool=None, browser_manager=None):
"""Common patch set for driving the Gateway lifespan in a unit test."""
channel_service = MagicMock()
channel_service.get_status.return_value = {}
patches = [
patch("app.gateway.app.get_app_config", return_value=startup_config),
patch("app.gateway.app.get_gateway_config", return_value=MagicMock(host="x", port=0)),
patch("app.gateway.app.langgraph_runtime", _noop_langgraph_runtime),
patch("app.gateway.app.auth.close_oidc_service", AsyncMock()),
patch("app.channels.service.start_channel_service", AsyncMock(return_value=channel_service)),
patch("app.channels.service.stop_channel_service", AsyncMock()),
patch("deerflow.skills.projection.ensure_public_skill_projection"),
patch("deerflow.agents.memory.get_memory_manager", return_value=MagicMock()),
]
if pool is not None:
patches.append(patch("deerflow.mcp.session_pool.get_session_pool", return_value=pool))
if browser_manager is not None:
patches.append(
patch(
"deerflow.community.browser_automation.get_browser_session_manager",
return_value=browser_manager,
)
)
return patches
def test_lifespan_closes_pooled_mcp_sessions_on_shutdown():
"""Pooled MCP sessions must be closed while the worker is still shutting down.
Each pooled session owns a live transport (a stdio subprocess, or an
SSE/HTTP connection) held by a dedicated owner task. Nothing else can reach
that task once the event loop stops, so lifespan shutdown is the only place
its ``__aexit__`` can run. The browser session manager is closed here for
the same reason; the MCP pool used to be skipped.
"""
from app.gateway.app import lifespan
async def scenario():
app = FastAPI()
startup_config = MagicMock()
startup_config.log_level = "INFO"
startup_config.memory.enabled = False
startup_config.memory.shutdown_flush_timeout_seconds = 5.0
pool = MagicMock()
pool.close_all = AsyncMock()
with ExitStack() as stack:
for patcher in _gateway_lifespan_patches(startup_config, pool=pool):
stack.enter_context(patcher)
async with lifespan(app):
pass
pool.close_all.assert_awaited_once()
asyncio.run(scenario())
def test_lifespan_continues_when_mcp_close_fails():
"""A failing MCP close must not abort the remaining shutdown hooks.
Shutdown is best-effort by design: every hook is isolated so one broken
teardown cannot strand the others.
"""
from app.gateway.app import lifespan
async def scenario():
app = FastAPI()
startup_config = MagicMock()
startup_config.log_level = "INFO"
startup_config.memory.enabled = False
startup_config.memory.shutdown_flush_timeout_seconds = 5.0
pool = MagicMock()
pool.close_all = AsyncMock(side_effect=RuntimeError("close failed"))
browser_manager = MagicMock()
browser_manager.close_all_sessions = AsyncMock(return_value=0)
with ExitStack() as stack:
for patcher in _gateway_lifespan_patches(startup_config, pool=pool, browser_manager=browser_manager):
stack.enter_context(patcher)
async with lifespan(app):
pass
pool.close_all.assert_awaited_once()
browser_manager.close_all_sessions.assert_awaited_once()
asyncio.run(scenario())
def test_lifespan_closes_mcp_sessions_created_during_run_drain():
"""An active run can acquire a session after other shutdown hooks begin."""
from app.gateway.app import lifespan
from deerflow.mcp.session_pool import MCPSessionPool
from deerflow.runtime import RunManager, RunStatus
async def scenario():
app = FastAPI()
startup_config = MagicMock()
startup_config.log_level = "INFO"
startup_config.memory.enabled = False
startup_config.memory.shutdown_flush_timeout_seconds = 5.0
pool = MCPSessionPool()
run_manager = RunManager()
record = await run_manager.create("mcp-shutdown-race")
await run_manager.set_status(record.run_id, RunStatus.running)
resume = asyncio.Event()
session_created = asyncio.Event()
transport_closed = asyncio.Event()
class SessionContext:
async def __aenter__(self):
self.owner = asyncio.current_task()
return SimpleNamespace(initialize=AsyncMock())
async def __aexit__(self, *_args):
assert asyncio.current_task() is self.owner
transport_closed.set()
async def active_run():
await resume.wait()
await pool.get_session("server", "thread", {"transport": "stdio", "command": "unused"})
session_created.set()
await asyncio.Event().wait()
@asynccontextmanager
async def runtime(_app, _config):
record.task = asyncio.create_task(active_run())
try:
yield
finally:
await run_manager.shutdown(timeout=1.0)
async def shutdown_memory_backend(**_kwargs):
# This hook runs before langgraph_runtime drains the active run.
resume.set()
await asyncio.wait_for(session_created.wait(), timeout=5)
try:
with ExitStack() as stack:
for patcher in _gateway_lifespan_patches(startup_config, pool=pool):
stack.enter_context(patcher)
stack.enter_context(patch("app.gateway.app.langgraph_runtime", runtime))
stack.enter_context(patch("app.gateway.app._shutdown_memory_backend", shutdown_memory_backend))
stack.enter_context(patch("langchain_mcp_adapters.sessions.create_session", side_effect=lambda _connection: SessionContext()))
async with lifespan(app):
pass
assert record.task is not None and record.task.done()
assert transport_closed.is_set()
assert not pool._entries
finally:
resume.set()
if record.task is not None and not record.task.done():
record.task.cancel()
await asyncio.gather(record.task, return_exceptions=True)
await pool.close_all()
asyncio.run(scenario())