From 8cea8ae1de85d2caf13bcbbed780e27c7358877e Mon Sep 17 00:00:00 2001 From: samuel-pullely <153273475+samuel-pullely@users.noreply.github.com> Date: Tue, 29 Jul 2025 22:30:42 +0200 Subject: [PATCH] fix(prebuilt): update ToolNode to allow `Command` update to remove all messages (#5678) ## Description Previously, when a tool returned `Command` to update the graph's state, the `_validate_tool_command` method in `ToolNode` would raise a `ValueError` if the `messages_update` list contained only a `RemoveMessage(id=REMOVE_ALL_MESSAGES)` object. This was because the validation logic expected a matching `ToolMessage` for the tool call and did not account for this specific state-clearing scenario. This commit modifies the validation logic to check if the `messages_update` list contains a single `RemoveMessage(id=REMOVE_ALL_MESSAGES)` element. If this condition is met, the `ToolMessage` validation is bypassed, allowing a tool to clear the entire message history without causing a validation error. A new test case, `test_tool_node_command_remove_all_messages`, has been added to `tests/test_tool_node.py` to verify this change and prevent future regressions. ## Example Here is a self-contained example that illustrates the problem and the fix. Without this change, the code block for `Example 2` would raise a `ValueError`. ```python from typing import Annotated, List from langchain_core.messages import ( AIMessage, AnyMessage, HumanMessage, RemoveMessage, ToolMessage, ) from langchain_core.tools import InjectedToolCallId, tool from langchain_openai import ChatOpenAI from langgraph.graph import END, StateGraph, add_messages from langgraph.graph.message import REMOVE_ALL_MESSAGES from langgraph.prebuilt import InjectedState, ToolNode from langgraph.types import Command from pydantic import BaseModel, Field # Agent state tracks current and all messages class AgentState(BaseModel): messages: Annotated[List[AnyMessage], add_messages] = Field( default_factory=list, description="Current conversation messages." ) all_messages: Annotated[List[AnyMessage], add_messages] = Field( default_factory=list, description="All messages, including removed ones." ) # Tool to clear history if long enough, otherwise returns a warning @tool def clear_history_tool( state: Annotated[AgentState, InjectedState], tool_call_id: Annotated[str, InjectedToolCallId], ): """Clears message history if it's long enough.""" if len(state.messages) < 3: return Command( update={ "messages": [ ToolMessage( "History is not long enough to be cleared. Please try again.", tool_call_id=tool_call_id, ) ] } ) else: return Command( update={ "messages": [RemoveMessage(id=REMOVE_ALL_MESSAGES)], "all_messages": state.messages + [ ToolMessage( "History has been successfully cleared.", tool_call_id=tool_call_id, ) ], } ) # Bind the tool to the model model = ChatOpenAI(model="gpt-4o-mini").bind_tools([clear_history_tool]) def model_node(state: AgentState): return {"messages": [model.invoke(state.messages)]} # Build the agent graph graph_builder = StateGraph(AgentState) graph_builder.add_node("model", model_node) graph_builder.add_node("tools", ToolNode([clear_history_tool])) graph_builder.set_entry_point("model") graph_builder.add_edge("model", "tools") graph_builder.add_edge("tools", END) graph = graph_builder.compile() def print_messages(header, messages): print(f"\n{header}") for message in messages: message.pretty_print() ### Example 1: Not enough history to clear state_1 = AgentState( messages=[HumanMessage(content="Please clear my message history.")] ) output_1 = graph.invoke(state_1) print_messages("First call: State 'messages'", output_1["messages"]) print_messages("First call: State 'all_messages'", output_1["all_messages"]) ### Example 2: History is cleared state_2 = AgentState( messages=[ HumanMessage(content="Will this PR get merged?"), AIMessage(content="Maybe, if it's good enough."), HumanMessage(content="Please clear my message history."), ] ) # Without the changes in this PR, the following line will raise a ValueError output_2 = graph.invoke(state_2) print_messages("Second call: State 'messages'", output_2["messages"]) print_messages("Second call: State 'all_messages'", output_2["all_messages"]) ``` ### Outputs *Without the changes in this PR:* ``` First call: State 'messages' ================================ Human Message ================================= Please clear my message history. ================================== Ai Message ================================== Tool Calls: clear_history_tool (ba421ac3-1e1a-4208-a8f6-c5500ee0abcc) Call ID: ba421ac3-1e1a-4208-a8f6-c5500ee0abcc Args: ================================= Tool Message ================================= Name: clear_history_tool History is not long enough to be cleared. Please try again. First call: State 'all_messages' Traceback (most recent call last): File "main.py", line 114, in output_2 = graph.invoke(state_2) ^^^^^^^^^^^^^^^^^^^^^ File ".venv/lib/python3.11/site-packages/langgraph/pregel/__init__.py", line 2844, in invoke for chunk in self.stream( File ".venv/lib/python3.11/site-packages/langgraph/pregel/__init__.py", line 2534, in stream for _ in runner.tick( File ".venv/lib/python3.11/site-packages/langgraph/prebuilt/tool_node.py", line 241, in _func outputs = [ ^ File ".venv/lib/python3.11/concurrent/futures/_base.py", line 619, in result_iterator yield _result_or_cancel(fs.pop()) ^^^^^^^^^^^^^^^^^^^^^^^^^^^ File ".venv/lib/python3.11/concurrent/futures/_base.py", line 317, in _result_or_cancel return fut.result(timeout) ^^^^^^^^^^^^^^^^^^^ File ".venv/lib/python3.11/concurrent/futures/_base.py", line 449, in result return self.__get_result() ^^^^^^^^^^^^^^^^^^^ File ".venv/lib/python3.11/concurrent/futures/_base.py", line 401, in __get_result raise self._exception File ".venv/lib/python3.11/concurrent/futures/thread.py", line 58, in run result = self.fn(*self.args, **self.kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File ".venv/lib/python3.11/site-packages/langchain_core/runnables/config.py", line 555, in _wrapped_fn return contexts.pop().run(fn, *args) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File ".venv/lib/python3.11/site-packages/langgraph/prebuilt/tool_node.py", line 353, in _run_one return self._validate_tool_command(response, call, input_type) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File ".venv/lib/python3.11/site-packages/langgraph/prebuilt/tool_node.py", line 616, in _validate_tool_command raise ValueError( ValueError: Expected to have a matching ToolMessage in Command.update for tool 'clear_history_tool', got: [RemoveMessage(content='', additional_kwargs={}, response_metadata={}, id='__remove_all__')]. Every tool call (LLM requesting to call a tool) in the message history MUST have a corresponding ToolMessage. You can fix it by modifying the tool to return `Command(update={"messages": [ToolMessage("Success", tool_call_id=tool_call_id), ...]}, ...)`. ``` *With the changes in this PR:* ``` First call: State 'messages' ================================ Human Message ================================= Please clear my message history. ================================== Ai Message ================================== Tool Calls: clear_history_tool (ba421ac3-1e1a-4208-a8f6-c5500ee0abcc) Call ID: ba421ac3-1e1a-4208-a8f6-c5500ee0abcc Args: ================================= Tool Message ================================= Name: clear_history_tool History is not long enough to be cleared. Please try again. First call: State 'all_messages' Second call: State 'messages' Second call: State 'all_messages' ================================ Human Message ================================= Will this PR get merged? ================================== Ai Message ================================== Maybe, if it's good enough. ================================ Human Message ================================= Please clear my message history. ================================== Ai Message ================================== Tool Calls: clear_history_tool (499b1be3-6df1-493f-85e5-8d7e429dead8) Call ID: 499b1be3-6df1-493f-85e5-8d7e429dead8 Args: ================================= Tool Message ================================= Name: clear_history_tool History has been successfully cleared. ``` ## Twitter handle [@samuelpullely](https://x.com/samuelpullely) --------- Co-authored-by: Eugene Yurtsev --- libs/prebuilt/langgraph/prebuilt/tool_node.py | 7 +++++ libs/prebuilt/tests/test_tool_node.py | 27 +++++++++++++++++++ 2 files changed, 34 insertions(+) diff --git a/libs/prebuilt/langgraph/prebuilt/tool_node.py b/libs/prebuilt/langgraph/prebuilt/tool_node.py index 12f751d6f..db45f8f46 100644 --- a/libs/prebuilt/langgraph/prebuilt/tool_node.py +++ b/libs/prebuilt/langgraph/prebuilt/tool_node.py @@ -52,6 +52,7 @@ from typing import ( from langchain_core.messages import ( AIMessage, AnyMessage, + RemoveMessage, ToolCall, ToolMessage, convert_to_messages, @@ -72,6 +73,7 @@ from typing_extensions import Annotated, get_args, get_origin from langgraph._internal._runnable import RunnableCallable from langgraph.errors import GraphBubbleUp +from langgraph.graph.message import REMOVE_ALL_MESSAGES from langgraph.prebuilt._internal import ToolCallWithContext from langgraph.store.base import BaseStore from langgraph.types import Command, Send @@ -754,6 +756,11 @@ class ToolNode(RunnableCallable): # convert to message objects if updates are in a dict format messages_update = convert_to_messages(messages_update) + + # no validation needed if all messages are being removed + if messages_update == [RemoveMessage(id=REMOVE_ALL_MESSAGES)]: + return updated_command + has_matching_tool_message = False for message in messages_update: if not isinstance(message, ToolMessage): diff --git a/libs/prebuilt/tests/test_tool_node.py b/libs/prebuilt/tests/test_tool_node.py index ec7c91fdd..2b6dfbebe 100644 --- a/libs/prebuilt/tests/test_tool_node.py +++ b/libs/prebuilt/tests/test_tool_node.py @@ -7,6 +7,7 @@ from typing import ( import pytest from langchain_core.messages import ( AIMessage, + RemoveMessage, ToolMessage, ) from langchain_core.tools import BaseTool, ToolException @@ -15,6 +16,7 @@ from pydantic import BaseModel, ValidationError from pydantic.v1 import ValidationError as ValidationErrorV1 from langgraph.errors import GraphBubbleUp, GraphInterrupt +from langgraph.graph.message import REMOVE_ALL_MESSAGES from langgraph.prebuilt import ToolNode from langgraph.prebuilt.tool_node import TOOL_CALL_ERROR_TEMPLATE from langgraph.types import Command, Send @@ -1129,3 +1131,28 @@ def test_tool_node_parent_command_with_send(): graph=Command.PARENT, ) ] + + +async def test_tool_node_command_remove_all_messages(): + from langchain_core.tools.base import InjectedToolCallId + + @dec_tool + def remove_all_messages_tool(tool_call_id: Annotated[str, InjectedToolCallId]): + """A tool that removes all messages.""" + return Command(update={"messages": [RemoveMessage(id=REMOVE_ALL_MESSAGES)]}) + + tool_node = ToolNode([remove_all_messages_tool]) + tool_call = { + "name": "remove_all_messages_tool", + "args": {}, + "id": "tool_call_123", + } + result = await tool_node.ainvoke( + {"messages": [AIMessage(content="", tool_calls=[tool_call])]} + ) + + assert isinstance(result, list) + assert len(result) == 1 + command = result[0] + assert isinstance(command, Command) + assert command.update == {"messages": [RemoveMessage(id=REMOVE_ALL_MESSAGES)]}