diff --git a/libs/langgraph/tests/test_channels.py b/libs/langgraph/tests/test_channels.py index aee6a1cea..f73e12930 100644 --- a/libs/langgraph/tests/test_channels.py +++ b/libs/langgraph/tests/test_channels.py @@ -237,3 +237,74 @@ def test_diff_channel_unsupported_saver_raises() -> None: raw_delta = DiffDelta(delta=[], prev_version=None) with pytest.raises(ValueError, match="DiffChannel received a raw DiffDelta"): spec.from_checkpoint(raw_delta) + + +def test_diff_channel_remove_message_delta_and_replay() -> None: + """RemoveMessage stored in a delta must round-trip correctly through the chain.""" + from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage + from langgraph.checkpoint.base import DiffChainValue, DiffDelta + + from langgraph.channels.diff import DiffChannel + from langgraph.graph.message import add_messages + + spec = DiffChannel(add_messages) + ch = spec.from_checkpoint(MISSING) + ch.after_checkpoint(None) + + # Step 1: add two messages + ch.update([HumanMessage(content="hi", id="h1")]) + ch.update([AIMessage(content="hello", id="a1")]) + d1 = ch.checkpoint() + assert isinstance(d1, DiffDelta) + ch.after_checkpoint("v1") + assert ch.get() == [ + HumanMessage(content="hi", id="h1"), + AIMessage(content="hello", id="a1"), + ] + + # Step 2: remove the AI message + ch.update([RemoveMessage(id="a1")]) + d2 = ch.checkpoint() + assert isinstance(d2, DiffDelta) + assert d2.prev_version == "v1" + assert any(isinstance(w, RemoveMessage) for w in d2.delta) + ch.after_checkpoint("v2") + assert ch.get() == [HumanMessage(content="hi", id="h1")] + + # Replay the full chain from scratch — must reproduce the post-remove state + chain = DiffChainValue(base=None, deltas=[d1.delta, d2.delta]) + ch2 = spec.from_checkpoint(chain) + assert ch2.get() == [HumanMessage(content="hi", id="h1")] + + +def test_diff_channel_update_by_id_delta_and_replay() -> None: + """Updating a message by ID stored in a delta must round-trip correctly.""" + from langchain_core.messages import HumanMessage + from langgraph.checkpoint.base import DiffChainValue, DiffDelta + + from langgraph.channels.diff import DiffChannel + from langgraph.graph.message import add_messages + + spec = DiffChannel(add_messages) + ch = spec.from_checkpoint(MISSING) + ch.after_checkpoint(None) + + # Step 1: add a message + ch.update([HumanMessage(content="original", id="h1")]) + d1 = ch.checkpoint() + assert isinstance(d1, DiffDelta) + ch.after_checkpoint("v1") + + # Step 2: update the same message by ID + ch.update([HumanMessage(content="updated", id="h1")]) + d2 = ch.checkpoint() + assert isinstance(d2, DiffDelta) + assert d2.prev_version == "v1" + ch.after_checkpoint("v2") + assert ch.get() == [HumanMessage(content="updated", id="h1")] + + # Replay the full chain — must produce the updated message, not the original + chain = DiffChainValue(base=None, deltas=[d1.delta, d2.delta]) + ch2 = spec.from_checkpoint(chain) + assert len(ch2.get()) == 1 + assert ch2.get()[0].content == "updated" diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 84e68d653..e975d02ab 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -9500,3 +9500,90 @@ async def test_diff_channel_time_travel() -> None: assert msgs[0].content == "h1" assert msgs[1].content == "ai-1" assert msgs[2].content == "h3" + + +async def test_diff_channel_remove_message_end_to_end() -> None: + """RemoveMessage inside a DiffChannel graph must persist and reload correctly.""" + from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage + from langgraph.checkpoint.memory import InMemorySaver + + from langgraph.channels.diff import DiffChannel + from langgraph.graph import START, StateGraph + from langgraph.graph.message import add_messages + + class State(TypedDict): + messages: Annotated[list, DiffChannel(add_messages)] + + def respond(state: State) -> dict: + return {"messages": [AIMessage(content="reply", id="ai-1")]} + + def delete_first(state: State) -> dict: + # removes the first message + return {"messages": [RemoveMessage(id=state["messages"][0].id)]} + + builder = StateGraph(State) + builder.add_node("respond", respond) + builder.add_node("delete_first", delete_first) + builder.add_edge(START, "respond") + builder.add_edge("respond", "delete_first") + graph = builder.compile(checkpointer=InMemorySaver()) + + config = {"configurable": {"thread_id": "diff-remove-test"}} + graph.invoke({"messages": [HumanMessage(content="hello", id="h1")]}, config) + + state = graph.get_state(config) + msgs = state.values["messages"] + # h1 was removed, only ai-1 should remain + assert len(msgs) == 1, f"expected 1 message, got {len(msgs)}: {msgs}" + assert msgs[0].id == "ai-1" + + # A subsequent turn must reconstruct from the checkpoint correctly + graph.invoke({"messages": [HumanMessage(content="again", id="h2")]}, config) + state = graph.get_state(config) + msgs = state.values["messages"] + # ai-1 + h2 + ai-1(second reply, same id overwrites) + h2 removed + # more simply: after second run we expect ai-1 updated + h2 remaining minus deleted h2 + # just assert h1 is still gone + assert all(m.id != "h1" for m in msgs), ( + "h1 should still be absent after second turn" + ) + + +async def test_diff_channel_update_by_id_end_to_end() -> None: + """Updating a message by ID via DiffChannel must persist and reload correctly.""" + from langchain_core.messages import HumanMessage + from langgraph.checkpoint.memory import InMemorySaver + + from langgraph.channels.diff import DiffChannel + from langgraph.graph import START, StateGraph + from langgraph.graph.message import add_messages + + class State(TypedDict): + messages: Annotated[list, DiffChannel(add_messages)] + + def update_msg(state: State) -> dict: + # re-send h1 with updated content + return {"messages": [HumanMessage(content="updated", id="h1")]} + + builder = StateGraph(State) + builder.add_node("update_msg", update_msg) + builder.add_edge(START, "update_msg") + graph = builder.compile(checkpointer=InMemorySaver()) + + config = {"configurable": {"thread_id": "diff-update-id-test"}} + graph.invoke({"messages": [HumanMessage(content="original", id="h1")]}, config) + + state = graph.get_state(config) + msgs = state.values["messages"] + assert len(msgs) == 1, f"expected 1 message, got {len(msgs)}: {msgs}" + assert msgs[0].content == "updated" + assert msgs[0].id == "h1" + + # Second turn: verify the updated state is the base for further accumulation + graph.invoke({"messages": [HumanMessage(content="new", id="h2")]}, config) + state = graph.get_state(config) + msgs = state.values["messages"] + ids = [m.id for m in msgs] + assert "h1" in ids # h1 persists (updated, not duplicated) + assert "h2" in ids + assert ids.count("h1") == 1, "h1 must not be duplicated"