feat(langgraph): manage standalone ToolNode calls (#902)

#### Overview

Add managed NeMo Relay execution for standalone LangGraph `ToolNode` calls.

- [x] I confirm this contribution is my own work, or I have the right to submit it under this project's license.
- [x] I searched existing issues and open pull requests, and this does not duplicate existing work.

#### Details

- Add `create_tool_node`, `wrap_tool_call`, and `awrap_tool_call` to the LangGraph integration.
- Preserve LangGraph argument injection, parallel execution, error handling, and `Command` results while routing tool calls through Relay.
- Add focused integration coverage and standalone ToolNode documentation.

#### Where should the reviewer start?

Start with `python/nemo_relay/integrations/langgraph/tool_node.py`; the tests demonstrate the sync, async, interception, injection, parallel, command, and error paths.

#### Related Issues: (use one of the action keywords Closes / Fixes / Resolves / Relates to)

- Relates to: none


## Summary by CodeRabbit

* **New Features**
  * Added managed LangGraph `ToolNode` integration for synchronous and asynchronous tool calls.
  * Added helpers for creating and wrapping tool nodes with managed execution.
  * Preserved tool-call metadata, command results, nested tool messages, and graph interruptions.
  * Added support for tool-error handling and callback instrumentation with `NemoRelayCallbackHandler`.

* **Documentation**
  * Added examples covering tool registration, graph wiring, invocation, callbacks, and direct tool composition.

Authors:
  - Will Killian (https://github.com/willkill07)

Approvers:
  - Eric Evans II (https://github.com/ericevans-nv)

URL: #902

Author: willkill07

docs/supported-integrations/langgraph.mdx

@@ -64,6 +64,53 @@ with nemo_relay.scope.scope("langgraph-request", nemo_relay.ScopeType.Agent):
 print(result)
 ```
 
+## Managed ToolNode Calls
+
+For a standalone LangGraph `ToolNode`, construct the node with
+`create_tool_node`. It routes each model-requested tool call through Relay's
+managed execution pipeline while preserving LangGraph state, store, and runtime
+argument injection.
+
+```python
+from langchain_core.messages import AIMessage
+from langchain_core.tools import tool
+from langgraph.graph import END, START, MessagesState, StateGraph
+from nemo_relay.integrations.langgraph import (
+    NemoRelayCallbackHandler,
+    create_tool_node,
+)
+
+@tool
+def get_weather(location: str) -> str:
+    """Return the weather for a location."""
+    return f"Sunny in {location}"
+
+builder = StateGraph(MessagesState)
+builder.add_node("tools", create_tool_node([get_weather]))
+builder.add_edge(START, "tools")
+builder.add_edge("tools", END)
+graph = builder.compile()
+
+with nemo_relay.scope.scope("langgraph-request", nemo_relay.ScopeType.Agent):
+    result = graph.invoke(
+        {
+            "messages": [
+                AIMessage(
+                    content="",
+                    tool_calls=[{"name": "get_weather", "args": {"location": "Boston"}, "id": "call-1"}],
+                )
+            ]
+        },
+        config={"callbacks": [NemoRelayCallbackHandler()]},
+    )
+```
+
+`NemoRelayCallbackHandler` records the graph lifecycle and provides its scopes.
+`create_tool_node` provides managed tool execution; use both for complete
+standalone LangGraph instrumentation. For custom ToolNode wrapper composition,
+construct `ToolNode` directly and pass `wrap_tool_call` and `awrap_tool_call`
+from this integration.
+
 For LangChain agents inside a LangGraph workflow, use `NemoRelayMiddleware` from
 this package the same way as the LangChain integration and pass the LangGraph
 `config` into the nested agent call:

python/nemo_relay/integrations/langgraph/__init__.py

@@ -5,8 +5,12 @@
 
 from nemo_relay.integrations.langchain import NemoRelayMiddleware
 from nemo_relay.integrations.langgraph.callbacks import NemoRelayCallbackHandler
+from nemo_relay.integrations.langgraph.tool_node import awrap_tool_call, create_tool_node, wrap_tool_call
 
 __all__ = [
     "NemoRelayCallbackHandler",
     "NemoRelayMiddleware",
+    "awrap_tool_call",
+    "create_tool_node",
+    "wrap_tool_call",
 ]

python/nemo_relay/integrations/langgraph/tool_node.py

@@ -0,0 +1,225 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Managed NeMo Relay wrappers for standalone LangGraph ``ToolNode`` objects."""
+
+from __future__ import annotations
+
+from collections.abc import Awaitable, Callable, Sequence
+from dataclasses import replace
+from typing import Any, cast
+
+from langchain_core.messages import ToolMessage
+from langchain_core.tools import BaseTool
+from langgraph.errors import GraphBubbleUp
+from langgraph.prebuilt import ToolNode
+from langgraph.prebuilt.tool_node import ToolCallRequest, ToolInvocationError
+from langgraph.types import Command
+
+import nemo_relay
+from nemo_relay.typed import Codec
+from nemo_relay.utils import run_sync
+
+
+class _ToolNodeResultCodec(nemo_relay.typed.BestEffortAnyCodec):
+    """Restore ToolMessages nested in LangGraph Command updates."""
+
+    _LIST_TAG = "__nemo_relay_langgraph_list__"
+
+    def to_json(self, value: object) -> nemo_relay.Json:
+        if isinstance(value, list):
+            return {self._LIST_TAG: [self.to_json(item) for item in value]}
+        return super().to_json(value)
+
+    def from_json(self, data: nemo_relay.Json) -> object:
+        if isinstance(data, dict) and isinstance(data.get(self._LIST_TAG), list):
+            return [self.from_json(item) for item in data[self._LIST_TAG]]
+
+        result = super().from_json(data)
+        if not isinstance(result, Command):
+            return result
+
+        if isinstance(result.update, dict):
+            messages = result.update.get("messages")
+            if isinstance(messages, list):
+                result.update["messages"] = _restore_tool_messages(messages)
+        elif isinstance(result.update, list):
+            return replace(result, update=_restore_tool_messages(result.update))
+        return result
+
+
+def _restore_tool_messages(messages: list[object]) -> list[object]:
+    """Reconstruct serialized ToolMessages in a LangGraph command update."""
+    return [ToolMessage.model_validate(message) if isinstance(message, dict) else message for message in messages]
+
+
+_DEFAULT_HANDLE_TOOL_ERRORS = object()
+_TOOL_CALL_ERROR_TEMPLATE = "Error: {error}\n Please fix your mistakes."
+_GRAPH_BUBBLE_RESULT = {"__nemo_relay_langgraph_graph_bubble_up__": True}
+
+
+def _handle_tool_error(error: Exception, policy: object) -> str:
+    """Preserve ToolNode error handling while allowing graph bubbles to escape."""
+    if isinstance(error, GraphBubbleUp):
+        raise error
+    if policy is _DEFAULT_HANDLE_TOOL_ERRORS:
+        if isinstance(error, ToolInvocationError):
+            return error.message
+        raise error
+    if isinstance(policy, tuple):
+        if not isinstance(error, policy):
+            raise error
+        return _TOOL_CALL_ERROR_TEMPLATE.format(error=repr(error))
+    if isinstance(policy, type) and issubclass(policy, Exception):
+        if not isinstance(error, policy):
+            raise error
+        return _TOOL_CALL_ERROR_TEMPLATE.format(error=repr(error))
+    if policy is True:
+        return _TOOL_CALL_ERROR_TEMPLATE.format(error=repr(error))
+    if isinstance(policy, str):
+        return policy
+    if callable(policy):
+        return cast(Callable[[Exception], str], policy)(error)
+    raise ValueError(f"unexpected handle_tool_errors value: {policy}")
+
+
+def _tool_details(request: ToolCallRequest) -> tuple[nemo_relay.ScopeHandle, str, dict[str, Any], str | None]:
+    """Extract the model-controlled tool-call fields managed by Relay."""
+    return (
+        nemo_relay.scope.get_handle(),
+        request.tool_call["name"],
+        request.tool_call.get("args") or {},
+        request.tool_call.get("id"),
+    )
+
+
+def wrap_tool_call(
+    request: ToolCallRequest,
+    execute: Callable[[ToolCallRequest], ToolMessage | Command[Any]],
+) -> ToolMessage | Command[Any]:
+    """Run one synchronous LangGraph tool call through NeMo Relay.
+
+    Args:
+        request: LangGraph's tool-call request.
+        execute: LangGraph callback that invokes the requested tool.
+
+    Returns:
+        The LangGraph tool result after managed Relay execution.
+    """
+    parent, tool_name, tool_args, tool_call_id = _tool_details(request)
+    args_codec = cast(Codec[dict[str, Any]], nemo_relay.typed.BestEffortAnyCodec())
+    result_codec = cast(Codec[ToolMessage | Command[Any] | dict[str, bool]], _ToolNodeResultCodec())
+    graph_bubble: GraphBubbleUp | None = None
+
+    def _call(args: dict[str, Any]) -> nemo_relay.ToolExecutionResult[ToolMessage | Command[Any] | dict[str, bool]]:
+        nonlocal graph_bubble
+        try:
+            result = execute(request.override(tool_call={**request.tool_call, "args": args}))
+        except GraphBubbleUp as error:
+            # Relay's native callback boundary cannot propagate arbitrary Python
+            # exceptions. Preserve the original graph bubble locally and re-raise
+            # it after managed execution returns.
+            graph_bubble = error
+            result = _GRAPH_BUBBLE_RESULT
+        return nemo_relay.ToolExecutionResult(result)
+
+    outcome = run_sync(
+        nemo_relay.typed.tool_execute(
+            name=tool_name,
+            args=tool_args,
+            func=_call,
+            args_codec=args_codec,
+            result_codec=result_codec,
+            handle=parent,
+            tool_call_id=tool_call_id,
+        )
+    )
+    if graph_bubble is not None:
+        raise graph_bubble
+    return cast(ToolMessage | Command[Any], outcome.result)
+
+
+async def awrap_tool_call(
+    request: ToolCallRequest,
+    execute: Callable[[ToolCallRequest], Awaitable[ToolMessage | Command[Any]]],
+) -> ToolMessage | Command[Any]:
+    """Run one asynchronous LangGraph tool call through NeMo Relay.
+
+    Args:
+        request: LangGraph's tool-call request.
+        execute: Async LangGraph callback that invokes the requested tool.
+
+    Returns:
+        The LangGraph tool result after managed Relay execution.
+    """
+    parent, tool_name, tool_args, tool_call_id = _tool_details(request)
+    args_codec = cast(Codec[dict[str, Any]], nemo_relay.typed.BestEffortAnyCodec())
+    result_codec = cast(Codec[ToolMessage | Command[Any] | dict[str, bool]], _ToolNodeResultCodec())
+    graph_bubble: GraphBubbleUp | None = None
+
+    async def _call(
+        args: dict[str, Any],
+    ) -> nemo_relay.ToolExecutionResult[ToolMessage | Command[Any] | dict[str, bool]]:
+        nonlocal graph_bubble
+        try:
+            result = await execute(request.override(tool_call={**request.tool_call, "args": args}))
+        except GraphBubbleUp as error:
+            # See the synchronous wrapper: retain the original exception instead
+            # of letting the native callback boundary convert it to RuntimeError.
+            graph_bubble = error
+            result = _GRAPH_BUBBLE_RESULT
+        return nemo_relay.ToolExecutionResult(result)
+
+    outcome = await nemo_relay.typed.tool_execute(
+        name=tool_name,
+        args=tool_args,
+        func=_call,
+        args_codec=args_codec,
+        result_codec=result_codec,
+        handle=parent,
+        tool_call_id=tool_call_id,
+    )
+    if graph_bubble is not None:
+        raise graph_bubble
+    return cast(ToolMessage | Command[Any], outcome.result)
+
+
+def create_tool_node(
+    tools: Sequence[BaseTool | Callable[..., Any]],
+    **tool_node_kwargs: Any,
+) -> ToolNode:
+    """Create a LangGraph ``ToolNode`` whose tool calls use managed Relay execution.
+
+    For custom LangGraph wrapper composition, construct ``ToolNode`` directly
+    and pass :func:`wrap_tool_call` and :func:`awrap_tool_call` explicitly.
+
+    Args:
+        tools: LangGraph tools available to the returned node.
+        **tool_node_kwargs: Remaining native ``ToolNode`` constructor options,
+            excluding the Relay-managed wrapper options.
+
+    Returns:
+        A native ``ToolNode`` configured with Relay's sync and async wrappers.
+    """
+    configured_wrappers = {"wrap_tool_call", "awrap_tool_call"}.intersection(tool_node_kwargs)
+    if configured_wrappers:
+        names = ", ".join(sorted(configured_wrappers))
+        raise ValueError(f"create_tool_node configures {names}; construct ToolNode directly to compose custom wrappers")
+    error_policy = tool_node_kwargs.pop("handle_tool_errors", _DEFAULT_HANDLE_TOOL_ERRORS)
+    if error_policy is False:
+        error_handler: bool | Callable[[Exception], str] = False
+    else:
+
+        def error_handler(error: Exception) -> str:
+            return _handle_tool_error(error, error_policy)
+
+    return ToolNode(
+        tools,
+        wrap_tool_call=wrap_tool_call,
+        awrap_tool_call=awrap_tool_call,
+        handle_tool_errors=error_handler,
+        **tool_node_kwargs,
+    )
+
+
+__all__ = ["awrap_tool_call", "create_tool_node", "wrap_tool_call"]