diff --git a/libs/langgraph/langgraph/graph/message.py b/libs/langgraph/langgraph/graph/message.py index 1e5fe1ca8..0e0bc3c7a 100644 --- a/libs/langgraph/langgraph/graph/message.py +++ b/libs/langgraph/langgraph/graph/message.py @@ -249,18 +249,15 @@ def _messages_delta_reducer( ) -> list[AnyMessage]: """**Experimental.** Batch reducer for use with `DeltaChannel`. - Processes all writes in one pass — dedup by ID, `RemoveMessage` - tombstoning — without calling `add_messages`. + Provides full `add_messages` parity: dedup by ID, `RemoveMessage` + tombstoning, `REMOVE_ALL_MESSAGES` reset, `BaseMessageChunk` coercion, + and UUID assignment for ID-less messages — all in a single batched pass. This reducer is batching-invariant, as required by `DeltaChannel`: `reducer(reducer(state, xs), ys) == reducer(state, xs + ys)`. Raw dict / string / tuple inputs are coerced to typed `BaseMessage` - objects so that HTTP-driven graphs work without a separate coercion - step. Messages without IDs are assigned UUIDs (matching `add_messages` - behavior) so that message eviction and `RemoveMessage` tombstoning work - correctly. `REMOVE_ALL_MESSAGES` and `BaseMessageChunk` conversion are - not handled here. + objects so that HTTP-driven graphs work without a separate coercion step. Example:: @@ -281,14 +278,31 @@ def _messages_delta_reducer( flat.extend(w) else: flat.append(w) - # Steady state: the reducer's own output is already typed, so skip - # `convert_to_messages` on state when the first element is a BaseMessage. + # Steady state: the reducer's own output is already typed BaseMessages + # (never chunks), so skip convert_to_messages on the fast path. # Only raw input (initial dicts, deserialized blobs) hits the slow path. if state and isinstance(state[0], BaseMessage): - state_msgs = state + state_msgs = cast("list[AnyMessage]", state) else: - state_msgs = cast("list[AnyMessage]", convert_to_messages(state)) - msgs = cast("list[AnyMessage]", convert_to_messages(flat)) + state_msgs = cast( + "list[AnyMessage]", + [message_chunk_to_message(cast(BaseMessageChunk, m)) for m in convert_to_messages(state)], + ) + # Coerce chunks to full messages — streaming nodes can emit BaseMessageChunk. + msgs = cast( + "list[AnyMessage]", + [message_chunk_to_message(cast(BaseMessageChunk, m)) for m in convert_to_messages(flat)], + ) + + # REMOVE_ALL_MESSAGES resets everything; find the last sentinel and + # discard all state plus all writes before it. + remove_all_idx = None + for idx, m in enumerate(msgs): + if isinstance(m, RemoveMessage) and m.id == REMOVE_ALL_MESSAGES: + remove_all_idx = idx + if remove_all_idx is not None: + state_msgs = [] + msgs = msgs[remove_all_idx + 1 :] # Build index and assign missing IDs in one pass (parity with add_messages # so that eviction and RemoveMessage tombstoning work on ID-less messages). diff --git a/libs/langgraph/tests/test_channels.py b/libs/langgraph/tests/test_channels.py index ecede949c..8c6ccb842 100644 --- a/libs/langgraph/tests/test_channels.py +++ b/libs/langgraph/tests/test_channels.py @@ -3,7 +3,7 @@ from collections.abc import Sequence from typing import Annotated import pytest -from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage +from langchain_core.messages import AIMessage, AIMessageChunk, HumanMessage, RemoveMessage from langgraph.checkpoint.memory import InMemorySaver from langgraph.checkpoint.serde.types import _DeltaSnapshot from typing_extensions import NotRequired, TypedDict @@ -16,7 +16,7 @@ from langgraph.channels.topic import Topic from langgraph.channels.untracked_value import UntrackedValue from langgraph.errors import EmptyChannelError, InvalidUpdateError from langgraph.graph import START, StateGraph -from langgraph.graph.message import _messages_delta_reducer +from langgraph.graph.message import REMOVE_ALL_MESSAGES, _messages_delta_reducer from langgraph.graph.state import _get_channel from langgraph.types import Overwrite @@ -306,6 +306,31 @@ def test_messages_delta_reducer_assigns_uuid_to_id_less_messages() -> None: assert result2[0].content == "hi" +def test_messages_delta_reducer_remove_all_messages() -> None: + """REMOVE_ALL_MESSAGES sentinel clears all state and preceding writes.""" + state = [HumanMessage(content="old", id="h1"), AIMessage(content="prior", id="a1")] + + # Sentinel mid-batch: everything before it (including state) is discarded. + result = _messages_delta_reducer(state, [[RemoveMessage(id=REMOVE_ALL_MESSAGES), HumanMessage(content="fresh", id="h2")]]) + assert len(result) == 1 + assert result[0].content == "fresh" + + # Batching-invariant: split across two calls must equal one combined call. + step1 = _messages_delta_reducer(state, [[RemoveMessage(id=REMOVE_ALL_MESSAGES)]]) + step2 = _messages_delta_reducer(step1, [[HumanMessage(content="fresh", id="h2")]]) + assert step2 == result + + +def test_messages_delta_reducer_coerces_message_chunks() -> None: + """BaseMessageChunk writes are coerced to full messages.""" + chunk = AIMessageChunk(content="hello", id="a1") + result = _messages_delta_reducer([], [[chunk]]) + assert len(result) == 1 + assert not isinstance(result[0], AIMessageChunk) + assert result[0].content == "hello" + assert result[0].id == "a1" + + def test_delta_channel_checkpoint_returns_missing() -> None: """checkpoint() always returns MISSING regardless of state.