mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-09-25 05:56:18 +00:00
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:
parent
cb349da014
commit
eae70b5e40
@ -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.
|
||||
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 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.
|
||||
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.
|
||||
|
||||
@ -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
|
||||
calls. `session_init_timeout` covers server bring-up — tool discovery
|
||||
(subprocess spawn + `initialize` + `tools/list`) and persistent-session
|
||||
initialization — plus ephemeral HTTP/SSE task-session initialization. It
|
||||
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
|
||||
initialization — plus ephemeral HTTP/SSE task-session connection setup and
|
||||
initialization under a single deadline. This includes waiting for an SSE
|
||||
`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:
|
||||
|
||||
```json
|
||||
|
||||
@ -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.
|
||||
- **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.
|
||||
- **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()`
|
||||
- **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
|
||||
|
||||
@ -219,30 +219,35 @@ class McpTaskToolCaller:
|
||||
)
|
||||
captured: BaseException | None = None
|
||||
call_result: Any | None = None
|
||||
async with create_session(effective_connection) as remote_session:
|
||||
initialize = remote_session.initialize()
|
||||
if session_init_timeout_seconds is not None:
|
||||
await asyncio.wait_for(
|
||||
initialize,
|
||||
timeout=session_init_timeout_seconds,
|
||||
)
|
||||
else:
|
||||
await initialize
|
||||
try:
|
||||
call = remote_session.call_tool(
|
||||
request.name,
|
||||
request.args,
|
||||
**call_kwargs,
|
||||
)
|
||||
if timeout_seconds:
|
||||
call_result = await asyncio.wait_for(
|
||||
call,
|
||||
timeout=timeout_seconds,
|
||||
# Bound transport entry and initialization together, keeping the
|
||||
# adapter's AnyIO context managers in the same task for cleanup.
|
||||
try:
|
||||
async with (
|
||||
asyncio.timeout(session_init_timeout_seconds) as init_timeout,
|
||||
create_session(effective_connection) as remote_session,
|
||||
):
|
||||
await remote_session.initialize()
|
||||
# Tool calls have their own independent timeout below.
|
||||
init_timeout.reschedule(None)
|
||||
try:
|
||||
call = remote_session.call_tool(
|
||||
request.name,
|
||||
request.args,
|
||||
**call_kwargs,
|
||||
)
|
||||
else:
|
||||
call_result = await call
|
||||
except BaseException as exc: # preserve adapter disconnect semantics
|
||||
captured = exc
|
||||
if timeout_seconds:
|
||||
call_result = await asyncio.wait_for(
|
||||
call,
|
||||
timeout=timeout_seconds,
|
||||
)
|
||||
else:
|
||||
call_result = await call
|
||||
except BaseException as exc: # preserve adapter disconnect semantics
|
||||
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:
|
||||
raise captured
|
||||
if call_result is None:
|
||||
|
||||
@ -1,7 +1,7 @@
|
||||
import asyncio
|
||||
import sys
|
||||
from collections.abc import Coroutine
|
||||
from contextlib import suppress
|
||||
from contextlib import asynccontextmanager, suppress
|
||||
from datetime import timedelta
|
||||
from types import SimpleNamespace
|
||||
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:
|
||||
def __init__(self, session):
|
||||
self.session = session
|
||||
@ -65,13 +75,15 @@ class _SessionContext:
|
||||
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)
|
||||
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"
|
||||
with pytest.raises(TimeoutError):
|
||||
with pytest.raises(TimeoutError) as error:
|
||||
await task
|
||||
if expected_message is not None:
|
||||
assert str(error.value) == expected_message
|
||||
finally:
|
||||
if not task.done():
|
||||
task.cancel()
|
||||
@ -437,14 +449,15 @@ async def test_remote_task_session_initialization_respects_configured_timeout(tr
|
||||
MagicMock(return_value=_SessionContext(session)),
|
||||
):
|
||||
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",
|
||||
tool_name="status_report",
|
||||
arguments={"task_id": "remote-1"},
|
||||
user_id="user-1",
|
||||
thread_id="thread-1",
|
||||
thread_incarnation=None,
|
||||
)
|
||||
),
|
||||
)
|
||||
|
||||
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)),
|
||||
):
|
||||
await _assert_configured_timeout(
|
||||
caller.call_tool(
|
||||
expected_message="",
|
||||
awaitable=caller.call_tool(
|
||||
server_name="reports",
|
||||
tool_name="status_report",
|
||||
arguments={"task_id": "remote-1"},
|
||||
user_id="user-1",
|
||||
thread_id="thread-1",
|
||||
thread_incarnation=None,
|
||||
)
|
||||
),
|
||||
)
|
||||
|
||||
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"},
|
||||
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
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user