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

@@ -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.
 

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
 # ---------------------------------------------------------------------------