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 <module>
    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 <eyurtsev@gmail.com>

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):

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)]}