diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py
index 173304ed2..8419eb913 100644
--- a/libs/langgraph/langgraph/pregel/__init__.py
+++ b/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)
diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py
index c415c2cb3..1403ee90c 100644
--- a/libs/langgraph/tests/test_pregel_async.py
+++ b/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"}},