From a3471dcc127eb95054caf1fce8fe9a8fc0dc4642 Mon Sep 17 00:00:00 2001 From: Elior Nataf Lackritz Date: Sun, 11 Oct 2026 09:37:59 -0400 Subject: [PATCH] fix(langgraph): give a DeltaChannel message inside an Overwrite an id ensure_message_ids skipped an Overwrite, so the messages inside it were saved with no id. --- libs/langgraph/langgraph/pregel/_messages.py | 8 ++++++- .../tests/test_delta_channel_id_stability.py | 22 +++++++++++++++++++ .../tests/test_delta_channel_update_state.py | 16 +++++++++++++- 3 files changed, 44 insertions(+), 2 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/_messages.py b/libs/langgraph/langgraph/pregel/_messages.py index 2b61c8349..784615c2d 100644 --- a/libs/langgraph/langgraph/pregel/_messages.py +++ b/libs/langgraph/langgraph/pregel/_messages.py @@ -17,6 +17,7 @@ from pydantic import BaseModel from langgraph._internal._config import filter_to_user_tags from langgraph._internal._constants import NS_SEP +from langgraph.channels.binop import _get_overwrite from langgraph.constants import TAG_HIDDEN, TAG_NOSTREAM from langgraph.pregel.protocol import StreamChunk from langgraph.types import Command @@ -431,7 +432,7 @@ def ensure_message_ids(value: Any) -> None: BaseMessages; every get_state() replay then produces a different UUID and the same message appears with a different ID in each LangSmith trace. - Handles three input shapes: + Handles four input shapes: - BaseMessage objects: assign a UUID if id is None. - Dicts with a known "role" (OpenAI-style) or "type" (LangChain format) at the root level: stamp "id" into the dict in-place. The reducer's @@ -439,6 +440,7 @@ def ensure_message_ids(value: Any) -> None: - Lists of the above: apply the same logic to each element, replacing dict items with coerced BaseMessages so the shared list reference seen by checkpoint_pending_writes and the background thread both get typed messages. + - An `Overwrite` of any of the above: apply the same logic to its value. Mutating synchronously here (before the background thread is submitted) is safe: the serialised bytes always reflect the post-coercion state. @@ -459,3 +461,7 @@ def ensure_message_ids(value: Any) -> None: if msg.id is None: msg.id = str(uuid4()) value[i] = msg + else: + is_overwrite, overwrite_value = _get_overwrite(value) + if is_overwrite: + ensure_message_ids(overwrite_value) diff --git a/libs/langgraph/tests/test_delta_channel_id_stability.py b/libs/langgraph/tests/test_delta_channel_id_stability.py index 588616012..684c08725 100644 --- a/libs/langgraph/tests/test_delta_channel_id_stability.py +++ b/libs/langgraph/tests/test_delta_channel_id_stability.py @@ -17,6 +17,7 @@ from typing_extensions import TypedDict from langgraph.channels.delta import DeltaChannel from langgraph.graph import END, START, StateGraph +from langgraph.types import Overwrite pytestmark = pytest.mark.anyio @@ -138,3 +139,24 @@ def test_delta_channel_dict_style_message_gets_stable_id() -> None: assert len(set(ids)) == 1, ( f"dict-style HumanMessage id must be stable across get_state() calls; got {ids}" ) + + +def test_delta_channel_message_in_an_overwrite_gets_stable_id() -> None: + graph = _build_graph(InMemorySaver()) + config = {"configurable": {"thread_id": "overwrite-id-stability"}} + + graph.invoke({"messages": Overwrite([HumanMessage(content="hello")])}, config) + + ids = [ + next( + m.id + for m in graph.get_state(config).values["messages"] + if isinstance(m, HumanMessage) + ) + for _ in range(3) + ] + + assert ids[0] is not None, "a message inside an Overwrite should get a UUID" + assert len(set(ids)) == 1, ( + f"its id must be stable across get_state() calls; got {ids}" + ) diff --git a/libs/langgraph/tests/test_delta_channel_update_state.py b/libs/langgraph/tests/test_delta_channel_update_state.py index 1a2f04a81..abcffcb15 100644 --- a/libs/langgraph/tests/test_delta_channel_update_state.py +++ b/libs/langgraph/tests/test_delta_channel_update_state.py @@ -31,7 +31,7 @@ from langgraph.channels.last_value import LastValue from langgraph.graph import START, StateGraph from langgraph.graph.message import _messages_delta_reducer from langgraph.pregel import NodeBuilder, Pregel -from langgraph.types import StateSnapshot, StateUpdate +from langgraph.types import Overwrite, StateSnapshot, StateUpdate pytestmark = pytest.mark.anyio @@ -688,6 +688,20 @@ async def test_aupdate_state_gives_a_message_an_id_that_every_read_keeps() -> No assert first == second +def test_update_state_gives_a_message_in_an_overwrite_an_id_that_every_read_keeps() -> ( + None +): + graph = _build_graph(InMemorySaver()) + config = {"configurable": {"thread_id": "t"}} + graph.invoke({"messages": [HumanMessage("a", id="a")]}, config) + + graph.update_state(config, {"messages": Overwrite([HumanMessage("b")])}) + + first, second = _message_ids(graph, config), _message_ids(graph, config) + assert first[-1] is not None + assert first == second + + def test_update_state_on_an_older_checkpoint_gives_a_message_an_id() -> None: graph = _build_graph(InMemorySaver()) config = {"configurable": {"thread_id": "t"}}