8cea8ae1de 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>
2025-07-29 20:30:42 +00:00
x
2025-07-25 14:37:03 -04:00
2025-06-10 10:06:08 -07:00
2024-03-15 14:31:59 -07:00

LangGraph Logo

Version Downloads Open Issues Docs

Trusted by companies shaping the future of agents including Klarna, Replit, Elastic, and more LangGraph is a low-level orchestration framework for building, managing, and deploying long-running, stateful agents.

Get started

Install LangGraph:

pip install -U langgraph

Then, create an agent using prebuilt components:

# pip install -qU "langchain[anthropic]" to call the model

from langgraph.prebuilt import create_react_agent

def get_weather(city: str) -> str:
    """Get weather for a given city."""
    return f"It's always sunny in {city}!"

agent = create_react_agent(
    model="anthropic:claude-3-7-sonnet-latest",
    tools=[get_weather],
    prompt="You are a helpful assistant"
)

# Run the agent
agent.invoke(
    {"messages": [{"role": "user", "content": "what is the weather in sf"}]}
)

For more information, see the Quickstart. Or, to learn how to build an agent workflow with a customizable architecture, long-term memory, and other complex task handling, see the LangGraph basics tutorials.

Core benefits

LangGraph provides low-level supporting infrastructure for any long-running, stateful workflow or agent. LangGraph does not abstract prompts or architecture, and provides the following central benefits:

  • Durable execution: Build agents that persist through failures and can run for extended periods, automatically resuming from exactly where they left off.
  • Human-in-the-loop: Seamlessly incorporate human oversight by inspecting and modifying agent state at any point during execution.
  • Comprehensive memory: Create truly stateful agents with both short-term working memory for ongoing reasoning and long-term persistent memory across sessions.
  • Debugging with LangSmith: Gain deep visibility into complex agent behavior with visualization tools that trace execution paths, capture state transitions, and provide detailed runtime metrics.
  • Production-ready deployment: Deploy sophisticated agent systems confidently with scalable infrastructure designed to handle the unique challenges of stateful, long-running workflows.

LangGraphs ecosystem

While LangGraph can be used standalone, it also integrates seamlessly with any LangChain product, giving developers a full suite of tools for building agents. To improve your LLM application development, pair LangGraph with:

  • LangSmith — Helpful for agent evals and observability. Debug poor-performing LLM app runs, evaluate agent trajectories, gain visibility in production, and improve performance over time.
  • LangGraph Platform — Deploy and scale agents effortlessly with a purpose-built deployment platform for long running, stateful workflows. Discover, reuse, configure, and share agents across teams — and iterate quickly with visual prototyping in LangGraph Studio.
  • LangChain Provides integrations and composable components to streamline LLM application development.

Note

Looking for the JS version of LangGraph? See the JS repo and the JS docs.

Additional resources

  • Guides: Quick, actionable code snippets for topics such as streaming, adding memory & persistence, and design patterns (e.g. branching, subgraphs, etc.).
  • Reference: Detailed reference on core classes, methods, how to use the graph and checkpointing APIs, and higher-level prebuilt components.
  • Examples: Guided examples on getting started with LangGraph.
  • LangChain Forum: Connect with the community and share all of your technical questions, ideas, and feedback.
  • LangChain Academy: Learn the basics of LangGraph in our free, structured course.
  • Templates: Pre-built reference apps for common agentic workflows (e.g. ReAct agent, memory, retrieval etc.) that can be cloned and adapted.
  • Case studies: Hear how industry leaders use LangGraph to ship AI applications at scale.

Acknowledgements

LangGraph is inspired by Pregel and Apache Beam. The public interface draws inspiration from NetworkX. LangGraph is built by LangChain Inc, the creators of LangChain, but can be used without LangChain.

Languages
Python 99.6%
Makefile 0.2%
TypeScript 0.1%