Fix test

libs/langgraph/langgraph/pregel/__init__.py

@@ -1421,7 +1421,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
                 self.channels,
                 checkpoint,
             )
-            values, as_node = updates[0]
+            values, as_node, _ = updates[0]
 
             # no values as END, just clear all tasks
             if values is None and as_node == END:
@@ -1588,22 +1588,32 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
                     )
 
                     tasks_group_by = defaultdict(list)
+                    user_group_by: dict[str, list[StateUpdate]] = defaultdict(list)
+
                     for task in next_tasks.values():
                         tasks_group_by[task.name].append(task.id)
 
-                    user_group_by: dict[str, list[StateUpdate]] = defaultdict(list)
                     for item in values:
-                        if not isinstance(item, dict):
-                            continue
-
-                        values = item["values"]
-                        as_node = item["as_node"]
+                        if isinstance(item, dict):
+                            values = item["values"]
+                            as_node = item["as_node"]
+                        elif isinstance(item, StateUpdate):
+                            values = item.values
+                            as_node = item.as_node
+                        else:
+                            raise InvalidUpdateError(
+                                f"Invalid update item: {item} when copying checkpoint"
+                            )
 
                         user_group = user_group_by[as_node]
                         tasks_group = tasks_group_by[as_node]
-                        
+
                         target_idx = len(user_group)
-                        task_id = tasks_group[target_idx] if target_idx < len(tasks_group) else None
+                        task_id = (
+                            tasks_group[target_idx]
+                            if target_idx < len(tasks_group)
+                            else None
+                        )
 
                         user_group_by[as_node].append(
                             StateUpdate(values=values, as_node=as_node, task_id=task_id)
@@ -1870,7 +1880,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
                 self.channels,
                 checkpoint,
             )
-            values, as_node = updates[0]
+            values, as_node, _ = updates[0]
             # no values, just clear all tasks
             if values is None and as_node == END:
                 if len(updates) > 1:
@@ -1992,24 +2002,84 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
                     )
 
             # no values, copy checkpoint
-            if values is None and as_node == "__copy__":
+            if as_node == "__copy__":
                 if len(updates) > 1:
                     raise InvalidUpdateError(
                         "Cannot copy checkpoint with multiple updates"
                     )
 
+                if saved is None:
+                    raise InvalidUpdateError("Cannot copy a non-existent checkpoint")
+
                 next_checkpoint = create_checkpoint(checkpoint, None, step)
                 # copy checkpoint
                 next_config = await checkpointer.aput(
-                    saved.parent_config or saved.config if saved else checkpoint_config,
+                    saved.parent_config
+                    or patch_configurable(
+                        saved.config, {CONFIG_KEY_CHECKPOINT_ID: None}
+                    ),
                     next_checkpoint,
                     {
                         "source": "fork",
                         "step": step + 1,
-                        "parents": saved.metadata.get("parents", {}) if saved else {},
+                        "parents": saved.metadata.get("parents", {}),
                     },
                     {},
                 )
+
+                # we want to both clone a checkpoint and update state in one go.
+                # reuse the same task ID if possible.
+                if isinstance(values, list) and len(values) > 0:
+                    # figure out the task IDs for the next update checkpoint
+                    next_tasks = prepare_next_tasks(
+                        next_checkpoint,
+                        saved.pending_writes,
+                        self.nodes,
+                        channels,
+                        managed,
+                        next_config,
+                        step + 2,
+                        step + 4,
+                        for_execution=False,
+                    )
+
+                    tasks_group_by = defaultdict(list)
+                    user_group_by: dict[str, list[StateUpdate]] = defaultdict(list)
+
+                    for task in next_tasks.values():
+                        tasks_group_by[task.name].append(task.id)
+
+                    for item in values:
+                        if isinstance(item, dict):
+                            values = item["values"]
+                            as_node = item["as_node"]
+                        elif isinstance(item, StateUpdate):
+                            values = item.values
+                            as_node = item.as_node
+                        else:
+                            raise InvalidUpdateError(
+                                f"Invalid update item: {item} when copying checkpoint"
+                            )
+
+                        user_group = user_group_by[as_node]
+                        tasks_group = tasks_group_by[as_node]
+
+                        target_idx = len(user_group)
+                        task_id = (
+                            tasks_group[target_idx]
+                            if target_idx < len(tasks_group)
+                            else None
+                        )
+
+                        user_group_by[as_node].append(
+                            StateUpdate(values=values, as_node=as_node, task_id=task_id)
+                        )
+
+                    return await aperform_superstep(
+                        patch_checkpoint_map(next_config, saved.metadata),
+                        [item for lst in user_group_by.values() for item in lst],
+                    )
+
                 return patch_checkpoint_map(
                     next_config, saved.metadata if saved else None
                 )

