From d1d28d262668b6349164a0878902e29790bca303 Mon Sep 17 00:00:00 2001 From: Will Killian Date: Tue, 25 Aug 2026 18:34:52 -0400 Subject: [PATCH 1/4] feat: manage standalone LangGraph ToolNode calls Signed-off-by: Will Killian --- docs/supported-integrations/langgraph.mdx | 47 +++++ .../integrations/langgraph/__init__.py | 4 + .../integrations/langgraph/tool_node.py | 146 ++++++++++++++ .../test_langgraph_integration.py | 190 ++++++++++++++++++ 4 files changed, 387 insertions(+) create mode 100644 python/nemo_relay/integrations/langgraph/tool_node.py diff --git a/docs/supported-integrations/langgraph.mdx b/docs/supported-integrations/langgraph.mdx index 595a7e596..02debf25f 100644 --- a/docs/supported-integrations/langgraph.mdx +++ b/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: diff --git a/python/nemo_relay/integrations/langgraph/__init__.py b/python/nemo_relay/integrations/langgraph/__init__.py index 79c021c4c..a9da9218b 100644 --- a/python/nemo_relay/integrations/langgraph/__init__.py +++ b/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", ] diff --git a/python/nemo_relay/integrations/langgraph/tool_node.py b/python/nemo_relay/integrations/langgraph/tool_node.py new file mode 100644 index 000000000..e3dbccc9f --- /dev/null +++ b/python/nemo_relay/integrations/langgraph/tool_node.py @@ -0,0 +1,146 @@ +# 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 typing import Any, cast + +from langchain_core.messages import ToolMessage +from langchain_core.tools import BaseTool +from langgraph.prebuilt import ToolNode +from langgraph.prebuilt.tool_node import ToolCallRequest +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.""" + + def from_json(self, data: nemo_relay.Json) -> object: + result = super().from_json(data) + if not isinstance(result, Command) or not isinstance(result.update, dict): + return result + + messages = result.update.get("messages") + if not isinstance(messages, list): + return result + result.update["messages"] = [ + ToolMessage.model_validate(message) if isinstance(message, dict) else message for message in messages + ] + return result + + +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]], _ToolNodeResultCodec()) + + def _call(args: dict[str, Any]) -> nemo_relay.ToolExecutionResult[ToolMessage | Command[Any]]: + return nemo_relay.ToolExecutionResult(execute(request.override(tool_call={**request.tool_call, "args": args}))) + + return 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, + ) + ).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]], _ToolNodeResultCodec()) + + async def _call(args: dict[str, Any]) -> nemo_relay.ToolExecutionResult[ToolMessage | Command[Any]]: + return nemo_relay.ToolExecutionResult( + await execute(request.override(tool_call={**request.tool_call, "args": args})) + ) + + return ( + 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, + ) + ).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") + return ToolNode( + tools, + wrap_tool_call=wrap_tool_call, + awrap_tool_call=awrap_tool_call, + **tool_node_kwargs, + ) + + +__all__ = ["awrap_tool_call", "create_tool_node", "wrap_tool_call"] diff --git a/python/tests/integrations/langgraph_tests/test_langgraph_integration.py b/python/tests/integrations/langgraph_tests/test_langgraph_integration.py index 7696c241e..e18e0fd0c 100644 --- a/python/tests/integrations/langgraph_tests/test_langgraph_integration.py +++ b/python/tests/integrations/langgraph_tests/test_langgraph_integration.py @@ -25,6 +25,11 @@ class State(TypedDict): value: int +class ToolNodeState(TypedDict): + messages: Annotated[list[Any], operator.add] + offset: int + + def increment(state: State) -> State: return {"value": state["value"] + 1} @@ -86,6 +91,191 @@ def test_handler_type(callback_handler: NemoRelayCallbackHandler): assert isinstance(callback_handler, GraphCallbackHandler) +@pytest.mark.parametrize("use_async", [False, True]) +def test_create_tool_node_routes_standalone_tool_calls_through_relay( + use_async: bool, + subscribed_events: list[nemo_relay.Event], +): + from langchain_core.messages import AIMessage + from langchain_core.tools import tool + from langgraph.graph import END, START, StateGraph + from langgraph.prebuilt import InjectedState + + from nemo_relay.integrations.langgraph import create_tool_node + + def add_offset(value: int, state: Any) -> str: + """Add the graph state's offset to a model-provided value.""" + return str(value + state["offset"]) + + add_offset.__annotations__["state"] = Annotated[dict[str, Any], InjectedState] + add_offset_tool = tool(add_offset) + + async def rewrite_tool_args(_name: str, args: nemo_relay.Json, next_call: Any) -> Any: + downstream = await next_call({**cast(dict[str, Any], args), "value": 4}) + return nemo_relay.ToolExecutionInterceptOutcome( + downstream.result, + annotation=downstream.annotation, + ) + + node = create_tool_node([add_offset_tool], name="managed-tools") + builder = StateGraph(ToolNodeState) + builder.add_node("tools", node) + builder.add_edge(START, "tools") + builder.add_edge("tools", END) + graph = builder.compile() + nemo_relay.intercepts.register_tool_execution("langgraph-rewrite-tool-args", 1, rewrite_tool_args) + try: + with nemo_relay.scope.scope("request", nemo_relay.ScopeType.Agent) as request: + input_state = { + "offset": 3, + "messages": [ + AIMessage( + content="", + tool_calls=[ + {"name": "add_offset", "args": {"value": 1}, "id": "call-1"}, + {"name": "add_offset", "args": {"value": 2}, "id": "call-3"}, + ], + ) + ], + } + if use_async: + result = asyncio.run(graph.ainvoke(input_state)) + else: + result = graph.invoke(input_state) + finally: + nemo_relay.intercepts.deregister_tool_execution("langgraph-rewrite-tool-args") + + nemo_relay.subscribers.flush() + assert [message.content for message in result["messages"][-2:]] == ["7", "7"] + tool_events = [ + event for event in subscribed_events if isinstance(event, nemo_relay.ScopeEvent) and event.name == "add_offset" + ] + assert len(tool_events) == 4 + assert {event.scope_category for event in tool_events} == {"start", "end"} + assert all(event.parent_uuid == request.uuid for event in tool_events) + assert {event.category_profile["tool_call_id"] for event in tool_events if event.category_profile} == { + "call-1", + "call-3", + } + + +def test_exported_tool_node_wrappers_support_direct_tool_node_construction( + subscribed_events: list[nemo_relay.Event], +): + from langchain_core.messages import AIMessage + from langchain_core.tools import tool + from langgraph.graph import END, START, MessagesState, StateGraph + from langgraph.prebuilt import ToolNode + + from nemo_relay.integrations.langgraph import awrap_tool_call, wrap_tool_call + + @tool + def echo(value: str) -> str: + """Echo a value.""" + return value + + builder = StateGraph(MessagesState) + builder.add_node("tools", ToolNode([echo], wrap_tool_call=wrap_tool_call, awrap_tool_call=awrap_tool_call)) + builder.add_edge(START, "tools") + builder.add_edge("tools", END) + graph = builder.compile() + with nemo_relay.scope.scope("request", nemo_relay.ScopeType.Agent): + result = graph.invoke( + { + "messages": [ + AIMessage( + content="", + tool_calls=[{"name": "echo", "args": {"value": "managed"}, "id": "call-2"}], + ) + ] + } + ) + + nemo_relay.subscribers.flush() + assert result["messages"][-1].content == "managed" + assert any( + isinstance(event, nemo_relay.ScopeEvent) + and event.name == "echo" + and event.category_profile == {"tool_call_id": "call-2"} + for event in subscribed_events + ) + + +def test_create_tool_node_preserves_command_and_error_handling( + subscribed_events: list[nemo_relay.Event], +): + from langchain_core.messages import AIMessage, ToolMessage + from langchain_core.tools import tool + from langgraph.graph import END, START, StateGraph + from langgraph.types import Command + + from nemo_relay.integrations.langgraph import create_tool_node + + @tool + def update_offset(): + """Update graph state through a LangGraph command.""" + return Command( + update={ + "offset": 9, + "messages": [ToolMessage(content="updated", tool_call_id="call-command")], + } + ) + + @tool + def fail() -> str: + """Raise a tool failure handled by ToolNode.""" + raise ValueError("expected failure") + + command_builder = StateGraph(ToolNodeState) + command_builder.add_node("command", create_tool_node([update_offset])) + command_builder.add_edge(START, "command") + command_builder.add_edge("command", END) + command_graph = command_builder.compile() + + failure_builder = StateGraph(ToolNodeState) + failure_builder.add_node("failure", create_tool_node([fail], handle_tool_errors=True)) + failure_builder.add_edge(START, "failure") + failure_builder.add_edge("failure", END) + failure_graph = failure_builder.compile() + + with nemo_relay.scope.scope("request", nemo_relay.ScopeType.Agent): + command_result = command_graph.invoke( + { + "offset": 0, + "messages": [ + AIMessage( + content="", + tool_calls=[{"name": "update_offset", "args": {}, "id": "call-command"}], + ) + ], + } + ) + failure_result = failure_graph.invoke( + { + "offset": 0, + "messages": [AIMessage(content="", tool_calls=[{"name": "fail", "args": {}, "id": "call-failure"}])], + } + ) + + nemo_relay.subscribers.flush() + assert command_result["offset"] == 9 + assert command_result["messages"][-1].content == "updated" + assert failure_result["messages"][-1].status == "error" + assert {event.name for event in subscribed_events if isinstance(event, nemo_relay.ScopeEvent)} == { + "request", + "update_offset", + "fail", + } + + +@pytest.mark.parametrize("wrapper_name", ["wrap_tool_call", "awrap_tool_call"]) +def test_create_tool_node_rejects_custom_tool_wrappers(wrapper_name: str): + from nemo_relay.integrations.langgraph import create_tool_node + + with pytest.raises(ValueError, match="construct ToolNode directly"): + create_tool_node([], **{wrapper_name: lambda *_args: None}) + + class TestGraphCallbacks: _expected_events = [ "scope.start.request", From 9e2325cdad2f1c172f0a7f8d75a43d80de677946 Mon Sep 17 00:00:00 2001 From: Will Killian Date: Tue, 25 Aug 2026 18:54:55 -0400 Subject: [PATCH 2/4] fix: preserve LangGraph ToolNode interrupts Signed-off-by: Will Killian --- .../integrations/langgraph/tool_node.py | 137 +++++++++++++----- .../test_langgraph_integration.py | 81 +++++++++++ 2 files changed, 185 insertions(+), 33 deletions(-) diff --git a/python/nemo_relay/integrations/langgraph/tool_node.py b/python/nemo_relay/integrations/langgraph/tool_node.py index e3dbccc9f..8fc51f9f8 100644 --- a/python/nemo_relay/integrations/langgraph/tool_node.py +++ b/python/nemo_relay/integrations/langgraph/tool_node.py @@ -6,12 +6,14 @@ 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 +from langgraph.prebuilt.tool_node import ToolCallRequest, ToolInvocationError from langgraph.types import Command import nemo_relay @@ -22,20 +24,57 @@ 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) or not isinstance(result.update, dict): + if not isinstance(result, Command): return result - messages = result.update.get("messages") - if not isinstance(messages, list): - return result - result.update["messages"] = [ - ToolMessage.model_validate(message) if isinstance(message, dict) else message for message in messages - ] + 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, bool | tuple) or (isinstance(policy, type) and issubclass(policy, Exception)): + 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 ( @@ -61,12 +100,22 @@ def wrap_tool_call( """ 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]], _ToolNodeResultCodec()) - - def _call(args: dict[str, Any]) -> nemo_relay.ToolExecutionResult[ToolMessage | Command[Any]]: - return nemo_relay.ToolExecutionResult(execute(request.override(tool_call={**request.tool_call, "args": args}))) - - return run_sync( + 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, @@ -76,7 +125,10 @@ def _call(args: dict[str, Any]) -> nemo_relay.ToolExecutionResult[ToolMessage | handle=parent, tool_call_id=tool_call_id, ) - ).result + ) + if graph_bubble is not None: + raise graph_bubble + return cast(ToolMessage | Command[Any], outcome.result) async def awrap_tool_call( @@ -94,24 +146,34 @@ async def awrap_tool_call( """ 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]], _ToolNodeResultCodec()) - - async def _call(args: dict[str, Any]) -> nemo_relay.ToolExecutionResult[ToolMessage | Command[Any]]: - return nemo_relay.ToolExecutionResult( - await execute(request.override(tool_call={**request.tool_call, "args": args})) - ) - - return ( - 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, - ) - ).result + 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( @@ -135,10 +197,19 @@ def create_tool_node( 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, ) diff --git a/python/tests/integrations/langgraph_tests/test_langgraph_integration.py b/python/tests/integrations/langgraph_tests/test_langgraph_integration.py index e18e0fd0c..5bbbb9905 100644 --- a/python/tests/integrations/langgraph_tests/test_langgraph_integration.py +++ b/python/tests/integrations/langgraph_tests/test_langgraph_integration.py @@ -268,6 +268,87 @@ def fail() -> str: } +@pytest.mark.parametrize("use_async", [False, True]) +def test_create_tool_node_preserves_list_command_results(use_async: bool): + from langchain_core.messages import AIMessage, ToolMessage + from langchain_core.tools import tool + from langgraph.graph import END, START, StateGraph + from langgraph.types import Command + + from nemo_relay.integrations.langgraph import create_tool_node + + @tool + def update_offset(): + """Return a list containing a graph state update command.""" + return [ + Command( + update={ + "offset": 9, + "messages": [ToolMessage(content="updated", tool_call_id="call-list")], + } + ) + ] + + builder = StateGraph(ToolNodeState) + builder.add_node("tools", create_tool_node([update_offset])) + builder.add_edge(START, "tools") + builder.add_edge("tools", END) + graph = builder.compile() + input_state = { + "offset": 0, + "messages": [AIMessage(content="", tool_calls=[{"name": "update_offset", "args": {}, "id": "call-list"}])], + } + + if use_async: + result = asyncio.run(graph.ainvoke(input_state)) + else: + result = graph.invoke(input_state) + + assert result["offset"] == 9 + assert result["messages"][-1].content == "updated" + + +@pytest.mark.parametrize("use_async", [False, True]) +def test_create_tool_node_propagates_graph_interrupts(use_async: bool): + from langchain_core.messages import AIMessage + from langchain_core.tools import tool + from langgraph.checkpoint.memory import MemorySaver + from langgraph.graph import END, START, StateGraph + from langgraph.types import interrupt + + from nemo_relay.integrations.langgraph import create_tool_node + + @tool + def await_approval() -> str: + """Pause the graph until approval is supplied.""" + interrupt("approval required") + return "approved" + + builder = StateGraph(ToolNodeState) + builder.add_node("tools", create_tool_node([await_approval])) + builder.add_edge(START, "tools") + builder.add_edge("tools", END) + graph = builder.compile(checkpointer=MemorySaver()) + config = {"configurable": {"thread_id": str(uuid4())}} + input_state = { + "offset": 0, + "messages": [ + AIMessage(content="", tool_calls=[{"name": "await_approval", "args": {}, "id": "call-interrupt"}]) + ], + } + + if use_async: + + async def collect_updates() -> list[dict[str, Any]]: + return [update async for update in graph.astream(input_state, config, stream_mode="updates")] + + updates = asyncio.run(collect_updates()) + else: + updates = list(graph.stream(input_state, config, stream_mode="updates")) + + assert any("__interrupt__" in update for update in updates) + + @pytest.mark.parametrize("wrapper_name", ["wrap_tool_call", "awrap_tool_call"]) def test_create_tool_node_rejects_custom_tool_wrappers(wrapper_name: str): from nemo_relay.integrations.langgraph import create_tool_node From 113b1d2f7c26134c8b31515aba0eb9c0240b9fdd Mon Sep 17 00:00:00 2001 From: Will Killian Date: Tue, 25 Aug 2026 19:39:20 -0400 Subject: [PATCH 3/4] fix: preserve selective ToolNode error policies Signed-off-by: Will Killian --- .../integrations/langgraph/tool_node.py | 10 +++- .../test_langgraph_integration.py | 55 +++++++++++++++++++ 2 files changed, 64 insertions(+), 1 deletion(-) diff --git a/python/nemo_relay/integrations/langgraph/tool_node.py b/python/nemo_relay/integrations/langgraph/tool_node.py index 8fc51f9f8..78372e55f 100644 --- a/python/nemo_relay/integrations/langgraph/tool_node.py +++ b/python/nemo_relay/integrations/langgraph/tool_node.py @@ -66,7 +66,15 @@ def _handle_tool_error(error: Exception, policy: object) -> str: if isinstance(error, ToolInvocationError): return error.message raise error - if isinstance(policy, bool | tuple) or (isinstance(policy, type) and issubclass(policy, Exception)): + 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 diff --git a/python/tests/integrations/langgraph_tests/test_langgraph_integration.py b/python/tests/integrations/langgraph_tests/test_langgraph_integration.py index 5bbbb9905..2ae92150f 100644 --- a/python/tests/integrations/langgraph_tests/test_langgraph_integration.py +++ b/python/tests/integrations/langgraph_tests/test_langgraph_integration.py @@ -268,6 +268,61 @@ def fail() -> str: } +@pytest.mark.parametrize("use_async", [False, True]) +def test_create_tool_node_preserves_selective_error_handling(use_async: bool): + from langchain_core.messages import AIMessage + from langchain_core.tools import tool + from langgraph.graph import END, START, StateGraph + + from nemo_relay.integrations.langgraph import create_tool_node + + @tool + def value_error() -> str: + """Raise an error selected by the ToolNode policy.""" + raise ValueError("handled") + + @tool + def runtime_error() -> str: + """Raise an error outside the ToolNode policy.""" + raise RuntimeError("unhandled") + + def build_graph(tool: Any): + builder = StateGraph(ToolNodeState) + builder.add_node("tools", create_tool_node([tool], handle_tool_errors=ValueError)) + builder.add_edge(START, "tools") + builder.add_edge("tools", END) + return builder.compile() + + handled_graph = build_graph(value_error) + unhandled_graph = build_graph(runtime_error) + handled_input = { + "offset": 0, + "messages": [AIMessage(content="", tool_calls=[{"name": "value_error", "args": {}, "id": "call-value"}])], + } + unhandled_input = { + "offset": 0, + "messages": [AIMessage(content="", tool_calls=[{"name": "runtime_error", "args": {}, "id": "call-runtime"}])], + } + + if use_async: + + async def invoke_handled() -> ToolNodeState: + return await handled_graph.ainvoke(handled_input) + + async def invoke_unhandled() -> ToolNodeState: + return await unhandled_graph.ainvoke(unhandled_input) + + handled_result = asyncio.run(invoke_handled()) + with pytest.raises(RuntimeError, match="internal error"): + asyncio.run(invoke_unhandled()) + else: + handled_result = handled_graph.invoke(handled_input) + with pytest.raises(RuntimeError, match="internal error"): + unhandled_graph.invoke(unhandled_input) + + assert handled_result["messages"][-1].status == "error" + + @pytest.mark.parametrize("use_async", [False, True]) def test_create_tool_node_preserves_list_command_results(use_async: bool): from langchain_core.messages import AIMessage, ToolMessage From 848b1cfa26fa046a4c940cdbc3a9ace3ec6047a1 Mon Sep 17 00:00:00 2001 From: Will Killian Date: Tue, 25 Aug 2026 20:03:07 -0400 Subject: [PATCH 4/4] test: cover tuple ToolNode error policies Signed-off-by: Will Killian --- .../langgraph_tests/test_langgraph_integration.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/python/tests/integrations/langgraph_tests/test_langgraph_integration.py b/python/tests/integrations/langgraph_tests/test_langgraph_integration.py index 2ae92150f..57b845979 100644 --- a/python/tests/integrations/langgraph_tests/test_langgraph_integration.py +++ b/python/tests/integrations/langgraph_tests/test_langgraph_integration.py @@ -269,7 +269,11 @@ def fail() -> str: @pytest.mark.parametrize("use_async", [False, True]) -def test_create_tool_node_preserves_selective_error_handling(use_async: bool): +@pytest.mark.parametrize("policy", [ValueError, (ValueError,)]) +def test_create_tool_node_preserves_selective_error_handling( + use_async: bool, + policy: type[ValueError] | tuple[type[ValueError], ...], +): from langchain_core.messages import AIMessage from langchain_core.tools import tool from langgraph.graph import END, START, StateGraph @@ -288,7 +292,7 @@ def runtime_error() -> str: def build_graph(tool: Any): builder = StateGraph(ToolNodeState) - builder.add_node("tools", create_tool_node([tool], handle_tool_errors=ValueError)) + builder.add_node("tools", create_tool_node([tool], handle_tool_errors=policy)) builder.add_edge(START, "tools") builder.add_edge("tools", END) return builder.compile()