diff --git a/backend/packages/harness/deerflow/tools/sync.py b/backend/packages/harness/deerflow/tools/sync.py index 7521dd7b3..39fb39de2 100644 --- a/backend/packages/harness/deerflow/tools/sync.py +++ b/backend/packages/harness/deerflow/tools/sync.py @@ -58,6 +58,20 @@ def make_sync_tool_wrapper(coro: Callable[..., Any], tool_name: str) -> Callable else, such as ``run_config``, may collide with LangChain's injected ``config`` argument. Rename that user-facing field or extend this helper before using that signature. + + The returned wrapper is built with ``functools.wraps(coro)``, so it + copies ``__name__``, ``__qualname__``, ``__doc__``, ``__annotations__`` + and ``__dict__`` and sets ``__wrapped__``. This matters for LangGraph: + ``ToolNode._get_all_injected_args`` reads ``get_type_hints(tool.func)``, + and ``get_type_hints`` follows ``__wrapped__`` back to the coroutine's + own ``__globals__``. That is why injected parameters such as + ``runtime: Runtime`` are still detected here even when the caller was + compiled with ``from __future__ import annotations`` (mcp/tools.py, + skill_manage_tool.py) and the annotation is a string. Do not replace + ``wraps`` with a bare wrapper or a hand-copied signature: the string + annotation would then be evaluated against the wrapper's module globals, + or the ``runtime`` parameter would disappear entirely and the tool would + run with ``runtime=None``. """ config_param = _get_runnable_config_param(coro) @@ -79,6 +93,7 @@ def make_sync_tool_wrapper(coro: Callable[..., Any], tool_name: str) -> Callable if config_param: + @functools.wraps(coro) def sync_wrapper(*args: Any, config: RunnableConfig = None, **kwargs: Any) -> Any: if config is not None or config_param not in kwargs: kwargs[config_param] = config @@ -86,6 +101,7 @@ def make_sync_tool_wrapper(coro: Callable[..., Any], tool_name: str) -> Callable return sync_wrapper + @functools.wraps(coro) def sync_wrapper(*args: Any, **kwargs: Any) -> Any: return run_coroutine(*args, **kwargs) diff --git a/backend/tests/test_mcp_sync_wrapper.py b/backend/tests/test_mcp_sync_wrapper.py index b5f253b2f..89550a56d 100644 --- a/backend/tests/test_mcp_sync_wrapper.py +++ b/backend/tests/test_mcp_sync_wrapper.py @@ -3,8 +3,9 @@ import contextvars from unittest.mock import AsyncMock, MagicMock, patch import pytest +from langchain_core.messages import AIMessage from langchain_core.runnables import RunnableConfig -from langchain_core.tools import StructuredTool +from langchain_core.tools import InjectedToolArg, StructuredTool from pydantic import BaseModel, Field from deerflow.mcp.tools import get_mcp_tools @@ -185,3 +186,142 @@ def test_mcp_tool_sync_wrapper_exception_logging(): mock_log_error.assert_called_once() # Verify the tool name is in the log message assert mock_log_error.call_args[0][1] == "error_tool" + + +def test_func_patched_mcp_tool_keeps_toolnode_runtime_injection(tmp_path): + """The sync wrapper must not erase the coroutine's annotations, otherwise + LangGraph's ToolNode stops injecting the ToolRuntime into MCP tools. + + Regression test for the func patching in ``get_mcp_tools`` (and + ``_ensure_sync_invocable_tool``): the wrapper is attached after the MCP + adapter produced a coroutine whose ``runtime`` parameter carries an + ``InjectedToolArg`` annotation. Without wrapping via ``functools.wraps`` the + annotation is dropped, ``_get_all_injected_args`` returns + ``runtime_arg = None``, and ``tool.func`` runs with ``runtime=None``. + Through a real ``ToolNode`` this surfaces as ``resolve_runtime_user_id`` + returning the default user and, for the background-submit wrapper, + ``run_id``/``tool_call_id`` both hitting the ``None`` branch. + """ + from typing import Annotated + + from langgraph.graph import END, START, MessagesState, StateGraph + from langgraph.prebuilt import ToolNode + from langgraph.prebuilt.tool_node import _get_all_injected_args + from mcp.types import CallToolResult, TextContent + + from deerflow.mcp.tools import get_mcp_tools + + # Adapter-shaped coroutine: same signature langchain_mcp_adapters produces. + async def adapter_coro( + runtime: Annotated[object | None, InjectedToolArg()] = None, + **arguments: object, + ) -> str: + return "ok" + + discovered = StructuredTool( + name="mcp_server_navigate", + description="d", + args_schema={"type": "object", "properties": {"url": {"type": "string"}}, "required": ["url"]}, + coroutine=adapter_coro, + response_format="content_and_artifact", + ) + client = MagicMock() + client.get_tools = AsyncMock(return_value=[discovered]) + client.tool_interceptors = [] + client.callbacks = None + cfg = MagicMock() + cfg.mcp_servers = {} + + with ( + patch("langchain_mcp_adapters.client.MultiServerMCPClient", return_value=client), + patch("deerflow.config.extensions_config.ExtensionsConfig.from_file", return_value=cfg), + patch("deerflow.mcp.tools.validate_mcp_task_config_snapshot"), + patch( + "deerflow.mcp.tools.build_servers_config", + return_value={"pw": {"transport": "stdio", "command": "x", "args": []}}, + ), + patch("deerflow.mcp.tools.get_initial_oauth_headers", new_callable=AsyncMock, return_value={}), + patch("deerflow.mcp.tools.build_mcp_tool_interceptors", return_value=[]), + ): + from deerflow.mcp.session_pool import reset_session_pool + + reset_session_pool() + (tool,) = asyncio.run(get_mcp_tools()) + + # After func patching, LangGraph must still detect the runtime injection. + assert _get_all_injected_args(tool).runtime == "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) + + tool.coroutine = spy + + class FakeSession: + async def call_tool(self, *args: object, **kwargs: object): + return CallToolResult(content=[TextContent(type="text", text="ok")], isError=False) + + with ( + patch("deerflow.mcp.tools.get_paths") as gp, + patch( + "deerflow.mcp.tools.call_pooled_session_tool", + new=AsyncMock(return_value=CallToolResult(content=[TextContent(type="text", text="ok")], isError=False)), + ), + patch("deerflow.mcp.session_pool.MCPSessionPool.get_session", new=AsyncMock(return_value=FakeSession())), + ): + gp.return_value.ensure_thread_dirs = lambda *a, **k: None + gp.return_value.sandbox_work_dir = lambda *a, **k: tmp_path + gp.return_value.sandbox_user_data_dir = lambda *a, **k: tmp_path + graph = StateGraph(MessagesState, context_schema=dict) + graph.add_node("tools", ToolNode([tool])) + graph.add_edge(START, "tools") + graph.add_edge("tools", END) + ai = AIMessage( + content="", + tool_calls=[{"name": "mcp_server_navigate", "args": {"url": "x"}, "id": "c1", "type": "tool_call"}], + ) + asyncio.run( + graph.compile().ainvoke( + {"messages": [ai]}, + config={"configurable": {"thread_id": "T"}}, + context={"thread_id": "T", "run_id": "run-1", "user_id": "alice"}, + ) + ) + + assert seen["runtime"] is not None, "ToolNode did not inject runtime into the func-patched MCP tool" + assert seen["runtime"].context["user_id"] == "alice" + + +def test_sync_wrapped_builtin_tools_still_resolve_runtime(): + """Built-in tools expose ``runtime`` as a pydantic schema field, so + ``_get_all_injected_args`` detects it from the input schema rather than + from the wrapper's annotations. Wrapping their ``func`` with + ``make_sync_tool_wrapper`` must keep ``runtime`` resolved, otherwise a + sync-only caller would run them with ``runtime=None``. + + This pins the wrapper against the built-ins for the case that is easy to + break: one of them later drops ``runtime`` from its schema *and* a wrapper + refactor removes ``get_type_hints`` propagation. The MCP test above covers + the adapter tools, which carry no ``runtime`` schema field and therefore + depend purely on the wrapped coroutine annotations. + """ + import copy + + from langgraph.prebuilt.tool_node import _get_all_injected_args + + from deerflow.tools.builtins.background_tasks_tool import ( + cancel_background_task, + list_background_tasks, + ) + from deerflow.tools.builtins.batch_task_tool import batch_status, cancel_batch + + for tool in (list_background_tasks, cancel_background_task, batch_status, cancel_batch): + patched = copy.copy(tool) + # _ensure_sync_invocable_tool does exactly this to async-only tools. + patched.func = make_sync_tool_wrapper(patched.coroutine, patched.name) + assert _get_all_injected_args(patched).runtime == "runtime", f"sync wrapper dropped runtime resolution for built-in tool {patched.name}"