Repository navigation
fix(mcp): keep ToolRuntime injection for sync-wrapped MCP tools - #5164
Conversation
make_sync_tool_wrapper attached an annotation-less wrapper to tool.func, which made LangGraph's ToolNode stop detecting the coroutine's "runtime" parameter (_get_all_injected_args falls back to func first and its type hints are empty). Every MCP tool in a sync agent caller then ran with runtime=None: resolve_runtime_user_id fell through to the default user, and the background-submit wrapper lost run_id/tool_call_id on the TaskSubmitRequest, so completion notifications launched under the default lead agent instead of the thread's agent. Wrap the generator and both sync_wrapper variants with functools.wraps so get_type_hints still sees the original annotations. Adds a regression test that drives a func-patched pooled MCP tool through a real ToolNode and asserts the ToolRuntime is injected with the thread's user context. It fails on main (runtime=None) and passes with the fix.
|
@BlueX888 thanks for your contribution. Please click the CLA assitant button to sign the CLA before we start the review process. |
willem-bd
left a comment
There was a problem hiding this comment.
Verified against head 3796d4d. The root cause holds: langgraph 1.2.9 _get_all_injected_args reads tool.func first, so the annotation-less wrapper left MCP tools running with runtime=None. I also confirmed the fix is safe for PEP 563 string annotations — typing.get_type_hints follows __wrapped__ back to the coroutine's __globals__, so from __future__ import annotations modules (mcp/tools.py, skill_manage_tool.py, and the other _ensure_sync_invocable_tool clients) still resolve Runtime correctly — and that no in-repo code introspects tool.func signatures or dunder metadata. Reproduced the red/green claim: the new test fails against unfixed sync.py (runtime=None) and passes at head (8 passed), plus 61 tests across the other sync-wrapped built-in tools pass. No correctness issues found; comments are documentation, coverage and a test-name nit.
| """ | ||
| config_param = _get_runnable_config_param(coro) | ||
|
|
||
| @functools.wraps(coro) |
There was a problem hiding this comment.
Fix looks right. Worth extending the Notes block above to record the new contract: functools.wraps copies __name__, __qualname__, __doc__, __annotations__ and __dict__, and sets __wrapped__. Two consequences worth stating there — (1) typing.get_type_hints(wrapper) follows __wrapped__, so string annotations from modules compiled with from __future__ import annotations still resolve against the original coroutine's module globals; that is exactly why this works for mcp/tools.py and skill_manage_tool.py, where runtime: Runtime is a string annotation, and (2) anything introspecting tool.func now sees the coroutine signature instead of (*args, **kwargs). Separately, the wraps on the inner run_coroutine has no effect on the fix (only the outer sync_wrapper becomes tool.func), so it could be dropped to keep the wrap surface minimal.
There was a problem hiding this comment.
Done — expanded the Notes block and dropped the no-op wraps on the inner run_coroutine. Wording now states the functools.wraps contract (copies __name__/__qualname__/__doc__/__annotations__/__dict__, sets __wrapped__) and the reason it keeps get_type_hints resolving string annotations for PEP 563 callers (mcp/tools.py, skill_manage_tool.py). Also noted not to replace wraps with a bare wrapper. Pushed as part of d92f2b8.
| assert mock_log_error.call_args[0][1] == "error_tool" | ||
|
|
||
|
|
||
| def test_funopatched_mcp_tool_keeps_toolnode_runtime_injection(tmp_path): |
There was a problem hiding this comment.
Name typo: funopatched reads like a mangled func-patched. Since this becomes the canonical regression-test ID, test_func_patched_mcp_tool_keeps_toolnode_runtime_injection would be clearer.
There was a problem hiding this comment.
Renamed to test_func_patched_mcp_tool_keeps_toolnode_runtime_injection.
| (tool,) = asyncio.run(get_mcp_tools()) | ||
|
|
||
| # After func patching, LangGraph must still detect the runtime injection. | ||
| assert _get_all_injected_args(tool).runtime == "runtime" |
There was a problem hiding this comment.
This pins the MCP/session-pool path, but the same annotation erasure affected every tool patched by _ensure_sync_invocable_tool: skill_manage (its func wraps _skill_manage_impl), batch_status, cancel_batch, list_background_tasks, cancel_background_task, review_skill_package, update_agent and list_uploaded_files are all async-only with a runtime: Runtime parameter, so they ran with runtime=None in sync callers too. I verified at this head that _get_all_injected_args(skill_manage_tool).runtime == 'runtime', i.e. the fix does cover them, but nothing pins it — a one-line assertion (or a small loop over those impls) would stop a future wrapper refactor from silently regressing per-user resolution for the built-ins.
There was a problem hiding this comment.
Thanks for the careful read. On the built-in list: I verified at this head that all the built-ins you named (skill_manage, batch_status, cancel_batch, list_background_tasks, cancel_background_task, list_uploaded_files, update_agent, review_skill_package) expose runtime as a `get_input_schema()` field, so `_get_all_injected_args` resolves it from the schema channel and they were NOT affected by this bug — even a bare annotation-less wrapper on `func` leaves them at `runtime="runtime"`. The affected channel is only the MCP adapter tools, whose `args_schema` is a JSON dict with no `runtime` field and which therefore depend purely on the wrapped coroutine annotations (covered by the MCP regression test). I added `test_sync_wrapped_builtin_tools_still_resolve_runtime` as a pin that `make_sync_tool_wrapper` keeps `runtime` resolved for the built-ins too, so a future refactor cannot regress per-user resolution — but its docstring says what it actually guards. If you would prefer I also add an assertion that a bare wrapper drops the MCP path, that is already the red side of the MCP test.
… coverage Address review feedback on bytedance#5164: - Expand the Notes block in make_sync_tool_wrapper to state the functools.wraps contract (copies __name__/__qualname__/__doc__/__annotations__/__dict__ and sets __wrapped__) and why that is what keeps get_type_hints resolving string annotations from callers like mcp/tools.py and skill_manage_tool.py. Drop the no-op wraps on the inner run_coroutine so the wrap surface stays minimal. - Rename the regression test to test_func_patched_mcp_tool_keeps_toolnode_runtime_injection. - Add test_sync_wrapped_builtin_tools_still_resolve_runtime to pin that the built-in tools (which carry runtime as a pydantic schema field) keep resolving runtime after their func is wrapped by make_sync_tool_wrapper, so a future wrapper refactor cannot silently regress per-user resolution for them.
…ons' into fix/sync-wrapper-preserve-annotations
Why
LangGraph's
ToolNodedecides which parameters to inject aToolRuntimeinto via_get_all_injected_args, which reads type hints fromtool.funcfirst and falls back totool.coroutine.make_sync_tool_wrapperattaches an annotation-less wrapper astool.func, andget_mcp_tools(plus_ensure_sync_invocable_tool) patch every async-only tool with it. Becausefuncis now truthy,get_type_hints(func)returns{}and theruntimeparameter is no longer detected.Every MCP tool in a sync agent caller then runs with
runtime=None:resolve_runtime_user_id(None)skips theruntime.context["user_id"]channel and falls through to the contextvar/default, so per-user scope (session pool keys, workspace paths, user-scoped auth) resolves to the wrong identity in paths that only receive the injected runtime._make_background_submit_toolbuildsTaskSubmitRequestwithrun_id=None, tool_call_id=None(thread_id survives only via theget_config()fallback). Downstream,mcp_tasks/service.pyskips the source-run lookup, so the completion-notification run launches withassistant_id=None— i.e. under the default lead agent instead of the thread's custom agent.What changed
make_sync_tool_wrappernow wraps the inner runner and bothsync_wrappervariants withfunctools.wraps(coro).get_type_hintsthen sees the original annotations, so theInjectedToolArg-annotatedruntimeparameter is detected again andToolNodecontinues to inject it. No signature is synthesized, so the existingRunnableConfigpassthrough and the "no dynamic signature" constraint are preserved.Surface area
get_mcp_tools/_ensure_sync_invocable_toolfor MCP toolsBug fix verification
backend/tests/test_mcp_sync_wrapper.py::test_funopatched_mcp_tool_keeps_toolnode_runtime_injectionmainand green on this branch? yes — onmain_get_all_injected_args(tool).runtime == Noneand the tool runs withruntime=None; with the fix theToolNodeinjects the runtime andruntime.context["user_id"]is the thread's user.Validation
cd backend:uv run ruff check packages/harness/deerflow/tools/sync.py tests/test_mcp_sync_wrapper.py→ All checks passeduv run ruff format --check ...→ already formattedpytest tests/test_mcp_sync_wrapper.py→ 8 passedtest_mcp_sync_wrapper, test_mcp_session_pool, test_mcp_task_tool_wrapping, test_mcp_context_headers, test_mcp_custom_interceptors, test_mcp_user_scoped_auth, test_mcp_oauth, test_mcp_task_models) → 198 passedAI assistance
Tool(s) used: Claude Code
How you used it: Claude Code traced the root cause against the pinned langgraph 1.2.9 (
_get_all_injected_argsreadingfunctype hints), wrote thefunctools.wrapsfix, and added the regression test that drives a func-patched pooled MCP tool through a realToolNode. I read the diff, ran the tests above, and can explain each line.