fix(mcp): include connection setup in task session timeout (#5733)

* fix(mcp): include connection setup in task session timeout

* fix(mcp): identify remote session initialization timeouts

* style(tests): format MCP timeout assertions

---------

Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
This commit is contained in:
asts 2026-09-23 10:20:51 +08:00 committed by GitHub
parent cb349da014
commit eae70b5e40
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 273 additions and 34 deletions

View File

@ -623,6 +623,7 @@ already received by the browser, without an additional secret-redaction layer.
DeerFlow supports configurable MCP servers and skills to extend its capabilities. DeerFlow supports configurable MCP servers and skills to extend its capabilities.
For HTTP/SSE MCP servers, OAuth token flows are supported (`client_credentials`, `refresh_token`). For HTTP/SSE MCP servers, OAuth token flows are supported (`client_credentials`, `refresh_token`).
For stdio MCP servers, per-tool call timeouts can be configured with `tool_call_timeout`; durable background-task calls honor the same setting for HTTP/SSE servers as well. For stdio MCP servers, per-tool call timeouts can be configured with `tool_call_timeout`; durable background-task calls honor the same setting for HTTP/SSE servers as well.
For HTTP/SSE background-task calls, `session_init_timeout` separately bounds connection setup (including the SSE endpoint event) and MCP initialization together; it stops applying once the tool call begins. Initialization deadline errors identify the server and configured time limit.
Ordinary `task` subagents retain the parent run's captured thread incarnation for MCP calls, including legacy threads, so delegation preserves the same lifecycle scope. Ordinary `task` subagents retain the parent run's captured thread incarnation for MCP calls, including legacy threads, so delegation preserves the same lifecycle scope.
MCP tool names are prefixed with `<server_name>_` by default to prevent collisions across servers. If a server already namespaces its own tools, set `tool_name_prefix: false` on that server in `extensions_config.json` to keep the original names. Disable the prefix only when the resulting names remain unique across all enabled servers. MCP tool names are prefixed with `<server_name>_` by default to prevent collisions across servers. If a server already namespaces its own tools, set `tool_name_prefix: false` on that server in `extensions_config.json` to keep the original names. Disable the prefix only when the resulting names remain unique across all enabled servers.
Signed-in users' notification toggle, default model, conversation mode, and reasoning effort are saved to their account and restored on other browsers or after clearing browser storage. Browser notification permission still needs to be granted on each device. Changes retry after network failures; unsent changes survive a reload in the same tab. Concurrent edits to different fields are preserved; for the same field, the last server write wins. Existing unscoped browser preferences are not uploaded automatically because they have no account owner; reselect those settings once after upgrading. Static demos and auth-disabled development keep browser-local settings. Thread-specific model overrides and other display preferences remain local. Signed-in users' notification toggle, default model, conversation mode, and reasoning effort are saved to their account and restored on other browsers or after clearing browser storage. Browser notification permission still needs to be granted on each device. Changes retry after network failures; unsent changes survive a reload in the same tab. Concurrent edits to different fields are preserved; for the same field, the last server write wins. Existing unscoped browser preferences are not uploaded automatically because they have no account owner; reselect those settings once after upgrading. Static demos and auth-disabled development keep browser-local settings. Thread-specific model overrides and other display preferences remain local.

View File

@ -253,9 +253,12 @@ cannot finish transport cleanup on a loop that has already closed.
Two independent settings bound stdio MCP servers and durable HTTP/SSE task Two independent settings bound stdio MCP servers and durable HTTP/SSE task
calls. `session_init_timeout` covers server bring-up — tool discovery calls. `session_init_timeout` covers server bring-up — tool discovery
(subprocess spawn + `initialize` + `tools/list`) and persistent-session (subprocess spawn + `initialize` + `tools/list`) and persistent-session
initialization — plus ephemeral HTTP/SSE task-session initialization. It initialization — plus ephemeral HTTP/SSE task-session connection setup and
defaults to 60s so a hung server (e.g. `npx` blocked on a package download, or initialization under a single deadline. This includes waiting for an SSE
a server that never answers `initialize`) cannot block agent construction or `endpoint` event. Once initialization succeeds, this deadline is disabled;
the tool call uses its independent `tool_call_timeout`.
The initialization timeout defaults to 60s so a hung server (e.g. `npx` blocked
on a package download, or a server that never answers `initialize`) cannot block agent construction or
the task poller indefinitely. Set it to `null` to disable: the task poller indefinitely. Set it to `null` to disable:
```json ```json

View File

@ -5,6 +5,7 @@
- **Runtime availability boundary**: the installed process-local submitter is the source of truth for durable task-management tool exposure. `mcp_tasks` is startup-only; changing it on disk does not alter the live toolset until the Gateway restarts. - **Runtime availability boundary**: the installed process-local submitter is the source of truth for durable task-management tool exposure. `mcp_tasks` is startup-only; changing it on disk does not alter the live toolset until the Gateway restarts.
- **Long-running ordinary task driver**: `extensions_config.json -> mcpServers.<server>.task_toolsets` binds exact raw submit/status/cancel names; one raw tool may occupy only one role across that server's groups. `mcp/tools.py` hides status/cancel and replaces submit with a wrapper that returns only the local task ID after persistence. `ordinary.py` reads only MCP `structuredContent`, maps remote `running` to `working`, and treats `error_code=task_not_found` or malformed structured output as permanent failure. A status call with `isError=true` is a retryable call failure: the first text content block is retained as a bounded diagnostic, while a permanent remote-task outcome must arrive in a normal result with structured `status=failed`. `task_tool_caller.py` restores the same canonical `(server_name, versioned user/thread/incarnation scope)` stdio session as ordinary calls; non-NULL incarnations use an unambiguous versioned encoding, while an explicitly captured legacy NULL retains the pre-activation `user_id:thread_id` scope. A runtime with a missing/invalid incarnation key fails closed. HTTP/SSE calls remain ephemeral, apply `session_init_timeout` to initialization and `tool_call_timeout` to task calls, and support server-level OAuth refresh outside an Agent run. `McpTaskService` exponentially backs off transient status/cancel errors without a maximum attempt count, derives API `tracking_degraded` from the consecutive-error threshold, keeps `input_required` on a slower poll, and caps finite positive remote poll hints at 24 hours. Task-enabled server runtime/binding configuration and `mcpInterceptors` are frozen to the Gateway startup snapshot; hot drift fails clearly before tool discovery can diverge from background calls, while presentation-only fields and non-task servers remain reloadable. Configured task toolsets fail startup when the runtime is disabled or persistence is memory. Users still cannot submit an answer back to an `input_required` remote task. - **Long-running ordinary task driver**: `extensions_config.json -> mcpServers.<server>.task_toolsets` binds exact raw submit/status/cancel names; one raw tool may occupy only one role across that server's groups. `mcp/tools.py` hides status/cancel and replaces submit with a wrapper that returns only the local task ID after persistence. `ordinary.py` reads only MCP `structuredContent`, maps remote `running` to `working`, and treats `error_code=task_not_found` or malformed structured output as permanent failure. A status call with `isError=true` is a retryable call failure: the first text content block is retained as a bounded diagnostic, while a permanent remote-task outcome must arrive in a normal result with structured `status=failed`. `task_tool_caller.py` restores the same canonical `(server_name, versioned user/thread/incarnation scope)` stdio session as ordinary calls; non-NULL incarnations use an unambiguous versioned encoding, while an explicitly captured legacy NULL retains the pre-activation `user_id:thread_id` scope. A runtime with a missing/invalid incarnation key fails closed. HTTP/SSE calls remain ephemeral, apply `session_init_timeout` to initialization and `tool_call_timeout` to task calls, and support server-level OAuth refresh outside an Agent run. `McpTaskService` exponentially backs off transient status/cancel errors without a maximum attempt count, derives API `tracking_degraded` from the consecutive-error threshold, keeps `input_required` on a slower poll, and caps finite positive remote poll hints at 24 hours. Task-enabled server runtime/binding configuration and `mcpInterceptors` are frozen to the Gateway startup snapshot; hot drift fails clearly before tool discovery can diverge from background calls, while presentation-only fields and non-task servers remain reloadable. Configured task toolsets fail startup when the runtime is disabled or persistence is memory. Users still cannot submit an answer back to an `input_required` remote task.
- **Durable task payload bounds**: persisted task errors are capped at 4,000 characters. `input_required` and `result_artifact` must each serialize as valid JSON within 64 KiB; an invalid or oversized payload becomes a permanent protocol failure rather than being truncated and changing its semantics. Remote task IDs/task names are limited to 255 characters and task-enabled server names to 128, matching the SQL schema; an oversized submitted remote ID is rejected only after the Service has the handle so compensation cancellation still runs. Oversized results retain the existing bounded preview/truncation/artifact behavior. - **Durable task payload bounds**: persisted task errors are capped at 4,000 characters. `input_required` and `result_artifact` must each serialize as valid JSON within 64 KiB; an invalid or oversized payload becomes a permanent protocol failure rather than being truncated and changing its semantics. Remote task IDs/task names are limited to 255 characters and task-enabled server names to 128, matching the SQL schema; an oversized submitted remote ID is rejected only after the Service has the handle so compensation cancellation still runs. Oversized results retain the existing bounded preview/truncation/artifact behavior.
- **Remote task bring-up**: One `session_init_timeout` deadline covers transport context entry and MCP initialization. Report the server and configured bound on expiry; disable it before calling the tool. Enter and exit the adapter's AnyIO context managers in the same task; tests include real SSE endpoint/initialization stalls and connection cleanup.
- **Lazy initialization**: Tools loaded on first use via `get_cached_mcp_tools()` - **Lazy initialization**: Tools loaded on first use via `get_cached_mcp_tools()`
- **Loop-isolated stdio sessions**: Both the live registry and in-flight creations - **Loop-isolated stdio sessions**: Both the live registry and in-flight creations
are keyed by `(server_name, scope_key, owning_loop)`. Same-loop callers share are keyed by `(server_name, scope_key, owning_loop)`. Same-loop callers share

View File

@ -219,15 +219,16 @@ class McpTaskToolCaller:
) )
captured: BaseException | None = None captured: BaseException | None = None
call_result: Any | None = None call_result: Any | None = None
async with create_session(effective_connection) as remote_session: # Bound transport entry and initialization together, keeping the
initialize = remote_session.initialize() # adapter's AnyIO context managers in the same task for cleanup.
if session_init_timeout_seconds is not None: try:
await asyncio.wait_for( async with (
initialize, asyncio.timeout(session_init_timeout_seconds) as init_timeout,
timeout=session_init_timeout_seconds, create_session(effective_connection) as remote_session,
) ):
else: await remote_session.initialize()
await initialize # Tool calls have their own independent timeout below.
init_timeout.reschedule(None)
try: try:
call = remote_session.call_tool( call = remote_session.call_tool(
request.name, request.name,
@ -243,6 +244,10 @@ class McpTaskToolCaller:
call_result = await call call_result = await call
except BaseException as exc: # preserve adapter disconnect semantics except BaseException as exc: # preserve adapter disconnect semantics
captured = exc captured = exc
except TimeoutError:
if not init_timeout.expired():
raise
raise TimeoutError(f"MCP task session initialization for server {server_name!r} timed out after {session_init_timeout_seconds}s") from None
if captured is not None: if captured is not None:
raise captured raise captured
if call_result is None: if call_result is None:

View File

@ -1,7 +1,7 @@
import asyncio import asyncio
import sys import sys
from collections.abc import Coroutine from collections.abc import Coroutine
from contextlib import suppress from contextlib import asynccontextmanager, suppress
from datetime import timedelta from datetime import timedelta
from types import SimpleNamespace from types import SimpleNamespace
from typing import Any from typing import Any
@ -54,6 +54,16 @@ def _remote_config(transport: str = "http") -> ExtensionsConfig:
) )
def _remote_status_call(config: ExtensionsConfig) -> Coroutine[Any, Any, Any]:
return McpTaskToolCaller(config).call_tool(
server_name="reports",
tool_name="status_report",
arguments={"task_id": "remote-1"},
user_id="user-1",
thread_id="thread-1",
)
class _SessionContext: class _SessionContext:
def __init__(self, session): def __init__(self, session):
self.session = session self.session = session
@ -65,13 +75,15 @@ class _SessionContext:
return None return None
async def _assert_configured_timeout(awaitable: Coroutine[Any, Any, Any]) -> None: async def _assert_configured_timeout(awaitable: Coroutine[Any, Any, Any], *, wait_timeout: float = 0.25, expected_message: str | None = None) -> None:
task = asyncio.create_task(awaitable) task = asyncio.create_task(awaitable)
try: try:
done, _pending = await asyncio.wait({task}, timeout=0.25) done, _pending = await asyncio.wait({task}, timeout=wait_timeout)
assert task in done, "configured timeout was ignored" assert task in done, "configured timeout was ignored"
with pytest.raises(TimeoutError): with pytest.raises(TimeoutError) as error:
await task await task
if expected_message is not None:
assert str(error.value) == expected_message
finally: finally:
if not task.done(): if not task.done():
task.cancel() task.cancel()
@ -437,14 +449,15 @@ async def test_remote_task_session_initialization_respects_configured_timeout(tr
MagicMock(return_value=_SessionContext(session)), MagicMock(return_value=_SessionContext(session)),
): ):
await _assert_configured_timeout( await _assert_configured_timeout(
caller.call_tool( expected_message="MCP task session initialization for server 'reports' timed out after 0.01s",
awaitable=caller.call_tool(
server_name="reports", server_name="reports",
tool_name="status_report", tool_name="status_report",
arguments={"task_id": "remote-1"}, arguments={"task_id": "remote-1"},
user_id="user-1", user_id="user-1",
thread_id="thread-1", thread_id="thread-1",
thread_incarnation=None, thread_incarnation=None,
) ),
) )
session.call_tool.assert_not_awaited() session.call_tool.assert_not_awaited()
@ -476,14 +489,15 @@ async def test_remote_task_call_respects_configured_timeout(transport: str) -> N
MagicMock(return_value=_SessionContext(session)), MagicMock(return_value=_SessionContext(session)),
): ):
await _assert_configured_timeout( await _assert_configured_timeout(
caller.call_tool( expected_message="",
awaitable=caller.call_tool(
server_name="reports", server_name="reports",
tool_name="status_report", tool_name="status_report",
arguments={"task_id": "remote-1"}, arguments={"task_id": "remote-1"},
user_id="user-1", user_id="user-1",
thread_id="thread-1", thread_id="thread-1",
thread_incarnation=None, thread_incarnation=None,
) ),
) )
session.call_tool.assert_awaited_once_with( session.call_tool.assert_awaited_once_with(
@ -491,3 +505,218 @@ async def test_remote_task_call_respects_configured_timeout(transport: str) -> N
{"task_id": "remote-1"}, {"task_id": "remote-1"},
read_timeout_seconds=timedelta(seconds=0.01), read_timeout_seconds=timedelta(seconds=0.01),
) )
@pytest.mark.asyncio
@pytest.mark.parametrize("transport", ["http", "sse"])
@pytest.mark.parametrize("phase", ["connect", "initialize", "call", "cleanup"])
async def test_remote_task_unrelated_timeout_is_not_relabelled(transport: str, phase: str) -> None:
config = _remote_config(transport)
config.mcp_servers["reports"].session_init_timeout = 1.0
original = TimeoutError(f"original {phase} timeout")
session = SimpleNamespace(initialize=AsyncMock(), call_tool=AsyncMock(return_value={"ok": True}))
if phase == "initialize":
session.initialize.side_effect = original
elif phase == "call":
session.call_tool.side_effect = original
@asynccontextmanager
async def fake_session(_connection):
if phase == "connect":
raise original
yield session
if phase == "cleanup":
raise original
with patch("langchain_mcp_adapters.sessions.create_session", fake_session):
with pytest.raises(TimeoutError) as error:
await _remote_status_call(config)
assert error.value is original
@pytest.mark.asyncio
@pytest.mark.parametrize("transport", ["http", "sse"])
async def test_remote_task_session_open_respects_configured_timeout(transport: str) -> None:
config = _remote_config(transport)
config.mcp_servers["reports"].session_init_timeout = 0.01
session = SimpleNamespace(initialize=AsyncMock(), call_tool=AsyncMock())
closed = asyncio.Event()
@asynccontextmanager
async def slow_session(_connection):
try:
async with anyio.create_task_group():
await asyncio.Event().wait()
yield session
finally:
closed.set()
with patch("langchain_mcp_adapters.sessions.create_session", slow_session):
await _assert_configured_timeout(expected_message="MCP task session initialization for server 'reports' timed out after 0.01s", awaitable=_remote_status_call(config))
assert closed.is_set()
session.initialize.assert_not_awaited()
session.call_tool.assert_not_awaited()
@pytest.mark.asyncio
@pytest.mark.parametrize("transport", ["http", "sse"])
async def test_remote_task_connection_and_initialize_share_timeout(transport: str) -> None:
config = _remote_config(transport)
config.mcp_servers["reports"].session_init_timeout = 0.5
async def slow_initialize():
await asyncio.sleep(0.3)
session = SimpleNamespace(initialize=AsyncMock(side_effect=slow_initialize), call_tool=AsyncMock())
@asynccontextmanager
async def slow_session(_connection):
async with anyio.create_task_group():
await asyncio.sleep(0.3)
yield session
with patch("langchain_mcp_adapters.sessions.create_session", slow_session):
await _assert_configured_timeout(expected_message="MCP task session initialization for server 'reports' timed out after 0.5s", awaitable=_remote_status_call(config), wait_timeout=2)
session.initialize.assert_awaited_once()
session.call_tool.assert_not_awaited()
@pytest.mark.asyncio
@pytest.mark.parametrize("transport", ["http", "sse"])
@pytest.mark.parametrize("init_timeout", [None, 0.01])
async def test_remote_task_initialization_timeout_does_not_limit_tool_call(transport: str, init_timeout: float | None) -> None:
config = _remote_config(transport)
config.mcp_servers["reports"].session_init_timeout = init_timeout
config.mcp_servers["reports"].tool_call_timeout = 1.0
result = SimpleNamespace(structuredContent={"status": "completed"}, isError=False)
closed = asyncio.Event()
async def slow_call(*_args, **_kwargs):
await asyncio.sleep(0.04)
return result
session = SimpleNamespace(initialize=AsyncMock(), call_tool=AsyncMock(side_effect=slow_call))
@asynccontextmanager
async def task_group_session(_connection):
owner = asyncio.current_task()
async with anyio.create_task_group():
try:
yield session
finally:
assert asyncio.current_task() is owner
closed.set()
with patch("langchain_mcp_adapters.sessions.create_session", task_group_session):
assert await _remote_status_call(config) is result
assert closed.is_set()
session.call_tool.assert_awaited_once_with("status_report", {"task_id": "remote-1"}, read_timeout_seconds=timedelta(seconds=1))
@pytest.mark.asyncio
@pytest.mark.parametrize("transport", ["http", "sse"])
@pytest.mark.parametrize("phase", ["connect", "initialize", "call"])
async def test_remote_task_session_preserves_external_cancellation(transport: str, phase: str) -> None:
config = _remote_config(transport)
reached = asyncio.Event()
closed = asyncio.Event()
async def pause_at(current_phase):
if current_phase == phase:
reached.set()
await asyncio.Event().wait()
async def initialize():
await pause_at("initialize")
async def call_tool(*_args, **_kwargs):
await pause_at("call")
session = SimpleNamespace(initialize=AsyncMock(side_effect=initialize), call_tool=AsyncMock(side_effect=call_tool))
@asynccontextmanager
async def task_group_session(_connection):
owner = asyncio.current_task()
try:
async with anyio.create_task_group():
await pause_at("connect")
yield session
finally:
assert asyncio.current_task() is owner
closed.set()
with patch("langchain_mcp_adapters.sessions.create_session", task_group_session):
task = asyncio.create_task(_remote_status_call(config))
try:
await asyncio.wait_for(reached.wait(), 1)
task.cancel()
done, _ = await asyncio.wait({task}, timeout=1)
assert task in done, "external cancellation did not finish session cleanup"
with pytest.raises(asyncio.CancelledError):
await task
assert closed.is_set()
finally:
if not task.done():
task.cancel()
with suppress(asyncio.CancelledError):
await task
@pytest.mark.asyncio
@pytest.mark.parametrize("phase", ["endpoint", "initialize"])
async def test_sse_task_session_timeout_closes_real_connection(monkeypatch, phase: str) -> None:
monkeypatch.setenv("NO_PROXY", "127.0.0.1,localhost")
monkeypatch.setenv("no_proxy", "127.0.0.1,localhost")
connected = asyncio.Event()
initialized = asyncio.Event()
disconnected = asyncio.Event()
handlers: set[asyncio.Task] = set()
async def handle(reader: asyncio.StreamReader, writer: asyncio.StreamWriter):
handlers.add(asyncio.current_task())
is_sse = False
try:
request = await reader.readuntil(b"\r\n\r\n")
is_sse = request.startswith(b"GET /sse ")
if is_sse:
writer.write(b"HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\n\r\n: connected\n\n")
if phase == "initialize":
writer.write(b"event: endpoint\ndata: /messages/\n\n")
await writer.drain()
connected.set()
await reader.read()
else:
assert request.startswith(b"POST /messages/ ")
for header in request.split(b"\r\n"):
if header.lower().startswith(b"content-length:"):
await reader.readexactly(int(header.split(b":", 1)[1]))
break
writer.write(b"HTTP/1.1 202 Accepted\r\nContent-Length: 0\r\nConnection: close\r\n\r\n")
await writer.drain()
initialized.set()
finally:
writer.close()
await writer.wait_closed()
if is_sse:
disconnected.set()
server = await asyncio.start_server(handle, "127.0.0.1", 0)
config = _remote_config("sse")
config.mcp_servers["reports"].url = f"http://127.0.0.1:{server.sockets[0].getsockname()[1]}/sse"
config.mcp_servers["reports"].session_init_timeout = 0.5
try:
await _assert_configured_timeout(expected_message="MCP task session initialization for server 'reports' timed out after 0.5s", awaitable=_remote_status_call(config), wait_timeout=3)
assert connected.is_set()
assert initialized.is_set() == (phase == "initialize")
await asyncio.wait_for(disconnected.wait(), 1)
finally:
server.close()
for handler in handlers:
if not handler.done():
handler.cancel()
results = await asyncio.gather(*handlers, return_exceptions=True)
await server.wait_closed()
assert all(not isinstance(result, BaseException) or isinstance(result, asyncio.CancelledError) for result in results), results