From a7a465877d94fbf26bcd2dab72bd6b6da23d6ef3 Mon Sep 17 00:00:00 2001 From: Elior Nataf Lackritz Date: Tue, 6 Oct 2026 16:03:21 -0400 Subject: [PATCH] fix(langgraph): store a task path for each bulk_update_state update Updates applied together in one superstep are applied in the order given, but were stored without a task path, so a saver replaying by task path fell back to their task ids, which are hashes. Each update now gets the path (__interrupt__, i) and its writes store it. --- libs/langgraph/langgraph/pregel/main.py | 28 +++++++++-- .../tests/test_delta_channel_update_state.py | 50 +++++++++++++++++++ 2 files changed, 73 insertions(+), 5 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/main.py b/libs/langgraph/langgraph/pregel/main.py index 76f33b4d0..16232d39a 100644 --- a/libs/langgraph/langgraph/pregel/main.py +++ b/libs/langgraph/langgraph/pregel/main.py @@ -19,7 +19,7 @@ from collections.abc import ( from dataclasses import is_dataclass, replace from datetime import timedelta from functools import partial -from inspect import isclass +from inspect import isclass, signature from typing import ( Any, Generic, @@ -126,6 +126,7 @@ from langgraph.pregel._algo import ( apply_writes, local_read, prepare_next_tasks, + task_path_str, ) from langgraph.pregel._call import identifier from langgraph.pregel._checkpoint import ( @@ -1989,7 +1990,7 @@ class Pregel( if not writers: raise InvalidUpdateError(f"Node {as_node} has no writers") writes: deque[tuple[str, Any]] = deque() - task = PregelTaskWrites((), as_node, writes, [INTERRUPT]) + task = PregelTaskWrites((INTERRUPT, i), as_node, writes, [INTERRUPT]) # get the task ids that were prepared for this node # if a task id was provided in the StateUpdate, we use it # otherwise, we use the next available task id @@ -2050,7 +2051,10 @@ class Pregel( channel_writes = [w for w in task.writes if w[0] != PUSH] if channel_writes: checkpointer.put_writes( - checkpoint_config, channel_writes, task_id + checkpoint_config, + channel_writes, + task_id, + **_task_path_kwarg(checkpointer.put_writes, task), ) apply_writes( checkpoint, @@ -2477,7 +2481,7 @@ class Pregel( if not writers: raise InvalidUpdateError(f"Node {as_node} has no writers") writes: deque[tuple[str, Any]] = deque() - task = PregelTaskWrites((), as_node, writes, [INTERRUPT]) + task = PregelTaskWrites((INTERRUPT, i), as_node, writes, [INTERRUPT]) # get the task ids that were prepared for this node # if a task id was provided in the StateUpdate, we use it # otherwise, we use the next available task id @@ -2538,7 +2542,10 @@ class Pregel( channel_writes = [w for w in task.writes if w[0] != PUSH] if channel_writes: await checkpointer.aput_writes( - checkpoint_config, channel_writes, task_id + checkpoint_config, + channel_writes, + task_id, + **_task_path_kwarg(checkpointer.aput_writes, task), ) apply_writes( checkpoint, @@ -4237,6 +4244,17 @@ class Pregel( await self.cache.aclear(namespaces) +def _task_path_kwarg(put_writes: Callable[..., Any], task: PregelTaskWrites) -> dict: + """Pass the task's path to savers whose `put_writes` takes one. + + Savers that replay a checkpoint's writes in task path order then give back + updates applied together in the order they were given. + """ + if signature(put_writes).parameters.get("task_path") is None: + return {} + return {"task_path": task_path_str(task.path)} + + def _update_task_id(checkpoint_id: str, i: int) -> str: """Task id for the `i`th update of a superstep that has no task to reuse. diff --git a/libs/langgraph/tests/test_delta_channel_update_state.py b/libs/langgraph/tests/test_delta_channel_update_state.py index 41c3ac07d..b8424dfac 100644 --- a/libs/langgraph/tests/test_delta_channel_update_state.py +++ b/libs/langgraph/tests/test_delta_channel_update_state.py @@ -434,6 +434,56 @@ def test_bulk_update_state_keeps_every_update_next_to_a_pending_task( ) +class _TaskPathOrderSaver(InMemorySaver): + """Replays each checkpoint's writes by `(task_path, task_id, idx)`.""" + + def get_tuple(self, config: Any) -> Any: + tup = super().get_tuple(config) + if tup is None or not tup.pending_writes: + return tup + conf = tup.config["configurable"] + stored = self.writes[ + (conf["thread_id"], conf["checkpoint_ns"], conf["checkpoint_id"]) + ] + rows = sorted( + zip(stored.items(), tup.pending_writes), + key=lambda row: (row[0][1][3], *row[0][0]), + ) + return tup._replace(pending_writes=[write for _, write in rows]) + + get_delta_channel_history = BaseCheckpointSaver.get_delta_channel_history + aget_delta_channel_history = BaseCheckpointSaver.aget_delta_channel_history + + +GIVEN = ["u1", "u2", "u3", "u4", "u5", "u6"] + + +def _updates_in_given_order() -> list[list[StateUpdate]]: + return [ + [_update(c, "assistant" if i % 2 else "model") for i, c in enumerate(GIVEN)] + ] + + +def test_bulk_update_state_replays_updates_in_the_order_given() -> None: + graph = _build_graph(_TaskPathOrderSaver(), two_nodes=True) + config = {"configurable": {"thread_id": "bulk-order"}} + graph.invoke({"messages": [HumanMessage(content="hi", id="hi")]}, config) + + graph.bulk_update_state(config, _updates_in_given_order()) + + assert _contents(graph.get_state(config)) == ["hi", *GIVEN] + + +async def test_abulk_update_state_replays_updates_in_the_order_given() -> None: + graph = _build_graph(_TaskPathOrderSaver(), two_nodes=True) + config = {"configurable": {"thread_id": "bulk-order"}} + await graph.ainvoke({"messages": [HumanMessage(content="hi", id="hi")]}, config) + + await graph.abulk_update_state(config, _updates_in_given_order()) + + assert _contents(await graph.aget_state(config)) == ["hi", *GIVEN] + + # --------------------------------------------------------------------------- # Public-API observation of fresh-thread checkpoint shape # ---------------------------------------------------------------------------