diff --git a/README.md b/README.md index bad9db9e08..d282943b87 100644 --- a/README.md +++ b/README.md @@ -633,6 +633,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`). +Durable HTTP/SSE task status and cancellation calls select configured `user_auth` credentials using the persisted task owner, including after restart; per-request secrets are not retained for background calls. If a request-scoped credential overrides submit authentication, both credentials must authorize access to the same remote task. 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 file outputs, a bare filename is linked to a uniquely matching file created or changed by that call. Filenames embedded in unrelated paths, including Windows backslash paths, are left intact. 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. diff --git a/backend/docs/MCP_SERVER.md b/backend/docs/MCP_SERVER.md index d6535d0348..ed3386570b 100644 --- a/backend/docs/MCP_SERVER.md +++ b/backend/docs/MCP_SERVER.md @@ -392,13 +392,25 @@ must therefore persist its own tasks; multi-instance deployments should normally use an independently running HTTP/SSE service. Server-level OAuth works during background polling and refreshes normally. +When `user_auth` is enabled on an HTTP/SSE server, background status and +cancellation calls use the persisted task owner's configured credential, +including after a Gateway restart. +Only the user ID is carried from the task record; no request credential is stored. +An unmapped owner remains denied unless `user_auth.on_missing` is `passthrough`. Request-scoped secrets from a particular Agent run are not durable task -credentials and are unavailable to later background polls; use server-level +credentials and are unavailable to later background polls; use configured authentication for a task toolset. `headers_from_context` follows the same rule: submit is awaited inside the Agent run and carries the mapped headers, while status and cancel polls skip them and authenticate with the server's -static or OAuth credentials — so `on_missing: "deny"` guards the submit but not -those polls. Declaring both on one server logs a warning at startup. Restart DeerFlow after changing +static/OAuth credentials or the owner's configured `user_auth` credential — so +`headers_from_context.on_missing: "deny"` guards the submit but not +those polls. Declaring both on one server logs a warning at startup. +When a request header overrides `user_auth` on submit, ensure that both +credentials can access the same remote task. If the background credential +cannot access it and the status tool returns a normal structured +`error_code: "task_not_found"` result, the task becomes permanently `failed`, +not a retryable authentication error. +Restart DeerFlow after changing `mcp_tasks`, `task_toolsets`, `mcpInterceptors`, or any connection, authentication, transport, or timeout setting on a task-enabled server. DeerFlow rejects task-tool reloads that no longer match the Gateway's startup @@ -514,7 +526,8 @@ The caller supplies the values on each run request: - Durable background tasks are the one exception, and only half of one: a `task_toolsets` submit is awaited inside the Agent run and carries these headers, but the status and cancel polls run after that run ends, so they skip - them and use the server's static/OAuth credentials. See *Durable Background + them and use configured static/OAuth credentials or the persisted owner's + `user_auth` credential. See *Durable Background Tasks* above. Use `user_auth` instead when the credential belongs to a configured DeerFlow diff --git a/backend/packages/harness/deerflow/mcp/AGENTS.md b/backend/packages/harness/deerflow/mcp/AGENTS.md index 307b15f749..456e27117d 100644 --- a/backend/packages/harness/deerflow/mcp/AGENTS.md +++ b/backend/packages/harness/deerflow/mcp/AGENTS.md @@ -6,6 +6,7 @@ - **Long-running ordinary task driver**: `extensions_config.json -> mcpServers..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. +- **Background task identity**: `McpTaskToolCaller` binds an ID-only `_TaskOwner` in the existing user ContextVar only for background HTTP/SSE calls with enabled `user_auth`, restoring the previous value in `finally`. The `CurrentUser` contract guarantees only `.id`, not profile fields or roles. This lets `user_auth` select configured credentials without an active Agent run, including recovery from SQL; concurrent owners remain task-local. Submit, stdio, and calls without enabled `user_auth` keep their original context. Do not persist request secrets or enable `headers_from_context` for background status/cancel calls. - **Lazy initialization**: Tools loaded on first use via `get_cached_mcp_tools()`. Its `asyncio.run` fallback covers only `get_event_loop()` raising (no loop set, or a closed loop); a `RuntimeError` from discovery itself (e.g. `McpTaskConfigurationError`) is @@ -29,7 +30,7 @@ - **Per-server tool-name prefixing**: `mcpServers..tool_name_prefix` defaults to `true`, preserving the collision-safe `_` prefix. Servers whose tools already carry a stable namespace may set it to `false`; discovery then calls `langchain_mcp_adapters.tools.load_mcp_tools` with that server's flag. Source routing and stdio session-pool wrapping are based on the producing server and transport, never on whether the visible tool name starts with the server prefix. - **OAuth (HTTP/SSE)**: Supports token endpoint flows (`client_credentials`, `refresh_token`) with automatic token refresh + Authorization header injection. The rendered ` ` passes `mcp/headers.py::illegal_header_value_reason` inside `OAuthTokenManager.get_authorization_header` — the one boundary the tool interceptor, the initial discovery headers and the durable task path all read their value from — so a token endpoint returning something the transport would refuse fails closed instead of letting h11 echo the token into a model-visible tool error. The rendered value is what gets checked, not the two fields separately, because that is what the transport sees. The operator's static `headers` get the same check in `mcp/client.py::build_server_params`, where `build_servers_config` already drops just that server and logs the reason. - **Per-user credentials (HTTP/SSE)**: `mcpServers..user_auth` maps DeerFlow user ids to credential header values (`$ENV_VAR` references supported). `mcp/user_scoped_auth.py::build_user_scoped_auth_interceptor` rewrites the configured header on every tool call from the authenticated runtime user (registered after OAuth in `mcp/interceptors.py`, so its per-user value wins the header for servers declaring both). Fail-closed: an unmapped user or an empty resolved credential raises a `ToolException` unless `on_missing: "passthrough"` is set; a resolved credential the transport would refuse as a header value (line break, surrounding whitespace, non-ASCII — `mcp/headers.py::illegal_header_value_reason`) is always denied without echoing the value, since h11 renders the full value into its exception message on the line break and whitespace cases and tool errors are model-visible (httpx catches the non-ASCII case earlier, naming only the offending character). The server's static `headers` serve startup tool discovery only. Gateway GET masks `user_auth.users` values; PUT round-trips masked values by preserving stored credentials. -- **Per-request credentials (HTTP/SSE)**: `mcpServers..headers_from_context` maps HTTP header names to keys of the run request's `config.context.secrets` carrier, for credentials the caller chooses per request (multi-tenant gateways, per-run API keys) rather than per configured user. `mcp/context_headers.py::build_context_headers_interceptor` resolves the mapping on every tool call and rewrites those headers. Registered **after** `user_auth` in `mcp/interceptors.py`, so for a server declaring several sources the per-request value wins the final header: precedence is static `headers` < `oauth` < `user_auth` < `headers_from_context`. Fail-closed: a mapped key absent from the request secrets (or resolved empty) raises a `ToolException` naming only the missing key, unless `on_missing: "passthrough"` is set — a silent fallback would send one tenant's call under the discovery credential's authority. A resolved value the transport would refuse as a header value (line break, surrounding whitespace, non-ASCII) is always denied regardless of `on_missing`, without echoing the value: h11 renders the full value into its exception message on a line break or surrounding whitespace, and `ToolErrorHandlingMiddleware` copies tool errors into model-visible messages, so an unchecked bad credential would land the secret in the prompt, the checkpoint, and traces. `sse`/`http` only; a stdio server warns and is skipped, as with `user_auth`. The block stores names, never a credential, so the Gateway returns it unmasked and a `PUT` replaces the declared mapping verbatim; only `extra="allow"` keys inside it get sensitive-key masking, and those are restored from the stored block on a round-trip like every other masked extra. Durable `task_toolsets` calls split: submit is awaited inside the Agent run and carries the mapped headers (`McpTaskToolCaller.call_tool(request_scoped_headers=True)`, set only by `OrdinaryMcpTaskDriver.submit`), while status and cancel run after that run ended and keep the server-level credentials — so `on_missing: "deny"` covers submit but not those polls, which is what the startup warning is about. +- **Per-request credentials (HTTP/SSE)**: `mcpServers..headers_from_context` maps HTTP header names to keys of the run request's `config.context.secrets` carrier, for credentials the caller chooses per request (multi-tenant gateways, per-run API keys) rather than per configured user. `mcp/context_headers.py::build_context_headers_interceptor` resolves the mapping on every tool call and rewrites those headers. Registered **after** `user_auth` in `mcp/interceptors.py`, so for a server declaring several sources the per-request value wins the final header: precedence is static `headers` < `oauth` < `user_auth` < `headers_from_context`. Fail-closed: a mapped key absent from the request secrets (or resolved empty) raises a `ToolException` naming only the missing key, unless `on_missing: "passthrough"` is set — a silent fallback would send one tenant's call under the discovery credential's authority. A resolved value the transport would refuse as a header value (line break, surrounding whitespace, non-ASCII) is always denied regardless of `on_missing`, without echoing the value: h11 renders the full value into its exception message on a line break or surrounding whitespace, and `ToolErrorHandlingMiddleware` copies tool errors into model-visible messages, so an unchecked bad credential would land the secret in the prompt, the checkpoint, and traces. `sse`/`http` only; a stdio server warns and is skipped, as with `user_auth`. The block stores names, never a credential, so the Gateway returns it unmasked and a `PUT` replaces the declared mapping verbatim; only `extra="allow"` keys inside it get sensitive-key masking, and those are restored from the stored block on a round-trip like every other masked extra. Durable `task_toolsets` calls split: submit is awaited inside the Agent run and carries the mapped headers (`McpTaskToolCaller.call_tool(request_scoped_headers=True)`, set only by `OrdinaryMcpTaskDriver.submit`), while status and cancel run after that run ended and use configured static/OAuth credentials or the persisted owner's `user_auth` credential — so `headers_from_context.on_missing: "deny"` covers submit but not those polls, which is what the startup warning is about. - **Header names are case-insensitive** (`mcp/headers.py`): every credential interceptor writes through `apply_header_overrides`, which drops a key differing only in case and emits the spelling the connection already uses. Without it a static `authorization` and an injected `Authorization` both reach httpx — the adapter merges connection and override headers with a plain `{**static, **override}` splat — and a server reading the field with a single-value accessor gets the static entry, silently inverting the precedence above. `headers_from_context.headers` also rejects two spellings of one header at config load. - **Reading the run context from an interceptor**: use `request.runtime` (LangGraph's tool node injects a `ToolRuntime` into any tool parameter named `runtime`, which covers both the pooled stdio wrapper and `langchain-mcp-adapters`' own HTTP/SSE tool), falling back to ambient `langgraph.runtime.get_runtime()`. Do **not** use `langgraph.config.get_config()["context"]`: the run context rides the runtime, not the `RunnableConfig` propagated to child runnables, so that key is `None` inside a tool call. `tests/test_mcp_context_headers.py::test_adapter_tool_receives_the_runtime_langgraph_injects` pins the injection rule against an upstream rename by disabling the ambient fallback and driving a real adapter tool through a real graph. - **Routing hints**: `extensions_config.json -> mcpServers..routing` and diff --git a/backend/packages/harness/deerflow/mcp/context_headers.py b/backend/packages/harness/deerflow/mcp/context_headers.py index b985c8dfc4..18f0b35016 100644 --- a/backend/packages/harness/deerflow/mcp/context_headers.py +++ b/backend/packages/harness/deerflow/mcp/context_headers.py @@ -123,9 +123,10 @@ def build_context_headers_interceptor(extensions_config: ExtensionsConfig) -> An # the request secrets; the later status/cancel polls do not, because # the task runtime drives them long after that run ended. Those calls # deliberately skip this interceptor (see McpTaskToolCaller), so the - # background half authenticates with the server's own credentials. + # background half uses configured credentials, including user_auth + # selected by the persisted task owner. logger.warning( - "MCP server '%s' declares both headers_from_context and task_toolsets; background task status/cancel polls run outside an Agent run and will use this server's static/OAuth credentials instead of the per-request headers", + "MCP server '%s' declares both headers_from_context and task_toolsets; background task status/cancel polls run outside an Agent run and will use configured credentials instead of the per-request headers", server_name, ) mapping_by_server[server_name] = context_headers diff --git a/backend/packages/harness/deerflow/mcp/task_tool_caller.py b/backend/packages/harness/deerflow/mcp/task_tool_caller.py index 52309f5b0e..6290736377 100644 --- a/backend/packages/harness/deerflow/mcp/task_tool_caller.py +++ b/backend/packages/harness/deerflow/mcp/task_tool_caller.py @@ -5,6 +5,7 @@ from __future__ import annotations import asyncio import logging from collections.abc import Mapping +from dataclasses import dataclass from datetime import timedelta from typing import Any @@ -18,10 +19,18 @@ 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 MCPSessionPool, call_pooled_session_tool, get_session_pool from deerflow.mcp_scope import mcp_session_scope_key +from deerflow.runtime.user_context import reset_current_user, set_current_user logger = logging.getLogger(__name__) +@dataclass +class _TaskOwner: + """ID-only CurrentUser for background MCP auth; no profile or role data.""" + + id: str + + def _prepare_stdio_connection( connection: dict[str, Any], *, @@ -73,7 +82,8 @@ class McpTaskToolCaller: # own credential. The later status/cancel polls are driven by the task # runtime long after that run ended: there is no run context to read, so # the fail-closed interceptor would deny every poll. Those keep using - # server-level credentials (see docs/MCP_SERVER.md), which is what + # configured credentials, including user_auth for the persisted owner + # (see docs/MCP_SERVER.md), which is what # ``build_context_headers_interceptor`` warns about at startup. if context_headers_interceptor is None: self._interceptors = self._submit_interceptors @@ -98,7 +108,8 @@ class McpTaskToolCaller: inside the Agent run that carries the secrets, while status and cancel run after that run ended. """ - interceptors = self._submit_interceptors if request_scoped_headers else self._interceptors + is_background_call = not request_scoped_headers + interceptors = self._interceptors if is_background_call else self._submit_interceptors server_config = self._extensions_config.get_enabled_mcp_servers().get(server_name) if server_config is None: raise LookupError(f"MCP task server {server_name!r} is missing or disabled in the startup configuration") @@ -142,6 +153,7 @@ class McpTaskToolCaller: server_name=server_name, tool_name=tool_name, arguments=arguments, + background_user_id=None, timeout_seconds=server_config.tool_call_timeout, session_init_timeout_seconds=None, persistent_session=True, @@ -154,6 +166,9 @@ class McpTaskToolCaller: connection.get("headers") or {}, {"Authorization": authorization}, ) + # Only HTTP/SSE servers with enabled user_auth need an ambient owner. + # Leave other custom-interceptor contexts unchanged. + user_auth = server_config.user_auth return await self._invoke( session=None, pool=None, @@ -162,6 +177,7 @@ class McpTaskToolCaller: server_name=server_name, tool_name=tool_name, arguments=arguments, + background_user_id=user_id if is_background_call and user_auth is not None and user_auth.enabled else None, timeout_seconds=server_config.tool_call_timeout, session_init_timeout_seconds=server_config.session_init_timeout, persistent_session=False, @@ -178,6 +194,7 @@ class McpTaskToolCaller: server_name: str, tool_name: str, arguments: dict[str, Any], + background_user_id: str | None, timeout_seconds: float | None, session_init_timeout_seconds: float | None, persistent_session: bool, @@ -263,11 +280,21 @@ class McpTaskToolCaller: handler = wrapped - return await handler( - MCPToolCallRequest( - name=tool_name, - args=arguments, - server_name=server_name, - runtime=None, + # Durable status/cancel calls run after the originating Agent turn, so + # there is no LangGraph runtime from which the user-scoped auth + # interceptor can resolve an identity. Bind the persisted task owner for + # the duration of this call, leaving the live submit context untouched. + # ContextVar state keeps parallel polls for different users isolated. + user_context_token = set_current_user(_TaskOwner(id=background_user_id)) if background_user_id is not None else None + try: + return await handler( + MCPToolCallRequest( + name=tool_name, + args=arguments, + server_name=server_name, + runtime=None, + ) ) - ) + finally: + if user_context_token is not None: + reset_current_user(user_context_token) diff --git a/backend/tests/test_mcp_task_ordinary_e2e.py b/backend/tests/test_mcp_task_ordinary_e2e.py index 5e84836ca0..c40b41320a 100644 --- a/backend/tests/test_mcp_task_ordinary_e2e.py +++ b/backend/tests/test_mcp_task_ordinary_e2e.py @@ -1,11 +1,15 @@ +from contextlib import asynccontextmanager from datetime import UTC, datetime, timedelta from types import SimpleNamespace +from unittest.mock import AsyncMock import pytest import pytest_asyncio from app.mcp_tasks import McpTaskService from deerflow.config.database_config import DatabaseConfig +from deerflow.config.extensions_config import ExtensionsConfig +from deerflow.mcp.task_tool_caller import McpTaskToolCaller from deerflow.mcp.tasks import ( ORDINARY_MCP_TASK_DRIVER, McpTaskDriverRegistry, @@ -15,6 +19,7 @@ from deerflow.mcp.tasks import ( from deerflow.persistence.engine import close_engine, get_session_factory, init_engine_from_config from deerflow.persistence.mcp_tasks import McpTaskRepository from deerflow.persistence.thread_meta.model import ThreadMetaRow +from deerflow.runtime.user_context import get_current_user, reset_current_user, set_current_user @pytest_asyncio.fixture(autouse=True) @@ -207,3 +212,171 @@ async def test_status_tool_error_retries_with_detail_before_structured_failure_t assert failed["error"] == "report generation failed" assert failed["consecutive_poll_error_count"] == 0 assert failed["next_poll_at"] is None + + +@pytest.mark.asyncio +@pytest.mark.no_auto_user +@pytest.mark.parametrize("transport", ["http", "sse"]) +@pytest.mark.parametrize("operation", ["poll", "cancel"]) +async def test_recovered_task_authenticates_as_persisted_owner(tmp_path, monkeypatch, transport: str, operation: str) -> None: + await init_engine_from_config(DatabaseConfig(backend="sqlite", sqlite_dir=str(tmp_path))) + session_factory = get_session_factory() + assert session_factory is not None + repo = McpTaskRepository(session_factory) + await _create_thread(repo) + config = ExtensionsConfig.model_validate( + { + "mcpServers": { + "reports": { + "type": transport, + "url": "https://reports.example.com/mcp", + "headers": {"Authorization": "Bearer discovery"}, + "user_auth": {"users": {"user-1": "Bearer task-owner"}}, + "task_toolsets": [{"name": "reports", "submit_tool": "submit_report", "status_tool": "status_report", "cancel_tool": "cancel_report"}], + } + } + } + ) + fake_server = FakeMcpServer() + fake_server.status_results.append({"task_id": "remote-1", "status": "completed", "result": {"report": "ready"}}) + opened_headers = [] + tool_names = [] + + async def call_tool(name, arguments): + tool_names.append(name) + return await fake_server.call_tool(tool_name=name, arguments=arguments) + + @asynccontextmanager + async def create_session(connection): + opened_headers.append(dict(connection["headers"])) + yield SimpleNamespace(initialize=AsyncMock(), call_tool=call_tool) + + monkeypatch.setattr("langchain_mcp_adapters.sessions.create_session", create_session) + first_process = _service(repo, McpTaskToolCaller(config)) + foreground_user = SimpleNamespace(id="user-1") + user_token = set_current_user(foreground_user) + try: + created = await first_process.submit( + driver_name=ORDINARY_MCP_TASK_DRIVER, + request=_request("remote-1"), + now=datetime.now(UTC) - timedelta(seconds=2), + ) + assert get_current_user() is foreground_user + finally: + reset_current_user(user_token) + + # The foreground context and caller are gone. Recovery has only the owner + # persisted in SQL; no original run or request credential is rehydrated. + assert get_current_user() is None + recovered_process = _service(repo, McpTaskToolCaller(config)) + if operation == "cancel": + await recovered_process.cancel_task(task_id=created["id"], user_id="user-1", thread_id="thread-1", thread_incarnation="incarnation-1") + await recovered_process.run_once(now=datetime.now(UTC)) + + record = await repo.get(created["id"], user_id="user-1", thread_id="thread-1", thread_incarnation="incarnation-1") + assert record is not None + assert record["status"] == ("completed" if operation == "poll" else "cancelled") + assert tool_names == ["submit_report", "status_report" if operation == "poll" else "cancel_report"] + assert opened_headers == [{"Authorization": "Bearer task-owner"}] * 2 + assert get_current_user() is None + + +@pytest.mark.asyncio +@pytest.mark.no_auto_user +@pytest.mark.parametrize("transport", ["http", "sse"]) +@pytest.mark.parametrize( + ("operation", "owner_can_access"), + [("poll", True), ("cancel", True), ("poll", False)], + ids=["poll-shared-task", "cancel-shared-task", "poll-task-not-found"], +) +async def test_recovered_task_with_request_and_user_credentials(tmp_path, monkeypatch, transport: str, operation: str, owner_can_access: bool) -> None: + from langchain_core.messages import AIMessage + from langchain_core.tools import tool + from langgraph.graph import END, START, MessagesState, StateGraph + from langgraph.prebuilt import ToolNode + + await init_engine_from_config(DatabaseConfig(backend="sqlite", sqlite_dir=str(tmp_path))) + session_factory = get_session_factory() + assert session_factory is not None + repo = McpTaskRepository(session_factory) + await _create_thread(repo) + config = ExtensionsConfig.model_validate( + { + "mcpServers": { + "reports": { + "type": transport, + "url": "https://reports.example.com/mcp", + "headers": {"Authorization": "Bearer discovery"}, + "user_auth": {"users": {"user-1": "Bearer task-owner"}}, + "headers_from_context": {"headers": {"Authorization": "reports_token"}}, + "task_toolsets": [{"name": "reports", "submit_tool": "submit_report", "status_tool": "status_report", "cancel_tool": "cancel_report"}], + } + } + } + ) + fake_server = FakeMcpServer() + fake_server.status_results.append({"task_id": "remote-1", "status": "completed", "result": {"report": "ready"}}) + opened_headers = [] + tool_names = [] + + @asynccontextmanager + async def create_session(connection): + opened_headers.append(dict(connection["headers"])) + + async def call_tool(name, arguments): + tool_names.append(name) + if connection["headers"]["Authorization"] == "Bearer task-owner" and not owner_can_access: + return SimpleNamespace(structuredContent={"task_id": arguments["task_id"], "status": "failed", "error_code": "task_not_found"}, content=[], isError=False) + return await fake_server.call_tool(tool_name=name, arguments=arguments) + + yield SimpleNamespace(initialize=AsyncMock(), call_tool=call_tool) + + monkeypatch.setattr("langchain_mcp_adapters.sessions.create_session", create_session) + first_process = _service(repo, McpTaskToolCaller(config)) + + @tool + async def submit_report() -> str: + """Submit a durable report task.""" + created = await first_process.submit( + driver_name=ORDINARY_MCP_TASK_DRIVER, + request=_request("remote-1"), + now=datetime.now(UTC) - timedelta(seconds=2), + ) + return created["id"] + + builder = StateGraph(MessagesState, context_schema=dict) + builder.add_node("tools", ToolNode([submit_report])) + builder.add_edge(START, "tools") + builder.add_edge("tools", END) + graph = builder.compile() + user_token = set_current_user(SimpleNamespace(id="user-1")) + try: + submitted = await graph.ainvoke( + {"messages": [AIMessage(content="", tool_calls=[{"name": "submit_report", "args": {}, "id": "call-1", "type": "tool_call"}])]}, + context={"user_id": "user-1", "secrets": {"reports_token": "Bearer request"}}, + ) + finally: + reset_current_user(user_token) + + # The real graph supplied the submit secret. Its runtime and foreground + # identity are no longer present when a new caller recovers the SQL task. + assert get_current_user() is None + task_id = submitted["messages"][-1].content + recovered_process = _service(repo, McpTaskToolCaller(config)) + if operation == "cancel": + await recovered_process.cancel_task(task_id=task_id, user_id="user-1", thread_id="thread-1", thread_incarnation="incarnation-1") + await recovered_process.run_once(now=datetime.now(UTC)) + + record = await repo.get(task_id, user_id="user-1", thread_id="thread-1", thread_incarnation="incarnation-1") + assert record is not None + if owner_can_access: + assert record["status"] == ("completed" if operation == "poll" else "cancelled") + else: + assert record["status"] == "failed" + assert record["error"] == "Remote MCP task was not found" + assert record["next_poll_at"] is None + assert record["consecutive_poll_error_count"] == 0 + await recovered_process.run_once(now=datetime.now(UTC) + timedelta(seconds=60)) + assert tool_names == ["submit_report", "status_report" if operation == "poll" else "cancel_report"] + assert opened_headers == [{"Authorization": "Bearer request"}, {"Authorization": "Bearer task-owner"}] + assert get_current_user() is None diff --git a/backend/tests/test_mcp_task_tool_caller.py b/backend/tests/test_mcp_task_tool_caller.py index 7c27d85e45..8ba3330372 100644 --- a/backend/tests/test_mcp_task_tool_caller.py +++ b/backend/tests/test_mcp_task_tool_caller.py @@ -12,11 +12,12 @@ 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.extensions_config import ExtensionsConfig, McpUserScopedAuthConfig from deerflow.config.paths import Paths from deerflow.mcp.session_pool import MCPSessionPool from deerflow.mcp.task_tool_caller import McpTaskToolCaller from deerflow.mcp_scope import mcp_session_scope_key +from deerflow.runtime.user_context import get_current_user, reset_current_user, set_current_user def _config() -> ExtensionsConfig: @@ -143,6 +144,51 @@ async def test_stdio_task_call_reuses_exact_scope_and_raw_tool_name() -> None: pool.close_session.assert_not_awaited() +@pytest.mark.asyncio +@pytest.mark.no_auto_user +@pytest.mark.parametrize("has_ambient_user", [False, True], ids=["no-user", "existing-user"]) +@pytest.mark.parametrize( + ("transport", "user_auth_enabled"), + [("stdio", None), ("stdio", False), ("stdio", True), ("http", None), ("http", False), ("sse", None), ("sse", False)], +) +async def test_task_without_applicable_user_auth_preserves_interceptor_context(transport: str, user_auth_enabled: bool | None, has_ambient_user: bool) -> None: + config = _config() if transport == "stdio" else _remote_config(transport) + if user_auth_enabled is not None: + config.mcp_servers["reports"].user_auth = McpUserScopedAuthConfig(enabled=user_auth_enabled, users={"user-1": "Bearer task-owner"}) + caller = McpTaskToolCaller(config) + ambient_user = SimpleNamespace(id="other-user", system_role="member") if has_ambient_user else None + observed_users = [] + + async def inspect_context(request, handler): + observed_users.append(get_current_user()) + assert get_current_user() is ambient_user + return await handler(request) + + caller._interceptors.append(inspect_context) + session = SimpleNamespace(initialize=AsyncMock(), call_tool=AsyncMock(return_value="result")) + pool = SimpleNamespace(get_session=AsyncMock(return_value=session)) + user_token = set_current_user(ambient_user) if ambient_user is not None else None + try: + with ( + patch("deerflow.mcp.task_tool_caller.get_session_pool", return_value=pool), + patch("deerflow.mcp.task_tool_caller._prepare_stdio_connection", side_effect=lambda connection, **_kwargs: connection), + patch("langchain_mcp_adapters.sessions.create_session", return_value=_SessionContext(session)), + ): + result = await caller.call_tool( + server_name="reports", + tool_name="status_report", + arguments={"task_id": "remote-1"}, + user_id="user-1", + thread_id="thread-1", + ) + assert result == "result" + assert observed_users == [ambient_user] + assert get_current_user() is ambient_user + finally: + if user_token is not None: + reset_current_user(user_token) + + @pytest.mark.asyncio @pytest.mark.parametrize( "disconnect_error", @@ -423,6 +469,236 @@ async def test_http_task_call_authenticates_session_initialization() -> None: ) +@pytest.mark.asyncio +@pytest.mark.no_auto_user +@pytest.mark.parametrize("transport", ["http", "sse"]) +@pytest.mark.parametrize("operation", ["get_status", "cancel"]) +async def test_remote_task_uses_persisted_user_for_user_scoped_auth(transport: str, operation: str) -> None: + from deerflow.mcp.tasks import OrdinaryMcpTaskDriver, TaskReference, TaskStatus + + config = _remote_config(transport) + config.mcp_servers["reports"].user_auth = McpUserScopedAuthConfig(users={"user-1": "Bearer user-token"}) + remote_status = "running" if operation == "get_status" else "cancelled" + result = SimpleNamespace(structuredContent={"task_id": "remote-1", "status": remote_status}, isError=False) + session = SimpleNamespace(initialize=AsyncMock(), call_tool=AsyncMock(return_value=result)) + create_session = MagicMock(return_value=_SessionContext(session)) + caller = McpTaskToolCaller(config) + driver = OrdinaryMcpTaskDriver(caller) + reference = TaskReference( + local_task_id="local-1", + user_id="user-1", + thread_id="thread-1", + server_name="reports", + remote_task_id="remote-1", + driver_data={"status_tool": "status_report", "cancel_tool": "cancel_report"}, + ) + previous_user = get_current_user() + with patch("langchain_mcp_adapters.sessions.create_session", create_session): + snapshot = await getattr(driver, operation)(reference) + + assert get_current_user() is previous_user + assert snapshot.status == (TaskStatus.WORKING if operation == "get_status" else TaskStatus.CANCELLED) + create_session.assert_called_once_with( + { + "transport": transport, + "url": "https://reports.example.com/mcp", + "headers": { + "X-Static": "configured", + "Authorization": "Bearer user-token", + }, + } + ) + session.initialize.assert_awaited_once_with() + session.call_tool.assert_awaited_once_with("status_report" if operation == "get_status" else "cancel_report", {"task_id": "remote-1"}) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("transport", ["http", "sse"]) +async def test_concurrent_task_calls_keep_user_credentials_isolated(transport: str) -> None: + config = _remote_config(transport) + config.mcp_servers["reports"].user_auth = McpUserScopedAuthConfig(users={"user-1": "Bearer one", "user-2": "Bearer two"}) + config.mcp_servers["reports"].session_init_timeout = None + caller = McpTaskToolCaller(config) + opened: dict[str, str] = {} + entered = {user_id: asyncio.Event() for user_id in ("user-1", "user-2")} + + @asynccontextmanager + async def create_session(connection): + user = get_current_user() + assert user is not None + user_id = user.id + opened[user_id] = connection["headers"]["Authorization"] + entered[user_id].set() + await asyncio.wait_for(entered["user-2" if user_id == "user-1" else "user-1"].wait(), timeout=5) + assert get_current_user() is user + yield SimpleNamespace(initialize=AsyncMock(), call_tool=AsyncMock(return_value=user_id)) + assert get_current_user() is user + + previous_user = get_current_user() + with patch("langchain_mcp_adapters.sessions.create_session", create_session): + tasks = [ + asyncio.create_task( + caller.call_tool( + server_name="reports", + tool_name="status_report", + arguments={"task_id": f"remote-{user_id}"}, + user_id=user_id, + thread_id=f"thread-{user_id}", + ) + ) + for user_id in ("user-1", "user-2") + ] + try: + results = await asyncio.gather(*tasks) + finally: + for task in tasks: + if not task.done(): + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + + assert results == ["user-1", "user-2"] + assert opened == {"user-1": "Bearer one", "user-2": "Bearer two"} + assert get_current_user() is previous_user + + +@pytest.mark.asyncio +@pytest.mark.parametrize("transport", ["http", "sse"]) +@pytest.mark.parametrize("credential", [None, "", "Bearer invalid\n"]) +async def test_task_user_auth_denial_does_not_fall_back_to_another_user(transport: str, credential: str | None) -> None: + from langchain_core.tools import ToolException + + config = _remote_config(transport) + config.mcp_servers["reports"].headers = {"Authorization": "Bearer discovery"} + users = {"other-user": "Bearer other"} + if credential is not None: + users["user-1"] = credential + config.mcp_servers["reports"].user_auth = McpUserScopedAuthConfig(users=users) + caller = McpTaskToolCaller(config) + previous_user = SimpleNamespace(id="other-user") + token = set_current_user(previous_user) + try: + with patch("langchain_mcp_adapters.sessions.create_session") as create_session, pytest.raises(ToolException) as error: + await caller.call_tool( + server_name="reports", + tool_name="status_report", + arguments={"task_id": "remote-1"}, + user_id="user-1", + thread_id="thread-1", + ) + create_session.assert_not_called() + assert "user-1" in str(error.value) + assert "Bearer" not in str(error.value) + assert get_current_user() is previous_user + finally: + reset_current_user(token) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("transport", ["http", "sse"]) +@pytest.mark.parametrize("phase", ["initialize", "call_tool"]) +async def test_task_user_context_restored_after_remote_error(transport: str, phase: str) -> None: + config = _remote_config(transport) + config.mcp_servers["reports"].user_auth = McpUserScopedAuthConfig(users={"user-1": "Bearer user-token"}) + original_error = RuntimeError("remote service unavailable") + session = SimpleNamespace(initialize=AsyncMock(), call_tool=AsyncMock()) + getattr(session, phase).side_effect = original_error + previous_user = get_current_user() + + with patch("langchain_mcp_adapters.sessions.create_session", return_value=_SessionContext(session)): + with pytest.raises(RuntimeError) as error: + await _remote_status_call(config) + + assert error.value is original_error + assert get_current_user() is previous_user + + +@pytest.mark.asyncio +@pytest.mark.parametrize("transport", ["http", "sse"]) +async def test_task_user_context_restored_after_cancellation(transport: str) -> None: + config = _remote_config(transport) + config.mcp_servers["reports"].user_auth = McpUserScopedAuthConfig(users={"user-1": "Bearer user-token"}) + started = asyncio.Event() + context_restored = asyncio.Event() + + async def blocked_call(*_args): + started.set() + await asyncio.Event().wait() + + session = SimpleNamespace(initialize=AsyncMock(), call_tool=blocked_call) + + async def call_with_context_check(): + previous_user = get_current_user() + try: + await _remote_status_call(config) + finally: + assert get_current_user() is previous_user + context_restored.set() + + with patch("langchain_mcp_adapters.sessions.create_session", return_value=_SessionContext(session)): + task = asyncio.create_task(call_with_context_check()) + try: + await asyncio.wait_for(started.wait(), timeout=5) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + finally: + if not task.done(): + task.cancel() + with suppress(asyncio.CancelledError): + await task + + assert context_restored.is_set() + + +@pytest.mark.asyncio +@pytest.mark.no_auto_user +@pytest.mark.parametrize("transport", ["http", "sse"]) +async def test_task_user_auth_passthrough_keeps_server_credential(transport: str) -> None: + config = _remote_config(transport) + config.mcp_servers["reports"].headers = {"Authorization": "Bearer discovery"} + config.mcp_servers["reports"].user_auth = McpUserScopedAuthConfig(users={}, on_missing="passthrough") + session = SimpleNamespace(initialize=AsyncMock(), call_tool=AsyncMock(return_value="result")) + + with patch("langchain_mcp_adapters.sessions.create_session", return_value=_SessionContext(session)) as create_session: + await McpTaskToolCaller(config).call_tool( + server_name="reports", + tool_name="status_report", + arguments={"task_id": "remote-1"}, + user_id="user-1", + thread_id="thread-1", + ) + + assert create_session.call_args.args[0]["headers"] == {"Authorization": "Bearer discovery"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("transport", ["http", "sse"]) +async def test_task_submit_preserves_the_foreground_user_context(transport: str) -> None: + config = _remote_config(transport) + config.mcp_servers["reports"].user_auth = McpUserScopedAuthConfig(users={"user-1": "Bearer user-token"}) + foreground_user = SimpleNamespace(id="user-1", system_role="member") + + async def initialize(): + assert get_current_user() is foreground_user + + session = SimpleNamespace(initialize=initialize, call_tool=AsyncMock(return_value="submitted")) + token = set_current_user(foreground_user) + try: + with patch("langchain_mcp_adapters.sessions.create_session", return_value=_SessionContext(session)): + result = await McpTaskToolCaller(config).call_tool( + server_name="reports", + tool_name="submit_report", + arguments={}, + user_id="user-1", + thread_id="thread-1", + request_scoped_headers=True, + ) + assert result == "submitted" + assert get_current_user() is foreground_user + finally: + reset_current_user(token) + + @pytest.mark.asyncio @pytest.mark.parametrize("transport", ["http", "sse"]) async def test_remote_task_session_initialization_respects_configured_timeout(transport: str) -> None: