diff --git a/backend/packages/harness/deerflow/mcp/AGENTS.md b/backend/packages/harness/deerflow/mcp/AGENTS.md index b7314e2006..cb231194dd 100644 --- a/backend/packages/harness/deerflow/mcp/AGENTS.md +++ b/backend/packages/harness/deerflow/mcp/AGENTS.md @@ -55,7 +55,7 @@ The optional background-task `status` filter flows through `McpTaskSubmitter.lis - **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 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. +- **Reading the run context from an interceptor**: use `request.runtime` (LangGraph's tool node injects a `ToolRuntime` into a tool parameter by its `ToolRuntime` annotation: `langchain-mcp-adapters`' own HTTP/SSE tool takes it as `runtime`, while the pooled stdio wrapper and the task-submit wrapper take it as `deerflow_runtime` so a business argument named `runtime` still reaches the MCP server; code that reads the injected value from a wrapped tool, like the personal-tool guard in `user_tools.py`, looks the parameter name up per 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 `tools..routing` are soft preference metadata. The effective routing is resolved while `mcp/tools.py::get_mcp_tools()` still has both diff --git a/backend/packages/harness/deerflow/mcp/tools.py b/backend/packages/harness/deerflow/mcp/tools.py index 05444ff7a8..19f5ffcec5 100644 --- a/backend/packages/harness/deerflow/mcp/tools.py +++ b/backend/packages/harness/deerflow/mcp/tools.py @@ -683,15 +683,15 @@ def _make_session_pool_tool( effective_domain = binding.domain if binding is not None else ownership_domain async def call_with_persistent_session( - runtime: Runtime | None = None, + deerflow_runtime: Runtime | None = None, **arguments: Any, ) -> Any: - thread_id = _extract_thread_id(runtime) - user_id = resolve_runtime_user_id(runtime) + thread_id = _extract_thread_id(deerflow_runtime) + user_id = resolve_runtime_user_id(deerflow_runtime) # Scope the pooled session by user *and* thread. Filesystem isolation is # per-(user_id, thread_id), so a thread_id alone could otherwise let two # users with a colliding thread_id share one stateful MCP session. - thread_incarnation = runtime_thread_incarnation(runtime) + thread_incarnation = runtime_thread_incarnation(deerflow_runtime) scope_key = mcp_session_scope_key( user_id=user_id, thread_id=thread_id, @@ -810,7 +810,7 @@ def _make_session_pool_tool( name=original_name, args=arguments, server_name=server_name, - runtime=runtime, + runtime=deerflow_runtime, ) call_tool_result = await handler(request) else: @@ -877,16 +877,16 @@ def _make_background_submit_tool( background_contract = f"Submitted as durable background task {task_name!r}; returns a DeerFlow task ID immediately and status polling is handled automatically." async def submit_in_background( - runtime: Runtime | None = None, + deerflow_runtime: Runtime | None = None, **arguments: Any, ) -> dict[str, Any]: submitter = get_mcp_task_submitter() - thread_id = _extract_thread_id(runtime) - user_id = resolve_runtime_user_id(runtime) - thread_incarnation = runtime_thread_incarnation(runtime) - context = runtime.context if runtime is not None and runtime.context else {} + thread_id = _extract_thread_id(deerflow_runtime) + user_id = resolve_runtime_user_id(deerflow_runtime) + thread_incarnation = runtime_thread_incarnation(deerflow_runtime) + context = deerflow_runtime.context if deerflow_runtime is not None and deerflow_runtime.context else {} run_id = context.get("run_id") - tool_call_id = getattr(runtime, "tool_call_id", None) if runtime is not None else None + tool_call_id = getattr(deerflow_runtime, "tool_call_id", None) if deerflow_runtime is not None else None created = await submitter.submit( driver_name=ORDINARY_MCP_TASK_DRIVER, request=TaskSubmitRequest( diff --git a/backend/packages/harness/deerflow/mcp/user_tools.py b/backend/packages/harness/deerflow/mcp/user_tools.py index b76ced9689..551c2ed524 100644 --- a/backend/packages/harness/deerflow/mcp/user_tools.py +++ b/backend/packages/harness/deerflow/mcp/user_tools.py @@ -4,6 +4,7 @@ from __future__ import annotations import asyncio import functools +import inspect import threading from langchain_core.tools import ToolException @@ -15,8 +16,23 @@ from deerflow.tools.mcp_metadata import get_mcp_source from deerflow.tools.sync import make_sync_tool_wrapper +def _runtime_parameter(coroutine) -> str: + """Name of the parameter LangGraph injects the ``ToolRuntime`` into. + + Pooled stdio wrappers expose ``deerflow_runtime`` so a business argument + named ``runtime`` survives; HTTP/SSE tools keep ``langchain-mcp-adapters``' + own ``runtime`` parameter. + """ + try: + parameters = inspect.signature(coroutine).parameters + except (TypeError, ValueError): + return "runtime" + return "deerflow_runtime" if "deerflow_runtime" in parameters else "runtime" + + def _guard(tool, owner: str, server_name: str): original = tool.coroutine + runtime_parameter = _runtime_parameter(original) snapshot: PersonalMcpConfigSnapshot | None = None snapshot_lock = threading.Lock() @@ -29,7 +45,7 @@ def _guard(tool, owner: str, server_name: str): @functools.wraps(original) async def invoke(*args, **kwargs): - if resolve_runtime_user_id(kwargs.get("runtime")) != owner: + if resolve_runtime_user_id(kwargs.get(runtime_parameter)) != owner: raise ToolException("This MCP connection belongs to another user") server = await asyncio.to_thread(current_connection) if server is None: diff --git a/backend/tests/test_mcp_cwd.py b/backend/tests/test_mcp_cwd.py index 19831330d7..a0819be3a5 100644 --- a/backend/tests/test_mcp_cwd.py +++ b/backend/tests/test_mcp_cwd.py @@ -60,7 +60,7 @@ mcp.run(transport="stdio") try: tools = await mcp_tools.get_mcp_tools() assert [tool.name for tool in tools] == ["local_read_marker"] - content, _artifact = await tools[0].coroutine(runtime=runtime) + content, _artifact = await tools[0].coroutine(deerflow_runtime=runtime) assert content[0]["text"] == "configured-directory" finally: await pool.close_all() @@ -110,7 +110,7 @@ mcp.run(transport="stdio") tools = await mcp_tools.get_mcp_tools() assert [tool.name for tool in tools] == ["local_read_cwd"] assert Path(tools[0].description) == launch_dir.resolve() - content, _artifact = await tools[0].coroutine(runtime=runtime) + content, _artifact = await tools[0].coroutine(deerflow_runtime=runtime) assert Path(content[0]["text"]) == paths.sandbox_work_dir("thread", user_id="user").resolve() finally: await pool.close_all() diff --git a/backend/tests/test_mcp_session_pool.py b/backend/tests/test_mcp_session_pool.py index 344e360ef5..72c4a0261f 100644 --- a/backend/tests/test_mcp_session_pool.py +++ b/backend/tests/test_mcp_session_pool.py @@ -777,7 +777,7 @@ mcp.run(transport="stdio") with patch("deerflow.mcp.tools.get_paths", return_value=Paths(tmp_path)): wrapped = _make_session_pool_tool(original_tool, "crash", connection) with pytest.raises(McpError, match="Connection closed") as exc_info: - await wrapped.coroutine(runtime=runtime) + await wrapped.coroutine(deerflow_runtime=runtime) assert exc_info.value.error.code == CONNECTION_CLOSED scope_key = mcp_session_scope_key( @@ -787,7 +787,7 @@ mcp.run(transport="stdio") ) assert ("crash", scope_key) not in {k[:2] for k in get_session_pool()._entries} - content, _artifact = await wrapped.coroutine(runtime=runtime) + content, _artifact = await wrapped.coroutine(deerflow_runtime=runtime) assert content[0]["text"] == "recovered" @@ -812,11 +812,56 @@ async def test_session_pool_tool_evicts_session_after_transport_disconnect(tmp_p patch("deerflow.mcp.tools.get_paths", return_value=Paths(tmp_path)), pytest.raises(type(transport_error)), ): - await wrapped.coroutine(runtime=_legacy_tool_runtime(), value=1) + await wrapped.coroutine(deerflow_runtime=_legacy_tool_runtime(), value=1) pool.close_session_if_current.assert_awaited_once_with("srv", "test-user-autouse:default", session) +@pytest.mark.asyncio +async def test_session_pool_tool_keeps_a_business_argument_named_runtime(tmp_path): + """A schema property called ``runtime`` is the tool's own argument, not LangGraph's.""" + from langchain_core.messages import AIMessage + from langchain_core.tools import StructuredTool + from langgraph.graph import END, START, MessagesState, StateGraph + from langgraph.prebuilt import ToolNode + + from deerflow.config.paths import Paths + from deerflow.mcp.tools import _make_session_pool_tool + + original_tool = StructuredTool( + name="srv_run", + description="test", + args_schema={ + "type": "object", + "properties": {"runtime": {"type": "string"}, "code": {"type": "string"}}, + "required": ["runtime", "code"], + }, + coroutine=AsyncMock(), + response_format="content_and_artifact", + ) + session = AsyncMock() + session.call_tool = AsyncMock(return_value=CallToolResult(content=[TextContent(type="text", text="ok")], isError=False)) + pool = MagicMock() + pool.get_session = AsyncMock(return_value=session) + pool.close_session_if_current = AsyncMock() + with patch("deerflow.mcp.tools.get_session_pool", return_value=pool): + wrapped = _make_session_pool_tool(original_tool, "srv", {"transport": "stdio", "command": "x", "args": []}) + + builder = StateGraph(MessagesState, context_schema=dict) + builder.add_node("tools", ToolNode([wrapped])) + builder.add_edge(START, "tools") + builder.add_edge("tools", END) + graph = builder.compile() + + with patch("deerflow.mcp.tools.get_paths", return_value=Paths(tmp_path)): + await graph.ainvoke( + {"messages": [AIMessage(content="", tool_calls=[{"name": "srv_run", "args": {"runtime": "python", "code": "print(1)"}, "id": "call-1", "type": "tool_call"}])]}, + context={"thread_id": "thread-1", "thread_incarnation": None}, + ) + + assert session.call_tool.await_args.args[1] == {"runtime": "python", "code": "print(1)"} + + @pytest.mark.asyncio async def test_session_pool_tool_evicts_connection_closed_through_interceptor(tmp_path): """A passthrough interceptor must retain transport-failure recovery.""" @@ -837,7 +882,7 @@ async def test_session_pool_tool_evicts_connection_closed_through_interceptor(tm patch("deerflow.mcp.tools.get_paths", return_value=Paths(tmp_path)), pytest.raises(McpError, match="Connection closed"), ): - await wrapped.coroutine(runtime=_legacy_tool_runtime(), value=1) + await wrapped.coroutine(deerflow_runtime=_legacy_tool_runtime(), value=1) pool.close_session_if_current.assert_awaited_once_with("srv", "test-user-autouse:default", session) @@ -861,7 +906,7 @@ async def test_session_pool_tool_keeps_session_after_nonfatal_mcp_error(tmp_path patch("deerflow.mcp.tools.get_paths", return_value=Paths(tmp_path)), pytest.raises(McpError, match=str(error)), ): - await wrapped.coroutine(runtime=_legacy_tool_runtime(), value=1) + await wrapped.coroutine(deerflow_runtime=_legacy_tool_runtime(), value=1) pool.close_session_if_current.assert_not_awaited() @@ -880,7 +925,7 @@ async def test_session_pool_tool_preserves_disconnect_error_when_eviction_fails( patch("deerflow.mcp.tools.get_paths", return_value=Paths(tmp_path)), pytest.raises(anyio.ClosedResourceError) as exc_info, ): - await wrapped.coroutine(runtime=_legacy_tool_runtime(), value=1) + await wrapped.coroutine(deerflow_runtime=_legacy_tool_runtime(), value=1) assert exc_info.value is error pool.close_session_if_current.assert_awaited_once_with("srv", "test-user-autouse:default", session) @@ -944,7 +989,7 @@ async def test_session_pool_tool_keeps_session_after_tool_error_result(tmp_path) patch("deerflow.mcp.tools.get_paths", return_value=Paths(tmp_path)), pytest.raises(ToolException, match="invalid input"), ): - await wrapped.coroutine(runtime=_legacy_tool_runtime(), value=1) + await wrapped.coroutine(deerflow_runtime=_legacy_tool_runtime(), value=1) pool.close_session_if_current.assert_not_awaited() @@ -968,7 +1013,7 @@ async def test_session_pool_tool_keeps_session_after_interceptor_error(tmp_path) patch("deerflow.mcp.tools.get_paths", return_value=Paths(tmp_path)), pytest.raises(RuntimeError, match="interceptor failed"), ): - await wrapped.coroutine(runtime=_legacy_tool_runtime(), value=1) + await wrapped.coroutine(deerflow_runtime=_legacy_tool_runtime(), value=1) session.call_tool.assert_not_awaited() pool.close_session_if_current.assert_not_awaited() @@ -1031,15 +1076,15 @@ async def test_late_disconnect_from_old_session_does_not_evict_replacement(tmp_p {"transport": "stdio", "command": "x", "args": []}, ) runtime = _legacy_tool_runtime() - first_call = asyncio.create_task(wrapped.coroutine(runtime=runtime, value=1)) - late_call = asyncio.create_task(wrapped.coroutine(runtime=runtime, value=2)) + first_call = asyncio.create_task(wrapped.coroutine(deerflow_runtime=runtime, value=1)) + late_call = asyncio.create_task(wrapped.coroutine(deerflow_runtime=runtime, value=2)) await asyncio.wait_for(both_started.wait(), timeout=1) first_failure.set() with pytest.raises(anyio.ClosedResourceError): await first_call - await wrapped.coroutine(runtime=runtime, value=3) + await wrapped.coroutine(deerflow_runtime=runtime, value=3) late_failure.set() with pytest.raises(anyio.ClosedResourceError): await late_call @@ -1084,7 +1129,7 @@ async def test_session_pool_tool_wrapping(): mock_runtime.context = {"thread_id": "thread-42", "thread_incarnation": "incarnation-1"} mock_runtime.config = {} - await wrapped.coroutine(runtime=mock_runtime, url="https://example.com") + await wrapped.coroutine(deerflow_runtime=mock_runtime, url="https://example.com") mock_session.call_tool.assert_awaited_once_with("navigate", {"url": "https://example.com"}) @@ -1129,7 +1174,7 @@ async def test_session_pool_tool_pins_cwd_and_temp_env(tmp_path): patch("langchain_mcp_adapters.sessions.create_session", return_value=mock_cm) as create_session, ): wrapped = _make_session_pool_tool(original_tool, "playwright", connection) - await wrapped.coroutine(runtime=mock_runtime, url="https://example.com") + await wrapped.coroutine(deerflow_runtime=mock_runtime, url="https://example.com") session_connection = create_session.call_args.args[0] workspace = paths.sandbox_work_dir("thread-42", user_id="user-7") @@ -1185,7 +1230,7 @@ async def test_session_pool_tool_does_not_override_explicit_tmpdir(tmp_path): patch("langchain_mcp_adapters.sessions.create_session", return_value=mock_cm) as create_session, ): wrapped = _make_session_pool_tool(original_tool, "playwright", connection) - await wrapped.coroutine(runtime=mock_runtime, url="https://example.com") + await wrapped.coroutine(deerflow_runtime=mock_runtime, url="https://example.com") session_connection = create_session.call_args.args[0] # Operator-provided TMPDIR is preserved; TMP/TEMP still get our default. @@ -1232,7 +1277,7 @@ async def test_session_pool_tool_does_not_override_explicit_cwd(tmp_path): patch("langchain_mcp_adapters.sessions.create_session", return_value=mock_cm) as create_session, ): wrapped = _make_session_pool_tool(original_tool, "playwright", connection) - await wrapped.coroutine(runtime=mock_runtime, url="https://example.com") + await wrapped.coroutine(deerflow_runtime=mock_runtime, url="https://example.com") session_connection = create_session.call_args.args[0] workspace = paths.sandbox_work_dir("thread-42", user_id="user-7") @@ -1279,7 +1324,7 @@ async def test_session_pool_tool_skips_fs_work_for_non_stdio_transport(tmp_path) patch("langchain_mcp_adapters.sessions.create_session", return_value=mock_cm) as create_session, ): wrapped = _make_session_pool_tool(original_tool, "srv", connection) - await wrapped.coroutine(runtime=mock_runtime, url="https://example.com") + await wrapped.coroutine(deerflow_runtime=mock_runtime, url="https://example.com") session_connection = create_session.call_args.args[0] assert "cwd" not in session_connection @@ -1333,7 +1378,7 @@ async def test_session_pool_tool_skips_after_walk_when_no_text_content(tmp_path) patch("deerflow.mcp.tools._changed_workspace_files") as changed_files, ): wrapped = _make_session_pool_tool(original_tool, "playwright", connection) - await wrapped.coroutine(runtime=mock_runtime, url="https://example.com") + await wrapped.coroutine(deerflow_runtime=mock_runtime, url="https://example.com") changed_files.assert_not_called() @@ -1379,7 +1424,7 @@ async def test_session_pool_tool_runs_after_walk_when_text_content_present(tmp_p patch("deerflow.mcp.tools._changed_workspace_files", return_value=[]) as changed_files, ): wrapped = _make_session_pool_tool(original_tool, "playwright", connection) - await wrapped.coroutine(runtime=mock_runtime, url="https://example.com") + await wrapped.coroutine(deerflow_runtime=mock_runtime, url="https://example.com") changed_files.assert_called_once() @@ -1422,7 +1467,7 @@ async def test_session_pool_tool_forwards_interceptor_headers(): {"transport": "stdio", "command": "x", "args": []}, tool_interceptors=[header_interceptor], ) - await wrapped.coroutine(runtime=_legacy_tool_runtime(), x=1) + await wrapped.coroutine(deerflow_runtime=_legacy_tool_runtime(), x=1) mock_session.call_tool.assert_awaited_once_with("act", {"x": 1}, meta={"headers": {"X-User-Id": "u-42"}}) @@ -1471,7 +1516,7 @@ async def test_session_pool_interceptor_reads_request_scoped_secret(): {"transport": "stdio", "command": "x", "args": []}, tool_interceptors=[secret_header_interceptor], ) - await wrapped.coroutine(runtime=_legacy_tool_runtime(), x=1) + await wrapped.coroutine(deerflow_runtime=_legacy_tool_runtime(), x=1) mock_session.call_tool.assert_awaited_once_with( "act", @@ -1517,7 +1562,7 @@ async def test_session_pool_tool_no_headers_omits_meta(): {"transport": "stdio", "command": "x", "args": []}, tool_interceptors=[passthrough_interceptor], ) - await wrapped.coroutine(runtime=_legacy_tool_runtime(), x=1) + await wrapped.coroutine(deerflow_runtime=_legacy_tool_runtime(), x=1) mock_session.call_tool.assert_awaited_once_with("act", {"x": 1}) @@ -1561,7 +1606,7 @@ async def test_session_pool_tool_ignores_unsupported_header_type(caplog): {"transport": "stdio", "command": "x", "args": []}, tool_interceptors=[invalid_header_interceptor], ) - await wrapped.coroutine(runtime=_legacy_tool_runtime(), x=1) + await wrapped.coroutine(deerflow_runtime=_legacy_tool_runtime(), x=1) mock_session.call_tool.assert_awaited_once_with("act", {"x": 1}) assert "unsupported type" in caplog.text @@ -1599,7 +1644,7 @@ async def test_session_pool_tool_extracts_thread_id(): mock_runtime.context = {"thread_incarnation": "incarnation-1"} mock_runtime.config = {"configurable": {"thread_id": "from-config"}} - await wrapped.coroutine(runtime=mock_runtime, x=1) + await wrapped.coroutine(deerflow_runtime=mock_runtime, x=1) # Verify the session was created with the canonical versioned JSON scope; # the autouse fixture sets the effective user to "test-user-autouse". @@ -1641,7 +1686,7 @@ async def test_session_pool_tool_default_scope(): wrapped = _make_session_pool_tool(original_tool, "server", {"transport": "stdio", "command": "x", "args": []}) # No thread_id in runtime at all. - await wrapped.coroutine(runtime=None, x=1) + await wrapped.coroutine(deerflow_runtime=None, x=1) pool = get_session_pool() assert ("server", "test-user-autouse:default") in {k[:2] for k in pool._entries} @@ -1681,7 +1726,7 @@ async def test_session_pool_tool_get_config_fallback(): wrapped = _make_session_pool_tool(original_tool, "server", {"transport": "stdio", "command": "x", "args": []}) # runtime=None — get_config() fallback should provide thread_id - await wrapped.coroutine(runtime=None, x=1) + await wrapped.coroutine(deerflow_runtime=None, x=1) pool = get_session_pool() assert ("server", "test-user-autouse:from-langgraph-config") in {k[:2] for k in pool._entries} diff --git a/backend/tests/test_mcp_sync_wrapper.py b/backend/tests/test_mcp_sync_wrapper.py index e729f962bf..cd2d807866 100644 --- a/backend/tests/test_mcp_sync_wrapper.py +++ b/backend/tests/test_mcp_sync_wrapper.py @@ -292,16 +292,16 @@ def test_func_patched_mcp_tool_keeps_toolnode_runtime_injection(tmp_path): (tool,) = asyncio.run(get_mcp_tools()) # After func patching, LangGraph must still detect the runtime injection. - assert _get_all_injected_args(tool).runtime == "runtime" + assert _get_all_injected_args(tool).runtime == "deerflow_runtime" # And a real ToolNode must pass the ToolRuntime through to the coroutine. seen: dict[str, object] = {} orig = tool.coroutine - async def spy(runtime: object | None = None, **arguments: object) -> str: - seen["runtime"] = runtime - return await orig(runtime=runtime, **arguments) + async def spy(deerflow_runtime: object | None = None, **arguments: object) -> str: + seen["runtime"] = deerflow_runtime + return await orig(deerflow_runtime=deerflow_runtime, **arguments) tool.coroutine = spy diff --git a/backend/tests/test_mcp_task_tool_wrapping.py b/backend/tests/test_mcp_task_tool_wrapping.py index a300a280df..4ccbc76fb9 100644 --- a/backend/tests/test_mcp_task_tool_wrapping.py +++ b/backend/tests/test_mcp_task_tool_wrapping.py @@ -3,7 +3,10 @@ from types import SimpleNamespace from unittest.mock import patch import pytest +from langchain_core.messages import AIMessage from langchain_core.tools import StructuredTool +from langgraph.graph import END, START, MessagesState, StateGraph +from langgraph.prebuilt import ToolNode from pydantic import BaseModel from deerflow.config.extensions_config import ExtensionsConfig, McpServerConfig @@ -148,7 +151,7 @@ async def test_submit_wrapper_persists_before_returning_only_the_local_handle(co tool_call_id="call-1", ) - result = await submit_tool.coroutine(runtime=runtime, topic="MCP") + result = await submit_tool.coroutine(deerflow_runtime=runtime, topic="MCP") assert result == { "task_id": "mcp-task-local-1", @@ -174,6 +177,53 @@ async def test_submit_wrapper_persists_before_returning_only_the_local_handle(co set_mcp_task_submitter(None) +@pytest.mark.asyncio +async def test_submit_wrapper_keeps_a_business_argument_named_runtime() -> None: + """A tool whose own schema has a ``runtime`` property must receive its value.""" + submitter = FakeSubmitter() + set_mcp_task_submitter(submitter) + try: + + async def call(**arguments): + return arguments + + def tool(name: str) -> StructuredTool: + return StructuredTool( + name=name, + description=name, + args_schema={ + "type": "object", + "properties": {"runtime": {"type": "string"}, "code": {"type": "string"}}, + "required": ["runtime", "code"], + }, + coroutine=call, + ) + + configured = _configure_task_tools_for_server( + [tool("submit_report"), tool("get_report_status"), tool("cancel_report")], + server_name="reports", + server_config=_server_config(), + tool_name_prefix=False, + ) + builder = StateGraph(MessagesState, context_schema=dict) + builder.add_node("tools", ToolNode([configured[0]])) + builder.add_edge(START, "tools") + builder.add_edge("tools", END) + graph = builder.compile() + + await graph.ainvoke( + {"messages": [AIMessage(content="", tool_calls=[{"name": "submit_report", "args": {"runtime": "python", "code": "print(1)"}, "id": "call-1", "type": "tool_call"}])]}, + context={"thread_id": "thread-1", "thread_incarnation": "incarnation-1", "run_id": "run-1"}, + ) + + request = submitter.calls[0]["request"] + assert request.arguments == {"runtime": "python", "code": "print(1)"} + assert request.thread_id == "thread-1" + assert request.tool_call_id == "call-1" + finally: + set_mcp_task_submitter(None) + + @pytest.mark.asyncio async def test_submit_wrapper_fails_clearly_without_gateway_task_runtime() -> None: set_mcp_task_submitter(None) diff --git a/backend/tests/test_personal_mcp.py b/backend/tests/test_personal_mcp.py index 0f237a5387..75f5f4e034 100644 --- a/backend/tests/test_personal_mcp.py +++ b/backend/tests/test_personal_mcp.py @@ -480,6 +480,54 @@ async def test_personal_tool_guard_reuses_validation_and_rejects_file_replacemen reset_current_user(identity) +@pytest.mark.asyncio +@pytest.mark.parametrize("parameter", ["runtime", "deerflow_runtime"]) +async def test_personal_tool_guard_reads_the_injected_runtime_for_either_transport(personal_client, parameter): + """HTTP/SSE tools keep the adapter's ``runtime``; pooled stdio uses ``deerflow_runtime``.""" + from typing import Annotated + + from langchain_core.messages import AIMessage + from langchain_core.tools import InjectedToolArg, StructuredTool + from langgraph.graph import END, START, MessagesState, StateGraph + from langgraph.prebuilt import ToolNode + + from deerflow.mcp.user_tools import _guard + from deerflow.tools.types import Runtime + + assert create(personal_client, "alice").status_code == 200 + server_name = next(iter(load_user_mcp_config("alice").mcp_servers)) + calls = [] + + async def invoke_runtime(runtime: Annotated[Runtime | None, InjectedToolArg()] = None, **arguments): + calls.append(arguments) + return "ok" + + async def invoke_deerflow_runtime(deerflow_runtime: Runtime | None = None, **arguments): + calls.append(arguments) + return "ok" + + coroutine = invoke_runtime if parameter == "runtime" else invoke_deerflow_runtime + tool = _guard( + StructuredTool(name="personal_test", description="Personal test", args_schema={"type": "object", "properties": {"runtime": {"type": "string"}}}, coroutine=coroutine), + "alice", + server_name, + ) + builder = StateGraph(MessagesState, context_schema=dict) + builder.add_node("tools", ToolNode([tool])) + builder.add_edge(START, "tools") + builder.add_edge("tools", END) + graph = builder.compile() + message = AIMessage(content="", tool_calls=[{"name": "personal_test", "args": {"runtime": "python"}, "id": "call-1", "type": "tool_call"}]) + + # No ambient user: only the injected runtime says who is calling. + result = await graph.ainvoke({"messages": [message]}, context={"user_id": "alice"}) + assert result["messages"][-1].content == "ok" + assert calls == [{"runtime": "python"} if parameter == "deerflow_runtime" else {}] + + with pytest.raises(ToolException, match="belongs to another user"): + await graph.ainvoke({"messages": [message]}, context={"user_id": "bob"}) + + @pytest.mark.asyncio async def test_background_calls_resolve_only_persisted_task_owner(personal_client, monkeypatch): import deerflow.mcp.user_config as user_config @@ -636,7 +684,7 @@ async def test_gateway_registers_driver_for_personal_only_task_toolsets(personal identity = set_current_user(SimpleNamespace(id="alice")) try: result = await tools[0].coroutine( - runtime=SimpleNamespace(context={"thread_id": "thread-1", "thread_incarnation": "incarnation-1", "run_id": "run-1"}, config={}, tool_call_id="call-1"), + deerflow_runtime=SimpleNamespace(context={"thread_id": "thread-1", "thread_incarnation": "incarnation-1", "run_id": "run-1"}, config={}, tool_call_id="call-1"), topic="MCP", ) finally: