diff --git a/backend/app/gateway/AGENTS.md b/backend/app/gateway/AGENTS.md index f00ab9e27f..822b1b64fb 100644 --- a/backend/app/gateway/AGENTS.md +++ b/backend/app/gateway/AGENTS.md @@ -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 diff --git a/backend/app/gateway/app.py b/backend/app/gateway/app.py index 65d85a772b..f748c75bff 100644 --- a/backend/app/gateway/app.py +++ b/backend/app/gateway/app.py @@ -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 diff --git a/backend/tests/test_gateway_lifespan_shutdown.py b/backend/tests/test_gateway_lifespan_shutdown.py index f2a29dad49..2316bcf68e 100644 --- a/backend/tests/test_gateway_lifespan_shutdown.py +++ b/backend/tests/test_gateway_lifespan_shutdown.py @@ -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())