Update test

libs/langgraph/langgraph/pregel/__init__.py

@@ -917,7 +917,7 @@ class Pregel(
             }
         # find last node that updated the state, if not provided
         if values is None and as_node is None:
-            return await checkpointer.aput(
+            next_config = await checkpointer.aput(
                 checkpoint_config,
                 create_checkpoint(checkpoint, None, step),
                 {
@@ -928,6 +928,7 @@ class Pregel(
                 },
                 {},
             )
+            return patch_checkpoint_map(next_config, saved.metadata if saved else None)
         elif as_node is None and not saved:
             if (
                 isinstance(self.input_channels, str)

libs/langgraph/tests/test_pregel_async.py

@@ -8303,8 +8303,8 @@ async def test_doubly_nested_graph_state(
     ]
 
     # fork and replay
-    fork = app.update_state(grandchild_history[2].config, None)
-    assert [c for c in app.stream(None, fork, subgraphs=True)] == [
+    fork = await app.aupdate_state(grandchild_history[2].config, None)
+    assert [c async for c in app.astream(None, fork, subgraphs=True)] == [
         (
             (AnyStr("child:"), AnyStr("child_1:")),
             {"grandchild_1": {"my_key": "hi my value here"}},