libs/langgraph/tests/test_pregel.py

@@ -8051,15 +8051,12 @@ def test_timeout_with_parent_command(
 
 def test_fork_and_update_task_results(sync_checkpointer: BaseCheckpointSaver) -> None:
     """Test forking and updating task results with state history."""
-    two_count = 0
 
-    def checkpoint(snapshot: StateSnapshot) -> tuple[str, dict]:
-        """Convert a checkpoint to a tuple representation for comparison."""
-        return ("checkpoint", {"values": snapshot.values})
+    def checkpoint(values: dict[str, Any]):
+        return ("checkpoint", {"values": values})
 
-    def task(task_info: PregelTask) -> tuple[str, dict]:
-        """Convert a task to a tuple representation for comparison."""
-        return ("task", {"name": task_info.name, "result": task_info.result})
+    def task(name: str, result: Any):
+        return ("task", {"name": name, "result": result})
 
     def get_tree(history: list[StateSnapshot]) -> list:
         """Build a tree structure from state history for comparison."""
@@ -8081,15 +8078,17 @@ def test_fork_and_update_task_results(sync_checkpointer: BaseCheckpointSaver) ->
             node_map[checkpoint_id] = {"item": item, "children": []}
 
             parent = node_map.get(parent_checkpoint_id)
-            if parent:
-                parent["children"].append(node_map[checkpoint_id])
-            else:
-                root_nodes.append(node_map[checkpoint_id])
+            (parent["children"] if parent else root_nodes).append(
+                node_map[checkpoint_id]
+            )
 
         def node_to_tree(node: dict) -> list:
             """Convert a node to tree structure."""
-            result = [checkpoint(node["item"])] + [
-                task(task_info) for task_info in node["item"].tasks
+            result = [
+                checkpoint(node["item"].values),
+            ] + [
+                task(task_info.name, task_info.result)
+                for task_info in node["item"].tasks
             ]
 
             if len(node["children"]) > 1:
@@ -8100,9 +8099,10 @@ def test_fork_and_update_task_results(sync_checkpointer: BaseCheckpointSaver) ->
             else:
                 return result
 
-        # Process all root nodes
         if len(root_nodes) == 1:
+            # Process all root nodes
             return node_to_tree(root_nodes[0])
+
         elif len(root_nodes) > 1:
             # Multiple root nodes - treat as branches
             branches = [node_to_tree(node) for node in root_nodes]
@@ -8110,20 +8110,21 @@ def test_fork_and_update_task_results(sync_checkpointer: BaseCheckpointSaver) ->
         else:
             return []
 
+    class State(TypedDict):
+        name: Annotated[str, lambda a, b: " > ".join([a, b]) if a else b]
+
     # Define the graph with a sequence of nodes
-    def one(state: dict) -> Command:
+    def one(state: State) -> Command:
         return Command(goto=[Send("two", {})], update={"name": "one"})
 
-    def two(state: dict) -> dict:
-        nonlocal two_count
-        two_count += 1
-        return {"name": f"two {two_count}"}
+    def two(state: State) -> State:
+        return {"name": "two"}
 
-    def three(state: dict) -> dict:
+    def three(state: State) -> State:
         return {"name": "three"}
 
     graph = (
-        StateGraph({"name": str})
+        StateGraph(State)
         .add_node("one", one)
         .add_node("two", two)
         .add_node("three", three)
@@ -8141,16 +8142,16 @@ def test_fork_and_update_task_results(sync_checkpointer: BaseCheckpointSaver) ->
     history = list(graph.get_state_history(config))
 
     assert get_tree(history) == [
-        ("checkpoint", {"values": {}}),
-        ("task", {"name": "__start__", "result": {"name": "start"}}),
-        ("checkpoint", {"values": {"name": "start"}}),
-        ("task", {"name": "one", "result": {"name": "one"}}),
-        ("checkpoint", {"values": {"name": "start > one"}}),
-        ("task", {"name": "two", "result": {"name": "two 1"}}),
-        ("task", {"name": "two", "result": {"name": "two 2"}}),
-        ("checkpoint", {"values": {"name": "start > one > two 2 > two 1"}}),
-        ("task", {"name": "three", "result": {"name": "three"}}),
-        ("checkpoint", {"values": {"name": "start > one > two 2 > two 1 > three"}}),
+        checkpoint({"name": ""}),
+        task("__start__", {"name": "start"}),
+        checkpoint({"name": "start"}),
+        task("one", {"name": "one"}),
+        checkpoint({"name": "start > one"}),
+        task("two", {"name": "two"}),
+        task("two", {"name": "two"}),
+        checkpoint({"name": "start > one > two > two"}),
+        task("three", {"name": "three"}),
+        checkpoint({"name": "start > one > two > two > three"}),
     ]
 
     # Update the start state
@@ -8158,39 +8159,36 @@ def test_fork_and_update_task_results(sync_checkpointer: BaseCheckpointSaver) ->
         None,
         graph.update_state(
             history[4].config,
-            [StateUpdate(values={"name": "start*"}, as_node="__start__")],
-            "__copy__",
+            values=[StateUpdate(values={"name": "start*"}, as_node="__start__")],
+            as_node="__copy__",
         ),
     )
 
     history = list(graph.get_state_history(config))
     assert get_tree(history) == [
         [
-            ("checkpoint", {"values": {}}),
-            ("task", {"name": "__start__", "result": {"name": "start"}}),
-            ("checkpoint", {"values": {"name": "start"}}),
-            ("task", {"name": "one", "result": {"name": "one"}}),
-            ("checkpoint", {"values": {"name": "start > one"}}),
-            ("task", {"name": "two", "result": {"name": "two 1"}}),
-            ("task", {"name": "two", "result": {"name": "two 2"}}),
-            ("checkpoint", {"values": {"name": "start > one > two 2 > two 1"}}),
-            ("task", {"name": "three", "result": {"name": "three"}}),
-            ("checkpoint", {"values": {"name": "start > one > two 2 > two 1 > three"}}),
+            checkpoint({"name": ""}),
+            task("__start__", {"name": "start"}),
+            checkpoint({"name": "start"}),
+            task("one", {"name": "one"}),
+            checkpoint({"name": "start > one"}),
+            task("two", {"name": "two"}),
+            task("two", {"name": "two"}),
+            checkpoint({"name": "start > one > two > two"}),
+            task("three", {"name": "three"}),
+            checkpoint({"name": "start > one > two > two > three"}),
         ],
         [
-            ("checkpoint", {"values": {}}),
-            ("task", {"name": "__start__", "result": {"name": "start*"}}),
-            ("checkpoint", {"values": {"name": "start*"}}),
-            ("task", {"name": "one", "result": {"name": "one"}}),
-            ("checkpoint", {"values": {"name": "start* > one"}}),
-            ("task", {"name": "two", "result": {"name": "two 3"}}),
-            ("task", {"name": "two", "result": {"name": "two 4"}}),
-            ("checkpoint", {"values": {"name": "start* > one > two 4 > two 3"}}),
-            ("task", {"name": "three", "result": {"name": "three"}}),
-            (
-                "checkpoint",
-                {"values": {"name": "start* > one > two 4 > two 3 > three"}},
-            ),
+            checkpoint({"name": ""}),
+            task("__start__", {"name": "start*"}),
+            checkpoint({"name": "start*"}),
+            task("one", {"name": "one"}),
+            checkpoint({"name": "start* > one"}),
+            task("two", {"name": "two"}),
+            task("two", {"name": "two"}),
+            checkpoint({"name": "start* > one > two > two"}),
+            task("three", {"name": "three"}),
+            checkpoint({"name": "start* > one > two > two > three"}),
         ],
     ]
 
@@ -8212,54 +8210,44 @@ def test_fork_and_update_task_results(sync_checkpointer: BaseCheckpointSaver) ->
     history = list(graph.get_state_history(config))
     assert get_tree(history) == [
         [
-            ("checkpoint", {"values": {}}),
-            ("task", {"name": "__start__", "result": {"name": "start"}}),
-            ("checkpoint", {"values": {"name": "start"}}),
-            ("task", {"name": "one", "result": {"name": "one"}}),
-            ("checkpoint", {"values": {"name": "start > one"}}),
-            ("task", {"name": "two", "result": {"name": "two 1"}}),
-            ("task", {"name": "two", "result": {"name": "two 2"}}),
-            ("checkpoint", {"values": {"name": "start > one > two 2 > two 1"}}),
-            ("task", {"name": "three", "result": {"name": "three"}}),
-            ("checkpoint", {"values": {"name": "start > one > two 2 > two 1 > three"}}),
+            checkpoint({"name": ""}),
+            task("__start__", {"name": "start"}),
+            checkpoint({"name": "start"}),
+            task("one", {"name": "one"}),
+            checkpoint({"name": "start > one"}),
+            task("two", {"name": "two"}),
+            task("two", {"name": "two"}),
+            checkpoint({"name": "start > one > two > two"}),
+            task("three", {"name": "three"}),
+            checkpoint({"name": "start > one > two > two > three"}),
         ],
         [
-            ("checkpoint", {"values": {}}),
-            ("task", {"name": "__start__", "result": {"name": "start*"}}),
+            checkpoint({"name": ""}),
+            task("__start__", {"name": "start*"}),
             [
                 [
-                    ("checkpoint", {"values": {"name": "start*"}}),
-                    ("task", {"name": "one", "result": {"name": "one"}}),
-                    ("checkpoint", {"values": {"name": "start* > one"}}),
-                    ("task", {"name": "two", "result": {"name": "two 3"}}),
-                    ("task", {"name": "two", "result": {"name": "two 4"}}),
-                    (
-                        "checkpoint",
-                        {"values": {"name": "start* > one > two 4 > two 3"}},
-                    ),
-                    ("task", {"name": "three", "result": {"name": "three"}}),
-                    (
-                        "checkpoint",
-                        {"values": {"name": "start* > one > two 4 > two 3 > three"}},
-                    ),
+                    checkpoint({"name": "start*"}),
+                    task("one", {"name": "one"}),
+                    checkpoint({"name": "start* > one"}),
+                    task("two", {"name": "two"}),
+                    task("two", {"name": "two"}),
+                    checkpoint({"name": "start* > one > two > two"}),
+                    task("three", {"name": "three"}),
+                    checkpoint({"name": "start* > one > two > two > three"}),
                 ],
                 [
-                    ("checkpoint", {"values": {"name": "start*"}}),
-                    ("task", {"name": "one", "result": {"name": "one*"}}),
-                    ("checkpoint", {"values": {"name": "start* > one*"}}),
-                    ("task", {"name": "two", "result": {"name": "two 5"}}),
-                    ("checkpoint", {"values": {"name": "start* > one* > two 5"}}),
-                    ("task", {"name": "three", "result": {"name": "three"}}),
-                    (
-                        "checkpoint",
-                        {"values": {"name": "start* > one* > two 5 > three"}},
-                    ),
+                    checkpoint({"name": "start*"}),
+                    task("one", {"name": "one*"}),
+                    checkpoint({"name": "start* > one*"}),
+                    task("two", {"name": "two"}),
+                    checkpoint({"name": "start* > one* > two"}),
+                    task("three", {"name": "three"}),
+                    checkpoint({"name": "start* > one* > two > three"}),
                 ],
             ],
         ],
     ]
 
-    two_count = 0
     config = {"configurable": {"thread_id": "2"}}
 
     # Initialize the thread once again
@@ -8275,8 +8263,8 @@ def test_fork_and_update_task_results(sync_checkpointer: BaseCheckpointSaver) ->
         graph.update_state(
             history[2].config,
             [
-                StateUpdate(values={"name": "two 3"}, as_node="two"),
-                StateUpdate(values={"name": "two 4"}, as_node="two"),
+                StateUpdate(values={"name": "two"}, as_node="two"),
+                StateUpdate(values={"name": "two"}, as_node="two"),
             ],
             "__copy__",
         ),
@@ -8284,38 +8272,32 @@ def test_fork_and_update_task_results(sync_checkpointer: BaseCheckpointSaver) ->
 
     history = list(graph.get_state_history(config))
     assert get_tree(history) == [
-        ("checkpoint", {"values": {}}),
-        ("task", {"name": "__start__", "result": {"name": "start"}}),
-        ("checkpoint", {"values": {"name": "start"}}),
-        ("task", {"name": "one", "result": {"name": "one"}}),
+        checkpoint({"name": ""}),
+        task("__start__", {"name": "start"}),
+        checkpoint({"name": "start"}),
+        task("one", {"name": "one"}),
         [
             [
-                ("checkpoint", {"values": {"name": "start > one"}}),
-                ("task", {"name": "two", "result": {"name": "two 1"}}),
-                ("task", {"name": "two", "result": {"name": "two 2"}}),
-                ("checkpoint", {"values": {"name": "start > one > two 2 > two 1"}}),
-                ("task", {"name": "three", "result": {"name": "three"}}),
-                (
-                    "checkpoint",
-                    {"values": {"name": "start > one > two 2 > two 1 > three"}},
-                ),
+                checkpoint({"name": "start > one"}),
+                task("two", {"name": "two"}),
+                task("two", {"name": "two"}),
+                checkpoint({"name": "start > one > two > two"}),
+                task("three", {"name": "three"}),
+                checkpoint({"name": "start > one > two > two > three"}),
             ],
             [
-                ("checkpoint", {"values": {"name": "start > one"}}),
-                ("task", {"name": "two", "result": {"name": "two 3"}}),
-                ("task", {"name": "two", "result": {"name": "two 4"}}),
-                ("checkpoint", {"values": {"name": "start > one > two 3 > two 4"}}),
-                ("task", {"name": "three", "result": {"name": "three"}}),
-                (
-                    "checkpoint",
-                    {"values": {"name": "start > one > two 3 > two 4 > three"}},
-                ),
+                checkpoint({"name": "start > one"}),
+                task("two", {"name": "two"}),
+                task("two", {"name": "two"}),
+                checkpoint({"name": "start > one > two > two"}),
+                task("three", {"name": "three"}),
+                checkpoint({"name": "start > one > two > two > three"}),
             ],
         ],
     ]
 
     # Fork task three
-    assert history[1].values == {"name": "start > one > two 3 > two 4"}
+    assert history[1].values == {"name": "start > one > two > two"}
     assert len(history[1].tasks) == 1
     assert history[1].tasks[0].name == "three"
 
@@ -8330,52 +8312,33 @@ def test_fork_and_update_task_results(sync_checkpointer: BaseCheckpointSaver) ->
 
     history = list(graph.get_state_history(config))
     assert get_tree(history) == [
-        ("checkpoint", {"values": {}}),
-        ("task", {"name": "__start__", "result": {"name": "start"}}),
-        ("checkpoint", {"values": {"name": "start"}}),
-        ("task", {"name": "one", "result": {"name": "one"}}),
+        checkpoint({"name": ""}),
+        task("__start__", {"name": "start"}),
+        checkpoint({"name": "start"}),
+        task("one", {"name": "one"}),
         [
             [
-                ("checkpoint", {"values": {"name": "start > one"}}),
-                ("task", {"name": "two", "result": {"name": "two 1"}}),
-                ("task", {"name": "two", "result": {"name": "two 2"}}),
-                ("checkpoint", {"values": {"name": "start > one > two 2 > two 1"}}),
-                ("task", {"name": "three", "result": {"name": "three"}}),
-                (
-                    "checkpoint",
-                    {"values": {"name": "start > one > two 2 > two 1 > three"}},
-                ),
+                checkpoint({"name": "start > one"}),
+                task("two", {"name": "two"}),
+                task("two", {"name": "two"}),
+                checkpoint({"name": "start > one > two > two"}),
+                task("three", {"name": "three"}),
+                checkpoint({"name": "start > one > two > two > three"}),
             ],
             [
-                ("checkpoint", {"values": {"name": "start > one"}}),
-                ("task", {"name": "two", "result": {"name": "two 3"}}),
-                ("task", {"name": "two", "result": {"name": "two 4"}}),
+                checkpoint({"name": "start > one"}),
+                task("two", {"name": "two"}),
+                task("two", {"name": "two"}),
                 [
                     [
-                        (
-                            "checkpoint",
-                            {"values": {"name": "start > one > two 3 > two 4"}},
-                        ),
-                        ("task", {"name": "three", "result": {"name": "three"}}),
-                        (
-                            "checkpoint",
-                            {"values": {"name": "start > one > two 3 > two 4 > three"}},
-                        ),
+                        checkpoint({"name": "start > one > two > two"}),
+                        task("three", {"name": "three"}),
+                        checkpoint({"name": "start > one > two > two > three"}),
                     ],
                     [
-                        (
-                            "checkpoint",
-                            {"values": {"name": "start > one > two 3 > two 4"}},
-                        ),
-                        ("task", {"name": "three", "result": {"name": "three*"}}),
-                        (
-                            "checkpoint",
-                            {
-                                "values": {
-                                    "name": "start > one > two 3 > two 4 > three*"
-                                }
-                            },
-                        ),
+                        checkpoint({"name": "start > one > two > two"}),
+                        task("three", {"name": "three*"}),
+                        checkpoint({"name": "start > one > two > two > three*"}),
                     ],
                 ],
             ],
@@ -8383,7 +8346,7 @@ def test_fork_and_update_task_results(sync_checkpointer: BaseCheckpointSaver) ->
     ]
 
     # Regenerate task three
-    assert history[3].values == {"name": "start > one > two 3 > two 4"}
+    assert history[3].values == {"name": "start > one > two > two"}
     assert len(history[3].tasks) == 1
     assert history[3].tasks[0].name == "three"
 
@@ -8391,63 +8354,38 @@ def test_fork_and_update_task_results(sync_checkpointer: BaseCheckpointSaver) ->
 
     history = list(graph.get_state_history(config))
     assert get_tree(history) == [
-        ("checkpoint", {"values": {}}),
-        ("task", {"name": "__start__", "result": {"name": "start"}}),
-        ("checkpoint", {"values": {"name": "start"}}),
-        ("task", {"name": "one", "result": {"name": "one"}}),
+        checkpoint({"name": ""}),
+        task("__start__", {"name": "start"}),
+        checkpoint({"name": "start"}),
+        task("one", {"name": "one"}),
         [
             [
-                ("checkpoint", {"values": {"name": "start > one"}}),
-                ("task", {"name": "two", "result": {"name": "two 1"}}),
-                ("task", {"name": "two", "result": {"name": "two 2"}}),
-                ("checkpoint", {"values": {"name": "start > one > two 2 > two 1"}}),
-                ("task", {"name": "three", "result": {"name": "three"}}),
-                (
-                    "checkpoint",
-                    {"values": {"name": "start > one > two 2 > two 1 > three"}},
-                ),
+                checkpoint({"name": "start > one"}),
+                task("two", {"name": "two"}),
+                task("two", {"name": "two"}),
+                checkpoint({"name": "start > one > two > two"}),
+                task("three", {"name": "three"}),
+                checkpoint({"name": "start > one > two > two > three"}),
             ],
             [
-                ("checkpoint", {"values": {"name": "start > one"}}),
-                ("task", {"name": "two", "result": {"name": "two 3"}}),
-                ("task", {"name": "two", "result": {"name": "two 4"}}),
+                checkpoint({"name": "start > one"}),
+                task("two", {"name": "two"}),
+                task("two", {"name": "two"}),
                 [
                     [
-                        (
-                            "checkpoint",
-                            {"values": {"name": "start > one > two 3 > two 4"}},
-                        ),
-                        ("task", {"name": "three", "result": {"name": "three"}}),
-                        (
-                            "checkpoint",
-                            {"values": {"name": "start > one > two 3 > two 4 > three"}},
-                        ),
+                        checkpoint({"name": "start > one > two > two"}),
+                        task("three", {"name": "three"}),
+                        checkpoint({"name": "start > one > two > two > three"}),
                     ],
                     [
-                        (
-                            "checkpoint",
-                            {"values": {"name": "start > one > two 3 > two 4"}},
-                        ),
-                        ("task", {"name": "three", "result": {"name": "three*"}}),
-                        (
-                            "checkpoint",
-                            {
-                                "values": {
-                                    "name": "start > one > two 3 > two 4 > three*"
-                                }
-                            },
-                        ),
+                        checkpoint({"name": "start > one > two > two"}),
+                        task("three", {"name": "three*"}),
+                        checkpoint({"name": "start > one > two > two > three*"}),
                     ],
                     [
-                        (
-                            "checkpoint",
-                            {"values": {"name": "start > one > two 3 > two 4"}},
-                        ),
-                        ("task", {"name": "three", "result": {"name": "three"}}),
-                        (
-                            "checkpoint",
-                            {"values": {"name": "start > one > two 3 > two 4 > three"}},
-                        ),
+                        checkpoint({"name": "start > one > two > two"}),
+                        task("three", {"name": "three"}),
+                        checkpoint({"name": "start > one > two > two > three"}),
                     ],
                 ],
             ],

libs/langgraph/tests/test_pregel_async.py

@@ -8770,3 +8770,351 @@ async def test_timeout_with_parent_command(
         await graph.ainvoke({"value": "start"}, thread1)
     assert exc_info.value.args[0].goto == "test_cmd"
     assert exc_info.value.args[0].update == {"key": "value"}
+
+
+async def test_fork_and_update_task_results(
+    async_checkpointer: BaseCheckpointSaver,
+) -> None:
+    """Test forking and updating task results with state history."""
+
+    def checkpoint(values: dict[str, Any]):
+        return ("checkpoint", {"values": values})
+
+    def task(name: str, result: Any):
+        return ("task", {"name": name, "result": result})
+
+    def get_tree(history: list[StateSnapshot]) -> list:
+        """Build a tree structure from state history for comparison."""
+        if not history:
+            return []
+
+        # Build a tree structure similar to renderForks
+        node_map: dict[str, dict] = {}
+        root_nodes: list[dict] = []
+
+        # Second pass: establish parent-child relationships
+        for item in reversed(history):
+            checkpoint_id = item.config["configurable"]["checkpoint_id"]
+            parent_checkpoint_id = (
+                item.parent_config["configurable"]["checkpoint_id"]
+                if item.parent_config
+                else None
+            )
+            node_map[checkpoint_id] = {"item": item, "children": []}
+
+            parent = node_map.get(parent_checkpoint_id)
+            (parent["children"] if parent else root_nodes).append(
+                node_map[checkpoint_id]
+            )
+
+        def node_to_tree(node: dict) -> list:
+            """Convert a node to tree structure."""
+            result = [
+                checkpoint(node["item"].values),
+            ] + [
+                task(task_info.name, task_info.result)
+                for task_info in node["item"].tasks
+            ]
+
+            if len(node["children"]) > 1:
+                branches = [node_to_tree(child) for child in node["children"]]
+                return result + [branches]
+            elif len(node["children"]) == 1:
+                return result + node_to_tree(node["children"][0])
+            else:
+                return result
+
+        if len(root_nodes) == 1:
+            # Process all root nodes
+            return node_to_tree(root_nodes[0])
+
+        elif len(root_nodes) > 1:
+            # Multiple root nodes - treat as branches
+            branches = [node_to_tree(node) for node in root_nodes]
+            return branches
+        else:
+            return []
+
+    class State(TypedDict):
+        name: Annotated[str, lambda a, b: " > ".join([a, b]) if a else b]
+
+    # Define the graph with a sequence of nodes
+    def one(state: State) -> Command:
+        return Command(goto=[Send("two", {})], update={"name": "one"})
+
+    def two(state: State) -> State:
+        return {"name": "two"}
+
+    def three(state: State) -> State:
+        return {"name": "three"}
+
+    graph = (
+        StateGraph(State)
+        .add_node("one", one)
+        .add_node("two", two)
+        .add_node("three", three)
+        .add_edge(START, "one")
+        .add_edge("one", "two")
+        .add_edge("two", "three")
+        .compile(checkpointer=async_checkpointer)
+    )
+
+    config = {"configurable": {"thread_id": "1"}}
+    history: list[StateSnapshot] = []
+
+    # Initial run
+    await graph.ainvoke({"name": "start"}, config)
+    history = [c async for c in graph.aget_state_history(config)]
+
+    assert get_tree(history) == [
+        checkpoint({"name": ""}),
+        task("__start__", {"name": "start"}),
+        checkpoint({"name": "start"}),
+        task("one", {"name": "one"}),
+        checkpoint({"name": "start > one"}),
+        task("two", {"name": "two"}),
+        task("two", {"name": "two"}),
+        checkpoint({"name": "start > one > two > two"}),
+        task("three", {"name": "three"}),
+        checkpoint({"name": "start > one > two > two > three"}),
+    ]
+
+    # Update the start state
+    await graph.ainvoke(
+        None,
+        await graph.aupdate_state(
+            history[4].config,
+            values=[StateUpdate(values={"name": "start*"}, as_node="__start__")],
+            as_node="__copy__",
+        ),
+    )
+
+    history = [c async for c in graph.aget_state_history(config)]
+    assert get_tree(history) == [
+        [
+            checkpoint({"name": ""}),
+            task("__start__", {"name": "start"}),
+            checkpoint({"name": "start"}),
+            task("one", {"name": "one"}),
+            checkpoint({"name": "start > one"}),
+            task("two", {"name": "two"}),
+            task("two", {"name": "two"}),
+            checkpoint({"name": "start > one > two > two"}),
+            task("three", {"name": "three"}),
+            checkpoint({"name": "start > one > two > two > three"}),
+        ],
+        [
+            checkpoint({"name": ""}),
+            task("__start__", {"name": "start*"}),
+            checkpoint({"name": "start*"}),
+            task("one", {"name": "one"}),
+            checkpoint({"name": "start* > one"}),
+            task("two", {"name": "two"}),
+            task("two", {"name": "two"}),
+            checkpoint({"name": "start* > one > two > two"}),
+            task("three", {"name": "three"}),
+            checkpoint({"name": "start* > one > two > two > three"}),
+        ],
+    ]
+
+    # Fork from task "one"
+    # Start from the checkpoint that has the task "one"
+    assert history[3].values == {"name": "start*"}
+    assert len(history[3].tasks) == 1
+    assert history[3].tasks[0].name == "one"
+
+    await graph.ainvoke(
+        None,
+        await graph.aupdate_state(
+            history[3].config,
+            [StateUpdate(values={"name": "one*"}, as_node="one")],
+            "__copy__",
+        ),
+    )
+
+    history = [c async for c in graph.aget_state_history(config)]
+    assert get_tree(history) == [
+        [
+            checkpoint({"name": ""}),
+            task("__start__", {"name": "start"}),
+            checkpoint({"name": "start"}),
+            task("one", {"name": "one"}),
+            checkpoint({"name": "start > one"}),
+            task("two", {"name": "two"}),
+            task("two", {"name": "two"}),
+            checkpoint({"name": "start > one > two > two"}),
+            task("three", {"name": "three"}),
+            checkpoint({"name": "start > one > two > two > three"}),
+        ],
+        [
+            checkpoint({"name": ""}),
+            task("__start__", {"name": "start*"}),
+            [
+                [
+                    checkpoint({"name": "start*"}),
+                    task("one", {"name": "one"}),
+                    checkpoint({"name": "start* > one"}),
+                    task("two", {"name": "two"}),
+                    task("two", {"name": "two"}),
+                    checkpoint({"name": "start* > one > two > two"}),
+                    task("three", {"name": "three"}),
+                    checkpoint({"name": "start* > one > two > two > three"}),
+                ],
+                [
+                    checkpoint({"name": "start*"}),
+                    task("one", {"name": "one*"}),
+                    checkpoint({"name": "start* > one*"}),
+                    task("two", {"name": "two"}),
+                    checkpoint({"name": "start* > one* > two"}),
+                    task("three", {"name": "three"}),
+                    checkpoint({"name": "start* > one* > two > three"}),
+                ],
+            ],
+        ],
+    ]
+
+    config = {"configurable": {"thread_id": "2"}}
+
+    # Initialize the thread once again
+    await graph.ainvoke({"name": "start"}, config)
+    history = [c async for c in graph.aget_state_history(config)]
+
+    # Fork from task "two"
+    # Start from the checkpoint that has the task "two"
+    assert history[2].values == {"name": "start > one"}
+
+    await graph.ainvoke(
+        None,
+        await graph.aupdate_state(
+            history[2].config,
+            [
+                StateUpdate(values={"name": "two"}, as_node="two"),
+                StateUpdate(values={"name": "two"}, as_node="two"),
+            ],
+            "__copy__",
+        ),
+    )
+
+    history = [c async for c in graph.aget_state_history(config)]
+    assert get_tree(history) == [
+        checkpoint({"name": ""}),
+        task("__start__", {"name": "start"}),
+        checkpoint({"name": "start"}),
+        task("one", {"name": "one"}),
+        [
+            [
+                checkpoint({"name": "start > one"}),
+                task("two", {"name": "two"}),
+                task("two", {"name": "two"}),
+                checkpoint({"name": "start > one > two > two"}),
+                task("three", {"name": "three"}),
+                checkpoint({"name": "start > one > two > two > three"}),
+            ],
+            [
+                checkpoint({"name": "start > one"}),
+                task("two", {"name": "two"}),
+                task("two", {"name": "two"}),
+                checkpoint({"name": "start > one > two > two"}),
+                task("three", {"name": "three"}),
+                checkpoint({"name": "start > one > two > two > three"}),
+            ],
+        ],
+    ]
+
+    # Fork task three
+    assert history[1].values == {"name": "start > one > two > two"}
+    assert len(history[1].tasks) == 1
+    assert history[1].tasks[0].name == "three"
+
+    await graph.ainvoke(
+        None,
+        await graph.aupdate_state(
+            history[1].config,
+            [StateUpdate(values={"name": "three*"}, as_node="three")],
+            "__copy__",
+        ),
+    )
+
+    history = [c async for c in graph.aget_state_history(config)]
+    assert get_tree(history) == [
+        checkpoint({"name": ""}),
+        task("__start__", {"name": "start"}),
+        checkpoint({"name": "start"}),
+        task("one", {"name": "one"}),
+        [
+            [
+                checkpoint({"name": "start > one"}),
+                task("two", {"name": "two"}),
+                task("two", {"name": "two"}),
+                checkpoint({"name": "start > one > two > two"}),
+                task("three", {"name": "three"}),
+                checkpoint({"name": "start > one > two > two > three"}),
+            ],
+            [
+                checkpoint({"name": "start > one"}),
+                task("two", {"name": "two"}),
+                task("two", {"name": "two"}),
+                [
+                    [
+                        checkpoint({"name": "start > one > two > two"}),
+                        task("three", {"name": "three"}),
+                        checkpoint({"name": "start > one > two > two > three"}),
+                    ],
+                    [
+                        checkpoint({"name": "start > one > two > two"}),
+                        task("three", {"name": "three*"}),
+                        checkpoint({"name": "start > one > two > two > three*"}),
+                    ],
+                ],
+            ],
+        ],
+    ]
+
+    # Regenerate task three
+    assert history[3].values == {"name": "start > one > two > two"}
+    assert len(history[3].tasks) == 1
+    assert history[3].tasks[0].name == "three"
+
+    await graph.ainvoke(
+        None, await graph.aupdate_state(history[3].config, None, "__copy__")
+    )
+
+    history = [c async for c in graph.aget_state_history(config)]
+    assert get_tree(history) == [
+        checkpoint({"name": ""}),
+        task("__start__", {"name": "start"}),
+        checkpoint({"name": "start"}),
+        task("one", {"name": "one"}),
+        [
+            [
+                checkpoint({"name": "start > one"}),
+                task("two", {"name": "two"}),
+                task("two", {"name": "two"}),
+                checkpoint({"name": "start > one > two > two"}),
+                task("three", {"name": "three"}),
+                checkpoint({"name": "start > one > two > two > three"}),
+            ],
+            [
+                checkpoint({"name": "start > one"}),
+                task("two", {"name": "two"}),
+                task("two", {"name": "two"}),
+                [
+                    [
+                        checkpoint({"name": "start > one > two > two"}),
+                        task("three", {"name": "three"}),
+                        checkpoint({"name": "start > one > two > two > three"}),
+                    ],
+                    [
+                        checkpoint({"name": "start > one > two > two"}),
+                        task("three", {"name": "three*"}),
+                        checkpoint({"name": "start > one > two > two > three*"}),
+                    ],
+                    [
+                        checkpoint({"name": "start > one > two > two"}),
+                        task("three", {"name": "three"}),
+                        checkpoint({"name": "start > one > two > two > three"}),
+                    ],
+                ],
+            ],
+        ],
+    ]