fix(mcp): preserve pooled stdio sessions after task timeouts (#5027)

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:

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)``."""

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:

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,

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."""

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