diff --git a/backend/packages/harness/deerflow/mcp/AGENTS.md b/backend/packages/harness/deerflow/mcp/AGENTS.md index 275d94169..876c9172e 100644 --- a/backend/packages/harness/deerflow/mcp/AGENTS.md +++ b/backend/packages/harness/deerflow/mcp/AGENTS.md @@ -20,6 +20,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 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 333505fdb..d707f1fc0 100644 --- a/backend/packages/harness/deerflow/mcp/session_pool.py +++ b/backend/packages/harness/deerflow/mcp/session_pool.py @@ -364,6 +364,23 @@ class MCPSessionPool: loop, _ready, task, close_evt = inflight await self._shutdown_entry(loop, task, close_evt, cancel=True) + async def close_session_if_current( + self, + server_name: str, + scope_key: str, + session: ClientSession, + ) -> bool: + """Close *session* only if it is still the registered entry for the key.""" + key = (server_name, scope_key) + with self._lock: + entry = self._entries.get(key) + if entry is None or entry[0] is not session: + return False + self._entries.pop(key) + _session, loop, task, close_evt = entry + await self._shutdown_entry(loop, task, close_evt) + return True + async def close_server(self, server_name: str) -> None: """Close all sessions for a given server.""" with self._lock: diff --git a/backend/packages/harness/deerflow/mcp/tools.py b/backend/packages/harness/deerflow/mcp/tools.py index 73de63161..30d0de18e 100644 --- a/backend/packages/harness/deerflow/mcp/tools.py +++ b/backend/packages/harness/deerflow/mcp/tools.py @@ -11,8 +11,12 @@ 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 @@ -20,7 +24,7 @@ from deerflow.constants import DEFAULT_MCP_SESSION_INIT_TIMEOUT, MCP_TMP_SUBDIR from deerflow.mcp.client import build_servers_config 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 get_session_pool +from deerflow.mcp.session_pool import MCPSessionPool, get_session_pool from deerflow.mcp.tasks import ORDINARY_MCP_TASK_DRIVER, TaskSubmitRequest from deerflow.mcp.tasks.runtime import ( McpTaskConfigurationError, @@ -60,6 +64,43 @@ _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 @@ -556,10 +597,14 @@ 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 session.call_tool( - request.name, - request.args, - **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=kwargs, ) handler = compose_tool_interceptors(tool_interceptors, base_handler) @@ -572,10 +617,14 @@ def _make_session_pool_tool( ) call_tool_result = await handler(request) else: - call_tool_result = await session.call_tool( - original_name, - arguments, - **call_kwargs, + call_tool_result = await _call_pooled_session_tool( + session, + pool, + server_name=server_name, + scope_key=scope_key, + tool_name=original_name, + arguments=arguments, + call_kwargs=call_kwargs, ) # The after-call snapshot diff only feeds bare-filename correlation in diff --git a/backend/tests/test_mcp_session_pool.py b/backend/tests/test_mcp_session_pool.py index 01fa986dc..4fb5b29e6 100644 --- a/backend/tests/test_mcp_session_pool.py +++ b/backend/tests/test_mcp_session_pool.py @@ -3,10 +3,14 @@ import asyncio import logging import stat +import sys import threading from unittest.mock import AsyncMock, MagicMock, patch +import anyio 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 @@ -254,6 +258,317 @@ def test_reset_session_pool(): # --------------------------------------------------------------------------- +def _make_test_pool_tool(*, pool, call_tool, tool_interceptors=None): + from langchain_core.tools import StructuredTool + from pydantic import BaseModel, Field + + from deerflow.mcp.tools import _make_session_pool_tool + + class Args(BaseModel): + value: int = Field(..., description="value") + + original_tool = StructuredTool( + name="srv_act", + description="test", + args_schema=Args, + coroutine=AsyncMock(), + response_format="content_and_artifact", + ) + session = AsyncMock() + if isinstance(call_tool, BaseException): + session.call_tool = AsyncMock(side_effect=call_tool) + else: + session.call_tool = AsyncMock(return_value=call_tool) + pool.get_session = AsyncMock(return_value=session) + pool.close_session_if_current = AsyncMock() + + with patch("deerflow.mcp.tools.get_session_pool", return_value=pool): + wrapped = _make_session_pool_tool( + original_tool, + "srv", + {"transport": "stdio", "command": "x", "args": []}, + tool_interceptors=tool_interceptors, + ) + return wrapped, session + + +@pytest.mark.asyncio +async def test_session_pool_tool_reconnects_after_real_stdio_process_disconnect(tmp_path): + """A dead stdio subprocess must not poison later calls in the same scope.""" + from langchain_core.tools import StructuredTool + from pydantic import BaseModel + + from deerflow.config.paths import Paths + from deerflow.mcp.tools import _make_session_pool_tool + + server = """ +import os +import sys +from pathlib import Path + +from mcp.server.fastmcp import FastMCP + +marker = Path(sys.argv[1]) +mcp = FastMCP("crash-once") + +@mcp.tool() +def crash_once() -> str: + if not marker.exists(): + marker.write_text("crashed") + os._exit(17) + return "recovered" + +mcp.run(transport="stdio") +""" + + class Args(BaseModel): + pass + + original_tool = StructuredTool( + name="crash_crash_once", + description="crash once", + args_schema=Args, + coroutine=AsyncMock(), + response_format="content_and_artifact", + ) + marker = tmp_path / "crashed" + connection = { + "transport": "stdio", + "command": sys.executable, + "args": ["-c", server, str(marker)], + } + runtime = MagicMock() + runtime.context = {"thread_id": "thread", "user_id": "user"} + runtime.config = {} + + with patch("deerflow.mcp.tools.get_paths", return_value=Paths(tmp_path)): + wrapped = _make_session_pool_tool(original_tool, "crash", connection) + with pytest.raises(McpError, match="Connection closed") as exc_info: + await wrapped.coroutine(runtime=runtime) + + assert exc_info.value.error.code == CONNECTION_CLOSED + assert ("crash", "user:thread") not in get_session_pool()._entries + + content, _artifact = await wrapped.coroutine(runtime=runtime) + + assert content[0]["text"] == "recovered" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "transport_error", + [ + anyio.ClosedResourceError(), + anyio.BrokenResourceError(), + anyio.EndOfStream(), + ], +) +async def test_session_pool_tool_evicts_session_after_transport_disconnect(tmp_path, transport_error): + """Low-level closed-stream signals evict the exact pooled session.""" + from deerflow.config.paths import Paths + + pool = MagicMock() + wrapped, session = _make_test_pool_tool(pool=pool, call_tool=transport_error) + + with ( + patch("deerflow.mcp.tools.get_paths", return_value=Paths(tmp_path)), + pytest.raises(type(transport_error)), + ): + await wrapped.coroutine(value=1) + + pool.close_session_if_current.assert_awaited_once_with("srv", "test-user-autouse:default", session) + + +@pytest.mark.asyncio +async def test_session_pool_tool_evicts_connection_closed_through_interceptor(tmp_path): + """A passthrough interceptor must retain transport-failure recovery.""" + from deerflow.config.paths import Paths + + async def passthrough(request, handler): + return await handler(request) + + error = McpError(ErrorData(code=CONNECTION_CLOSED, message="Connection closed")) + pool = MagicMock() + wrapped, session = _make_test_pool_tool( + pool=pool, + call_tool=error, + tool_interceptors=[passthrough], + ) + + with ( + patch("deerflow.mcp.tools.get_paths", return_value=Paths(tmp_path)), + pytest.raises(McpError, match="Connection closed"), + ): + await wrapped.coroutine(value=1) + + pool.close_session_if_current.assert_awaited_once_with("srv", "test-user-autouse:default", session) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "error", + [ + McpError(ErrorData(code=408, message="request timed out")), + McpError(ErrorData(code=CONNECTION_CLOSED, message="server-specific failure")), + ], +) +async def test_session_pool_tool_keeps_session_after_nonfatal_mcp_error(tmp_path, error): + """Protocol errors such as timeouts do not prove that the session is dead.""" + from deerflow.config.paths import Paths + + pool = MagicMock() + wrapped, _session = _make_test_pool_tool(pool=pool, call_tool=error) + + with ( + patch("deerflow.mcp.tools.get_paths", return_value=Paths(tmp_path)), + pytest.raises(McpError, match=str(error)), + ): + await wrapped.coroutine(value=1) + + pool.close_session_if_current.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_session_pool_tool_preserves_disconnect_error_when_eviction_fails(tmp_path): + """Cleanup failure must not replace the transport error seen by the caller.""" + from deerflow.config.paths import Paths + + error = anyio.ClosedResourceError() + pool = MagicMock() + wrapped, session = _make_test_pool_tool(pool=pool, call_tool=error) + pool.close_session_if_current.side_effect = RuntimeError("cleanup failed") + + with ( + patch("deerflow.mcp.tools.get_paths", return_value=Paths(tmp_path)), + pytest.raises(anyio.ClosedResourceError) as exc_info, + ): + await wrapped.coroutine(value=1) + + assert exc_info.value is error + pool.close_session_if_current.assert_awaited_once_with("srv", "test-user-autouse:default", session) + + +@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.""" + from langchain_core.tools import ToolException + + from deerflow.config.paths import Paths + + result = CallToolResult( + content=[TextContent(type="text", text="invalid input")], + isError=True, + ) + pool = MagicMock() + wrapped, _session = _make_test_pool_tool(pool=pool, call_tool=result) + + with ( + patch("deerflow.mcp.tools.get_paths", return_value=Paths(tmp_path)), + pytest.raises(ToolException, match="invalid input"), + ): + await wrapped.coroutine(value=1) + + pool.close_session_if_current.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_session_pool_tool_keeps_session_after_interceptor_error(tmp_path): + """Interceptor failures happen outside the transport and must not evict it.""" + from deerflow.config.paths import Paths + + async def failing_interceptor(_request, _handler): + raise RuntimeError("interceptor failed") + + pool = MagicMock() + wrapped, session = _make_test_pool_tool( + pool=pool, + call_tool=CallToolResult(content=[], isError=False), + tool_interceptors=[failing_interceptor], + ) + + with ( + patch("deerflow.mcp.tools.get_paths", return_value=Paths(tmp_path)), + pytest.raises(RuntimeError, match="interceptor failed"), + ): + await wrapped.coroutine(value=1) + + session.call_tool.assert_not_awaited() + pool.close_session_if_current.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_late_disconnect_from_old_session_does_not_evict_replacement(tmp_path): + """Concurrent late failures must not close a replacement for the same key.""" + from langchain_core.tools import StructuredTool + from pydantic import BaseModel, Field + + from deerflow.config.paths import Paths + from deerflow.mcp.tools import _make_session_pool_tool + + first_failure = asyncio.Event() + late_failure = asyncio.Event() + both_started = asyncio.Event() + call_count = 0 + + async def old_call_tool(_name, arguments, **_kwargs): + nonlocal call_count + call_count += 1 + if call_count == 2: + both_started.set() + await (first_failure if arguments["value"] == 1 else late_failure).wait() + raise anyio.ClosedResourceError + + old_session = AsyncMock() + old_session.call_tool = AsyncMock(side_effect=old_call_tool) + replacement = AsyncMock() + replacement.call_tool = AsyncMock(return_value=CallToolResult(content=[], isError=False)) + sessions = iter([old_session, replacement]) + + def create_session(*_args, **_kwargs): + session = next(sessions) + context_manager = MagicMock() + context_manager.__aenter__ = AsyncMock(return_value=session) + context_manager.__aexit__ = AsyncMock(return_value=False) + return context_manager + + class Args(BaseModel): + value: int = Field(..., description="value") + + original_tool = StructuredTool( + name="srv_act", + description="test", + args_schema=Args, + coroutine=AsyncMock(), + response_format="content_and_artifact", + ) + pool = get_session_pool() + + with ( + patch("deerflow.mcp.tools.get_paths", return_value=Paths(tmp_path)), + patch("langchain_mcp_adapters.sessions.create_session", side_effect=create_session), + ): + wrapped = _make_session_pool_tool( + original_tool, + "srv", + {"transport": "stdio", "command": "x", "args": []}, + ) + first_call = asyncio.create_task(wrapped.coroutine(value=1)) + late_call = asyncio.create_task(wrapped.coroutine(value=2)) + await asyncio.wait_for(both_started.wait(), timeout=1) + + first_failure.set() + with pytest.raises(anyio.ClosedResourceError): + await first_call + + await wrapped.coroutine(value=3) + late_failure.set() + with pytest.raises(anyio.ClosedResourceError): + await late_call + + assert pool._entries[("srv", "test-user-autouse:default")][0] is replacement + await pool.close_all() + + @pytest.mark.asyncio async def test_session_pool_tool_wrapping(): """The wrapper tool delegates to a pool-managed session."""