Skip to content

fix(mcp): keep ToolRuntime injection for sync-wrapped MCP tools - #5164

Merged
WillemJiang merged 5 commits into
bytedance:mainfrom
BlueX888:fix/sync-wrapper-preserve-annotations
Sep 4, 2026
Merged

WillemJiang merged 5 commits into
bytedance:mainfrom
BlueX888:fix/sync-wrapper-preserve-annotations

Conversation

@BlueX888

@BlueX888 BlueX888 commented Sep 3, 2026

Copy link
Copy Markdown
Contributor

Why

LangGraph's ToolNode decides which parameters to inject a ToolRuntime into via _get_all_injected_args, which reads type hints from tool.func first and falls back to tool.coroutine. make_sync_tool_wrapper attaches an annotation-less wrapper as tool.func, and get_mcp_tools (plus _ensure_sync_invocable_tool) patch every async-only tool with it. Because func is now truthy, get_type_hints(func) returns {} and the runtime parameter is no longer detected.

Every MCP tool in a sync agent caller then runs with runtime=None:

  • resolve_runtime_user_id(None) skips the runtime.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_tool builds TaskSubmitRequest with run_id=None, tool_call_id=None (thread_id survives only via the get_config() fallback). Downstream, mcp_tasks/service.py skips the source-run lookup, so the completion-notification run launches with assistant_id=None — i.e. under the default lead agent instead of the thread's custom agent.

What changed

make_sync_tool_wrapper now wraps the inner runner and both sync_wrapper variants with functools.wraps(coro). get_type_hints then sees the original annotations, so the InjectedToolArg-annotated runtime parameter is detected again and ToolNode continues to inject it. No signature is synthesized, so the existing RunnableConfig passthrough and the "no dynamic signature" constraint are preserved.

Surface area

  • Frontend UI
  • Backend API
  • Agents / LangGraph — sync tool wrapper used by get_mcp_tools / _ensure_sync_invocable_tool for MCP tools
  • Sandbox
  • Skills
  • Dependencies
  • Default behavior change
  • Docs / tests / CI only

Bug fix verification

  • Test path: backend/tests/test_mcp_sync_wrapper.py::test_funopatched_mcp_tool_keeps_toolnode_runtime_injection
  • Did it go red on main and green on this branch? yes — on main _get_all_injected_args(tool).runtime == None and the tool runs with runtime=None; with the fix the ToolNode injects the runtime and runtime.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 passed
  • uv run ruff format --check ... → already formatted
  • pytest tests/test_mcp_sync_wrapper.py → 8 passed
  • Broader set (test_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 passed

AI 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_args reading func type hints), wrote the functools.wraps fix, and added the regression test that drives a func-patched pooled MCP tool through a real ToolNode. I read the diff, ran the tests above, and can explain each line.

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.
@github-actions github-actions Bot added area:backend Gateway / runtime / core backend under backend/ risk:medium Medium risk: regular code changes size/M PR changes 100-300 lines labels Sep 3, 2026
@CLAassistant

CLAassistant commented Sep 3, 2026 •

Copy link
Copy Markdown

CLA assistant check
All committers have signed the CLA.

@WillemJiang

Copy link
Copy Markdown
Collaborator

@BlueX888 thanks for your contribution. Please click the CLA assitant button to sign the CLA before we start the review process.

@willem-bd willem-bd left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread backend/tests/test_mcp_sync_wrapper.py Outdated
assert mock_log_error.call_args[0][1] == "error_tool"


def test_funopatched_mcp_tool_keeps_toolnode_runtime_injection(tmp_path):

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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"

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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
@WillemJiang WillemJiang added this to the 2.1.0 milestone Sep 4, 2026
@WillemJiang
WillemJiang merged commit dbe11dc into bytedance:main Sep 4, 2026
16 checks passed
@BlueX888
BlueX888 deleted the fix/sync-wrapper-preserve-annotations branch September 4, 2026 23:08
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:backend Gateway / runtime / core backend under backend/ risk:medium Medium risk: regular code changes size/M PR changes 100-300 lines

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants