fix(langgraph): do not apply pending writes when updating state (#6389)

PR #6195 fixed `bulk_update_state` to populate `task.result` by calling
`prepare_next_tasks` to discover task IDs. Before #6195,
prepare_next_tasks was gated by the condition `CONFIG_KEY_CHECKPOINT_ID
not in config[CONF]` - so it only ran if we were resuming from an empty
checkpoint. This check was removed in order to properly populate task
results. However, the removal of this check inadvertently applied
pending writes during manual state updates which caused issues when
forking:

- When you fork from a checkpoint by calling `update_state(config,
new_values, as_node="mynode")`, pending writes from the original
execution were being applied
- This caused stale data to leak into forked threads (eg. old tool call
results appearing in forked execution)

Changes

Removed pending writes application from `bulk_update_state` and
`abulk_update_state`:
- Still call `prepare_next_tasks` to discover task IDs, but skip the
code that applies null writes and regular pending writes

Tests

- Added `test_fork_does_not_apply_pending_writes` for sync and async
which verifies forking doesn't include stale pending writes from
original execution

Author: casparb

libs/langgraph/langgraph/pregel/main.py

@@ -1702,7 +1702,9 @@ def perform_superstep(
             # we use the task id generated by prepare_next_tasks
             node_to_task_ids: dict[str, deque[str]] = defaultdict(deque)
             if saved is not None and saved.pending_writes is not None:
-                # tasks for this checkpoint
+                # we call prepare_next_tasks to discover the task IDs that
+                # would have been generated, so we can reuse them and
+                # properly populate task.result in state history
                 next_tasks = prepare_next_tasks(
                     checkpoint,
                     saved.pending_writes,
@@ -1721,32 +1723,6 @@ def perform_superstep(
                 for t in next_tasks.values():
                     node_to_task_ids[t.name].append(t.id)
 
-                # apply null writes
-                if null_writes := [
-                    w[1:] for w in saved.pending_writes or [] if w[0] == NULL_TASK_ID
-                ]:
-                    apply_writes(
-                        checkpoint,
-                        channels,
-                        [PregelTaskWrites((), INPUT, null_writes, [])],
-                        checkpointer.get_next_version,
-                        self.trigger_to_nodes,
-                    )
-                # apply writes
-                for tid, k, v in saved.pending_writes:
-                    if k in (ERROR, INTERRUPT):
-                        continue
-                    if tid not in next_tasks:
-                        continue
-                    next_tasks[tid].writes.append((k, v))
-                if tasks := [t for t in next_tasks.values() if t.writes]:
-                    apply_writes(
-                        checkpoint,
-                        channels,
-                        tasks,
-                        checkpointer.get_next_version,
-                        self.trigger_to_nodes,
-                    )
             valid_updates: list[tuple[str, dict[str, Any] | None, str | None]] = []
             if len(updates) == 1:
                 values, as_node, task_id = updates[0]
@@ -2167,7 +2143,9 @@ async def aperform_superstep(
             # we use the task id generated by prepare_next_tasks
             node_to_task_ids: dict[str, deque[str]] = defaultdict(deque)
             if saved is not None and saved.pending_writes is not None:
-                # tasks for this checkpoint
+                # we call prepare_next_tasks to discover the task IDs that
+                # would have been generated, so we can reuse them and
+                # properly populate task.result in state history
                 next_tasks = prepare_next_tasks(
                     checkpoint,
                     saved.pending_writes,
@@ -2186,31 +2164,6 @@ async def aperform_superstep(
                 for t in next_tasks.values():
                     node_to_task_ids[t.name].append(t.id)
 
-                # apply null writes
-                if null_writes := [
-                    w[1:] for w in saved.pending_writes or [] if w[0] == NULL_TASK_ID
-                ]:
-                    apply_writes(
-                        checkpoint,
-                        channels,
-                        [PregelTaskWrites((), INPUT, null_writes, [])],
-                        checkpointer.get_next_version,
-                        self.trigger_to_nodes,
-                    )
-                for tid, k, v in saved.pending_writes:
-                    if k in (ERROR, INTERRUPT):
-                        continue
-                    if tid not in next_tasks:
-                        continue
-                    next_tasks[tid].writes.append((k, v))
-                if tasks := [t for t in next_tasks.values() if t.writes]:
-                    apply_writes(
-                        checkpoint,
-                        channels,
-                        tasks,
-                        checkpointer.get_next_version,
-                        self.trigger_to_nodes,
-                    )
             valid_updates: list[tuple[str, dict[str, Any] | None, str | None]] = []
             if len(updates) == 1:
                 values, as_node, task_id = updates[0]

libs/langgraph/tests/test_pregel.py

@@ -8805,3 +8805,43 @@ def node_c(state: State):
         InvalidUpdateError, match="Can receive only one Overwrite value per super-step."
     ):
         graph.invoke({"messages": ["START"]}, config)
+
+
+def test_fork_does_not_apply_pending_writes(
+    sync_checkpointer: BaseCheckpointSaver,
+) -> None:
+    """Test that forking with update_state does not apply pending writes from original execution."""
+
+    class State(TypedDict):
+        value: Annotated[int, operator.add]
+
+    def node_a(state: State) -> State:
+        return {"value": 10}
+
+    def node_b(state: State) -> State:
+        return {"value": 100}
+
+    graph = (
+        StateGraph(State)
+        .add_node("node_a", node_a)
+        .add_node("node_b", node_b)
+        .add_edge(START, "node_a")
+        .add_edge("node_a", "node_b")
+        .compile(checkpointer=sync_checkpointer)
+    )
+
+    thread1 = {"configurable": {"thread_id": "1"}}
+    graph.invoke({"value": 1}, thread1)
+
+    history = list(graph.get_state_history(thread1))
+    checkpoint_before_a = next(s for s in history if s.next == ("node_a",))
+
+    fork_config = graph.update_state(
+        checkpoint_before_a.config, {"value": 20}, as_node="node_a"
+    )
+
+    # Continue from fork (should run node_b)
+    result = graph.invoke(None, fork_config)
+
+    # Should be: 1 (input) + 20 (forked node_a) + 100 (node_b) = 121
+    assert result == {"value": 121}

libs/langgraph/tests/test_pregel_async.py

@@ -9249,3 +9249,41 @@ def first_task_result(history: list[StateSnapshot], node: str) -> Any:
 
     assert bulk_start_result == ref_start_result == {"num": 1, "text": "one"}
     assert bulk_double_result == ref_double_result == {"num": 2, "text": "oneone"}
+
+
+async def test_fork_does_not_apply_pending_writes(
+    async_checkpointer: BaseCheckpointSaver,
+) -> None:
+    """Test that forking with aupdate_state does not apply pending writes from original execution."""
+
+    class State(TypedDict):
+        value: Annotated[int, operator.add]
+
+    def node_a(state: State) -> State:
+        return {"value": 10}
+
+    def node_b(state: State) -> State:
+        return {"value": 100}
+
+    graph = (
+        StateGraph(State)
+        .add_node("node_a", node_a)
+        .add_node("node_b", node_b)
+        .add_edge(START, "node_a")
+        .add_edge("node_a", "node_b")
+        .compile(checkpointer=async_checkpointer)
+    )
+
+    thread1 = {"configurable": {"thread_id": "1"}}
+    await graph.ainvoke({"value": 1}, thread1)
+
+    history = [c async for c in graph.aget_state_history(thread1)]
+    checkpoint_before_a = next(s for s in history if s.next == ("node_a",))
+
+    fork_config = await graph.aupdate_state(
+        checkpoint_before_a.config, {"value": 20}, as_node="node_a"
+    )
+    result = await graph.ainvoke(None, fork_config)
+
+    # 1 (input) + 20 (forked node_a) + 100 (node_b) = 121
+    assert result == {"value": 121}