From 9c1dd111602b93396caaada34391df4fee1a4234 Mon Sep 17 00:00:00 2001 From: ChaseMoon <143682773+zhuiyueya@users.noreply.github.com> Date: Fri, 28 Aug 2026 10:37:01 +0800 Subject: [PATCH] fix(mcp): preserve pooled stdio sessions after task timeouts (#5027) --- .../packages/harness/deerflow/mcp/AGENTS.md | 2 +- .../harness/deerflow/mcp/session_pool.py | 68 +++++++ .../harness/deerflow/mcp/task_tool_caller.py | 48 +++-- .../packages/harness/deerflow/mcp/tools.py | 47 +---- backend/tests/test_mcp_session_pool.py | 42 +++- backend/tests/test_mcp_task_tool_caller.py | 181 +++++++++++++++++- 6 files changed, 318 insertions(+), 70 deletions(-) diff --git a/backend/packages/harness/deerflow/mcp/AGENTS.md b/backend/packages/harness/deerflow/mcp/AGENTS.md index 896c299a2..f65b975e4 100644 --- a/backend/packages/harness/deerflow/mcp/AGENTS.md +++ b/backend/packages/harness/deerflow/mcp/AGENTS.md @@ -23,7 +23,7 @@ references `tool_search` when a hinted MCP tool is currently deferred; do not add a parallel routing middleware for PR1-style preference hints. - **Stdio file outputs**: Persistent stdio sessions are scoped by `user_id:thread_id`. For stdio transports only, DeerFlow pins the subprocess default `cwd` to the thread workspace and `TMPDIR`/`TMP`/`TEMP` to `workspace/.mcp/tmp/`, unless the operator explicitly configured `cwd` or temp env values. `.mcp` is a DeerFlow-owned internal namespace: its temporary/debug files remain addressable when returned by a tool but are excluded from run workspace-change summaries — by directory name at any depth, consistent with the other reserved names in `EXCLUDED_DIR_NAMES` (`.git`, `node_modules`, …) and robust if a server ever creates a relative `.mcp` from a different cwd. Both launch paths pin it at the workspace root today. SSE/HTTP transports skip this filesystem prep entirely. -- **Stdio disconnect recovery**: An ordinary Agent tool call that receives the MCP SDK's explicit `Connection closed` error or an AnyIO closed-stream error evicts only that `(server_name, user_id:thread_id)` session when the registered entry is still the same `ClientSession` that failed. A late error from an old concurrent call cannot evict its replacement or a new in-flight creation. The failing call still surfaces its original error and is never replayed automatically; a later Agent retry creates a fresh subprocess/session. Protocol timeouts, normal `isError=true` tool results, and interceptor failures do not evict a healthy stateful session. +- **Stdio disconnect recovery**: Ordinary Agent tool calls and durable task submit/status/cancel calls that receive the MCP SDK's explicit `Connection closed` error or an AnyIO closed-stream error evict only that `(server_name, user_id:thread_id)` session when the registered entry is still the same `ClientSession` that failed. A late error from an old concurrent call cannot evict its replacement or a new in-flight creation. The failing call still surfaces its original error and is never replayed automatically; a later retry creates a fresh subprocess/session. Protocol timeouts, normal `isError=true` tool results, and interceptor failures do not evict a healthy stateful session. - **Stdio path translation**: MCP-returned local file references are not copied. If a `ResourceLink` or conservative free-text path resolves to an existing file inside the thread's mounted user-data tree, it is translated deterministically to `/mnt/user-data/...`; paths outside that tree remain unchanged. - **Runtime updates**: Gateway API saves to extensions_config.json; the Gateway-embedded runtime detects changes via the resolved-path + content-signature check above, so multi-worker / stale-mtime deployments still pick up an added/removed MCP server without a restart (`PUT /api/mcp/config` keeps whole-payload validation, while `PATCH /api/mcp/config` changes only one server's `enabled` field, normalizes the same `type`/MCP-spec `transport` alias as the runtime config model, and validates the target only when enabling it; either endpoint's reset clears the cache only in its own worker). MCP, skill, and embedded-client writers hold the process-local `extensions_config_write_lock` plus the sidecar advisory `extensions_config_file_lock` for the complete read-modify-write/reload cycle, then share `atomic_write_extensions_config()`, which writes and fsyncs a same-directory temporary file before `os.replace()` and preserves an existing file's mode and symlink target; failed serialization or replacement leaves the prior config intact and cleans up the temporary file. - **Stdio launch policy at the HTTP boundary** (`routers/mcp.py::_validate_mcp_update_request`, shared by `PUT` and the enable branch of `PATCH`): a config file may express anything, but the API is untrusted input, so an API-registered stdio server must (a) name a bare executable from the allowlist — `_DEFAULT_MCP_STDIO_COMMAND_ALLOWLIST` = `{npx, uvx}`, extended by `DEER_FLOW_MCP_STDIO_COMMAND_ALLOWLIST`, with path separators, whitespace, and shell metacharacters rejected in `command`; (b) carry no `args` flag in `_ARBITRARY_EXEC_ARGS`; and (c) set no `env` name in `_CODE_INJECTING_ENV_VARS`. Checks (b) and (c) exist because the command check alone names a binary without constraining what that binary runs. The `env` denylist applies to **every** allowlisted command, and both denylists match `--flag=value` as well as `--flag value`. The `args` denylist's **scope depends on the command**, because where a launcher stops parsing its own flags is what decides whether a token is an exec flag at all: diff --git a/backend/packages/harness/deerflow/mcp/session_pool.py b/backend/packages/harness/deerflow/mcp/session_pool.py index d707f1fc0..64fa091ce 100644 --- a/backend/packages/harness/deerflow/mcp/session_pool.py +++ b/backend/packages/harness/deerflow/mcp/session_pool.py @@ -39,10 +39,78 @@ import threading from collections import OrderedDict from typing import Any +import anyio from mcp import ClientSession +from mcp.shared.exceptions import McpError +from mcp.types import CONNECTION_CLOSED logger = logging.getLogger(__name__) +_MCP_CLOSED_STREAM_ERRORS = ( + anyio.ClosedResourceError, + anyio.BrokenResourceError, + anyio.EndOfStream, +) + + +def _is_mcp_transport_disconnect(error: Exception) -> bool: + if isinstance(error, _MCP_CLOSED_STREAM_ERRORS): + return True + return isinstance(error, McpError) and error.error.code == CONNECTION_CLOSED and error.error.message == "Connection closed" + + +async def _finish_session_cleanup(cleanup: asyncio.Task[Any], server_name: str) -> bool: + """Wait for cleanup despite cancellation and report whether it was requested.""" + cancelled = False + while not cleanup.done(): + try: + await asyncio.shield(cleanup) + except asyncio.CancelledError: + cancelled = True + except Exception: + logger.warning( + "Failed to close disconnected MCP session for server '%s'", + server_name, + exc_info=True, + ) + return cancelled + + if cleanup.cancelled(): + logger.warning( + "Disconnected MCP session cleanup was cancelled for server '%s'", + server_name, + ) + else: + cleanup_error = cleanup.exception() + if cleanup_error is not None: + logger.warning( + "Failed to close disconnected MCP session for server '%s'", + server_name, + exc_info=(type(cleanup_error), cleanup_error, cleanup_error.__traceback__), + ) + return cancelled + + +async def call_pooled_session_tool( + session: ClientSession, + pool: MCPSessionPool, + *, + server_name: str, + scope_key: str, + tool_name: str, + arguments: dict[str, Any], + call_kwargs: dict[str, Any], +) -> Any: + """Call a pooled session and evict it only after an explicit disconnect.""" + try: + return await session.call_tool(tool_name, arguments, **call_kwargs) + except Exception as error: + if _is_mcp_transport_disconnect(error): + cleanup = asyncio.create_task(pool.close_session_if_current(server_name, scope_key, session)) + if await _finish_session_cleanup(cleanup, server_name): + raise asyncio.CancelledError + raise + class MCPSessionPool: """Manages persistent MCP sessions scoped by ``(server_name, scope_key)``.""" diff --git a/backend/packages/harness/deerflow/mcp/task_tool_caller.py b/backend/packages/harness/deerflow/mcp/task_tool_caller.py index 3643061d0..fc4420742 100644 --- a/backend/packages/harness/deerflow/mcp/task_tool_caller.py +++ b/backend/packages/harness/deerflow/mcp/task_tool_caller.py @@ -16,7 +16,7 @@ from deerflow.mcp.context_headers import build_context_headers_interceptor from deerflow.mcp.headers import apply_header_overrides from deerflow.mcp.interceptors import build_mcp_tool_interceptors from deerflow.mcp.oauth import OAuthTokenManager, build_oauth_tool_interceptor -from deerflow.mcp.session_pool import get_session_pool +from deerflow.mcp.session_pool import MCPSessionPool, call_pooled_session_tool, get_session_pool logger = logging.getLogger(__name__) @@ -133,23 +133,19 @@ class McpTaskToolCaller: raise else: session = await pool.get_session(server_name, scope_key, connection) - try: - return await self._invoke( - session=session, - connection=connection, - server_name=server_name, - tool_name=tool_name, - arguments=arguments, - timeout_seconds=server_config.tool_call_timeout, - session_init_timeout_seconds=None, - persistent_session=True, - interceptors=interceptors, - ) - except Exception: - # A dead pooled subprocess must not poison every later status - # poll. The next retry recreates this exact scoped session. - await pool.close_session(server_name, scope_key) - raise + return await self._invoke( + session=session, + pool=pool, + scope_key=scope_key, + connection=connection, + server_name=server_name, + tool_name=tool_name, + arguments=arguments, + timeout_seconds=server_config.tool_call_timeout, + session_init_timeout_seconds=None, + persistent_session=True, + interceptors=interceptors, + ) authorization = await self._oauth_token_manager.get_authorization_header(server_name) if authorization: @@ -159,6 +155,8 @@ class McpTaskToolCaller: ) return await self._invoke( session=None, + pool=None, + scope_key=scope_key, connection=connection, server_name=server_name, tool_name=tool_name, @@ -173,6 +171,8 @@ class McpTaskToolCaller: self, *, session: Any | None, + pool: MCPSessionPool | None, + scope_key: str, connection: dict[str, Any], server_name: str, tool_name: str, @@ -191,7 +191,7 @@ class McpTaskToolCaller: call_kwargs["read_timeout_seconds"] = timedelta(seconds=timeout_seconds) if persistent_session: - assert session is not None + assert session is not None and pool is not None if request.headers: if isinstance(request.headers, Mapping): call_kwargs["meta"] = {"headers": dict(request.headers)} @@ -200,7 +200,15 @@ class McpTaskToolCaller: "Ignoring MCP interceptor headers with unsupported type: %s", type(request.headers).__name__, ) - return await session.call_tool(request.name, request.args, **call_kwargs) + return await call_pooled_session_tool( + session, + pool, + server_name=server_name, + scope_key=scope_key, + tool_name=request.name, + arguments=request.args, + call_kwargs=call_kwargs, + ) effective_connection = dict(connection) if request.headers: diff --git a/backend/packages/harness/deerflow/mcp/tools.py b/backend/packages/harness/deerflow/mcp/tools.py index b35eb7f83..0d488921b 100644 --- a/backend/packages/harness/deerflow/mcp/tools.py +++ b/backend/packages/harness/deerflow/mcp/tools.py @@ -11,12 +11,8 @@ from pathlib import Path from typing import Any from urllib.parse import unquote, urlparse -import anyio from langchain_core.tools import BaseTool, StructuredTool from langgraph.config import get_config -from mcp import ClientSession -from mcp.shared.exceptions import McpError -from mcp.types import CONNECTION_CLOSED from deerflow.config.extensions_config import ExtensionsConfig, McpServerConfig, resolve_effective_mcp_routing from deerflow.config.paths import VIRTUAL_PATH_PREFIX, Paths, get_paths @@ -25,7 +21,7 @@ from deerflow.mcp.client import build_servers_config from deerflow.mcp.headers import apply_header_overrides from deerflow.mcp.interceptors import build_mcp_tool_interceptors, compose_tool_interceptors from deerflow.mcp.oauth import build_oauth_tool_interceptor, get_initial_oauth_headers -from deerflow.mcp.session_pool import MCPSessionPool, get_session_pool +from deerflow.mcp.session_pool import call_pooled_session_tool, get_session_pool from deerflow.mcp.tasks import ORDINARY_MCP_TASK_DRIVER, TaskSubmitRequest from deerflow.mcp.tasks.runtime import ( McpTaskConfigurationError, @@ -65,43 +61,6 @@ _TEXT_PATH_TRAILING_CHARS = ".,;:!?)]}>\"'`" _FILE_SNAPSHOT = dict[Path, tuple[int, int]] -_MCP_CLOSED_STREAM_ERRORS = ( - anyio.ClosedResourceError, - anyio.BrokenResourceError, - anyio.EndOfStream, -) - - -def _is_mcp_transport_disconnect(error: Exception) -> bool: - if isinstance(error, _MCP_CLOSED_STREAM_ERRORS): - return True - return isinstance(error, McpError) and error.error.code == CONNECTION_CLOSED and error.error.message == "Connection closed" - - -async def _call_pooled_session_tool( - session: ClientSession, - pool: MCPSessionPool, - *, - server_name: str, - scope_key: str, - tool_name: str, - arguments: dict[str, Any], - call_kwargs: dict[str, Any], -) -> Any: - try: - return await session.call_tool(tool_name, arguments, **call_kwargs) - except Exception as error: - if _is_mcp_transport_disconnect(error): - try: - await pool.close_session_if_current(server_name, scope_key, session) - except Exception: - logger.warning( - "Failed to close disconnected MCP session for server '%s'", - server_name, - exc_info=True, - ) - raise - def _local_path_from_uri(uri: str, *, base_dir: Path | None = None) -> Path | None: """Return an absolute local filesystem ``Path`` if *uri* points to a local @@ -598,7 +557,7 @@ def _make_session_pool_tool( kwargs["meta"] = {"headers": dict(request.headers)} else: logger.warning("Ignoring MCP interceptor headers with unsupported type: %s", type(request.headers).__name__) - return await _call_pooled_session_tool( + return await call_pooled_session_tool( session, pool, server_name=server_name, @@ -618,7 +577,7 @@ def _make_session_pool_tool( ) call_tool_result = await handler(request) else: - call_tool_result = await _call_pooled_session_tool( + call_tool_result = await call_pooled_session_tool( session, pool, server_name=server_name, diff --git a/backend/tests/test_mcp_session_pool.py b/backend/tests/test_mcp_session_pool.py index 4fb5b29e6..b7337c7a5 100644 --- a/backend/tests/test_mcp_session_pool.py +++ b/backend/tests/test_mcp_session_pool.py @@ -12,7 +12,7 @@ import pytest from mcp.shared.exceptions import McpError from mcp.types import CONNECTION_CLOSED, CallToolResult, ErrorData, TextContent -from deerflow.mcp.session_pool import MCPSessionPool, get_session_pool, reset_session_pool +from deerflow.mcp.session_pool import MCPSessionPool, call_pooled_session_tool, get_session_pool, reset_session_pool @pytest.fixture(autouse=True) @@ -448,6 +448,46 @@ async def test_session_pool_tool_preserves_disconnect_error_when_eviction_fails( pool.close_session_if_current.assert_awaited_once_with("srv", "test-user-autouse:default", session) +@pytest.mark.asyncio +async def test_session_pool_disconnect_cleanup_survives_caller_cancellation(): + """A cancelled caller cannot interrupt disconnected-session teardown.""" + cleanup_started = asyncio.Event() + release_cleanup = asyncio.Event() + cleanup_finished = asyncio.Event() + + async def close_session_if_current(*_args): + cleanup_started.set() + await release_cleanup.wait() + cleanup_finished.set() + return True + + session = AsyncMock() + session.call_tool = AsyncMock(side_effect=anyio.ClosedResourceError()) + pool = MagicMock() + pool.close_session_if_current = close_session_if_current + + call = asyncio.create_task( + call_pooled_session_tool( + session, + pool, + server_name="srv", + scope_key="scope", + tool_name="act", + arguments={}, + call_kwargs={}, + ) + ) + await cleanup_started.wait() + call.cancel() + await asyncio.sleep(0) + call.cancel() + release_cleanup.set() + + with pytest.raises(asyncio.CancelledError): + await call + assert cleanup_finished.is_set() + + @pytest.mark.asyncio async def test_session_pool_tool_keeps_session_after_tool_error_result(tmp_path): """An MCP tool-level error is a valid response from a live session.""" diff --git a/backend/tests/test_mcp_task_tool_caller.py b/backend/tests/test_mcp_task_tool_caller.py index 064aeee9d..e0f8dad64 100644 --- a/backend/tests/test_mcp_task_tool_caller.py +++ b/backend/tests/test_mcp_task_tool_caller.py @@ -1,4 +1,5 @@ import asyncio +import sys from collections.abc import Coroutine from contextlib import suppress from datetime import timedelta @@ -6,9 +7,14 @@ from types import SimpleNamespace from typing import Any from unittest.mock import AsyncMock, MagicMock, patch +import anyio import pytest +from mcp.shared.exceptions import McpError +from mcp.types import CONNECTION_CLOSED, ErrorData from deerflow.config.extensions_config import ExtensionsConfig +from deerflow.config.paths import Paths +from deerflow.mcp.session_pool import MCPSessionPool from deerflow.mcp.task_tool_caller import McpTaskToolCaller, mcp_task_session_scope_key @@ -111,11 +117,21 @@ async def test_stdio_task_call_reuses_exact_scope_and_raw_tool_name() -> None: @pytest.mark.asyncio -async def test_broken_stdio_task_session_is_evicted_for_next_poll_reconnect() -> None: - session = SimpleNamespace(call_tool=AsyncMock(side_effect=ConnectionError("disconnected"))) +@pytest.mark.parametrize( + "disconnect_error", + [ + anyio.ClosedResourceError(), + anyio.BrokenResourceError(), + anyio.EndOfStream(), + McpError(ErrorData(code=CONNECTION_CLOSED, message="Connection closed")), + ], +) +async def test_broken_stdio_task_session_is_evicted_for_next_poll_reconnect(disconnect_error: Exception) -> None: + session = SimpleNamespace(call_tool=AsyncMock(side_effect=disconnect_error)) pool = MagicMock() pool.get_session = AsyncMock(return_value=session) pool.close_session = AsyncMock() + pool.close_session_if_current = AsyncMock() caller = McpTaskToolCaller(_config()) with ( @@ -124,7 +140,7 @@ async def test_broken_stdio_task_session_is_evicted_for_next_poll_reconnect() -> "deerflow.mcp.task_tool_caller._prepare_stdio_connection", return_value={"transport": "stdio", "command": "report-mcp"}, ), - pytest.raises(ConnectionError, match="disconnected"), + pytest.raises(type(disconnect_error)), ): await caller.call_tool( server_name="reports", @@ -134,7 +150,164 @@ async def test_broken_stdio_task_session_is_evicted_for_next_poll_reconnect() -> thread_id="thread-1", ) - pool.close_session.assert_awaited_once_with("reports", "user-1:thread-1") + pool.close_session_if_current.assert_awaited_once_with( + "reports", + "user-1:thread-1", + session, + ) + pool.close_session.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_stdio_task_timeout_keeps_healthy_stateful_session() -> None: + timeout_error = McpError(ErrorData(code=408, message="request timed out")) + session = SimpleNamespace(call_tool=AsyncMock(side_effect=timeout_error)) + pool = MagicMock() + pool.get_session = AsyncMock(return_value=session) + pool.close_session = AsyncMock() + pool.close_session_if_current = AsyncMock() + caller = McpTaskToolCaller(_config()) + + with ( + patch("deerflow.mcp.task_tool_caller.get_session_pool", return_value=pool), + patch( + "deerflow.mcp.task_tool_caller._prepare_stdio_connection", + return_value={"transport": "stdio", "command": "report-mcp"}, + ), + pytest.raises(McpError, match="request timed out"), + ): + await caller.call_tool( + server_name="reports", + tool_name="status_report", + arguments={"task_id": "remote-1"}, + user_id="user-1", + thread_id="thread-1", + ) + + pool.close_session_if_current.assert_not_awaited() + pool.close_session.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_stdio_task_interceptor_failure_keeps_healthy_session() -> None: + session = SimpleNamespace(call_tool=AsyncMock()) + pool = MagicMock() + pool.get_session = AsyncMock(return_value=session) + pool.close_session = AsyncMock() + pool.close_session_if_current = AsyncMock() + caller = McpTaskToolCaller(_config()) + + async def reject_call(_request, _handler): + raise RuntimeError("interceptor rejected call") + + caller._interceptors = [reject_call] + + with ( + patch("deerflow.mcp.task_tool_caller.get_session_pool", return_value=pool), + patch( + "deerflow.mcp.task_tool_caller._prepare_stdio_connection", + return_value={"transport": "stdio", "command": "report-mcp"}, + ), + pytest.raises(RuntimeError, match="interceptor rejected call"), + ): + await caller.call_tool( + server_name="reports", + tool_name="status_report", + arguments={"task_id": "remote-1"}, + user_id="user-1", + thread_id="thread-1", + ) + + session.call_tool.assert_not_awaited() + pool.close_session_if_current.assert_not_awaited() + pool.close_session.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_stdio_task_timeout_preserves_real_stateful_session(tmp_path) -> None: + server = """ +import asyncio +import os + +from mcp.server.fastmcp import FastMCP + +mcp = FastMCP("slow-status") +tasks = {} + + +@mcp.tool() +def submit_report() -> dict[str, object]: + tasks["remote-1"] = 0 + return {"task_id": "remote-1", "status": "running", "pid": os.getpid()} + + +@mcp.tool() +async def status_report(task_id: str) -> dict[str, object]: + if task_id not in tasks: + return {"task_id": task_id, "status": "failed", "error_code": "task_not_found"} + tasks[task_id] += 1 + if tasks[task_id] == 1: + await asyncio.sleep(0.2) + return { + "task_id": task_id, + "status": "completed", + "pid": os.getpid(), + "status_calls": tasks[task_id], + } + + +mcp.run(transport="stdio") +""" + config = _config() + server_config = config.mcp_servers["reports"] + server_config.command = sys.executable + server_config.args = ["-c", server] + server_config.tool_call_timeout = 1.0 + caller = McpTaskToolCaller(config) + pool = MCPSessionPool() + + try: + with ( + patch("deerflow.mcp.task_tool_caller.get_paths", return_value=Paths(tmp_path)), + patch("deerflow.mcp.task_tool_caller.get_session_pool", return_value=pool), + ): + submitted = await caller.call_tool( + server_name="reports", + tool_name="submit_report", + arguments={}, + user_id="user-1", + thread_id="thread-1", + ) + + server_config.tool_call_timeout = 0.05 + with pytest.raises(McpError, match="Timed out while waiting") as exc_info: + await caller.call_tool( + server_name="reports", + tool_name="status_report", + arguments={"task_id": submitted.structuredContent["task_id"]}, + user_id="user-1", + thread_id="thread-1", + ) + assert exc_info.value.error.code == 408 + + await asyncio.sleep(0.25) + server_config.tool_call_timeout = 1.0 + recovered = await caller.call_tool( + server_name="reports", + tool_name="status_report", + arguments={"task_id": submitted.structuredContent["task_id"]}, + user_id="user-1", + thread_id="thread-1", + ) + finally: + await pool.close_all() + + assert recovered.structuredContent == { + "task_id": "remote-1", + "status": "completed", + "pid": submitted.structuredContent["pid"], + "status_calls": 2, + } @pytest.mark.asyncio