From c5ec568cfbf9dcedc26a7d6de699c564383d0479 Mon Sep 17 00:00:00 2001
From: Tat Dat Duong <david@duong.cz>
Date: Wed, 9 Oct 2024 17:58:03 +0200
Subject: [PATCH] Patch config before entering map_debug_checkpoint

---
 libs/langgraph/langgraph/pregel/loop.py | 13 +++++++++----
 1 file changed, 9 insertions(+), 4 deletions(-)

diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py
index dadc54b7a..e472bdd1a 100644
--- a/libs/langgraph/langgraph/pregel/loop.py
+++ b/libs/langgraph/langgraph/pregel/loop.py
@@ -101,7 +101,7 @@ from langgraph.pregel.read import PregelNode
 from langgraph.pregel.utils import get_new_channel_versions
 from langgraph.store.base import BaseStore
 from langgraph.types import All, PregelExecutableTask, StreamMode
-from langgraph.utils.config import patch_configurable
+from langgraph.utils.config import patch_checkpoint_map, patch_configurable
 
 V = TypeVar("V")
 P = ParamSpec("P")
@@ -176,6 +176,7 @@ class PregelLoop:
     checkpoint_pending_writes: List[PendingWrite]
     checkpoint_previous_versions: dict[str, Union[str, float, int]]
     prev_checkpoint_config: Optional[RunnableConfig]
+    prev_checkpoint_metadata: Optional[CheckpointMetadata]
 
     step: int
     stop: int
@@ -362,14 +363,16 @@ class PregelLoop:
                 "debug",
                 map_debug_checkpoint,
                 self.step - 1,  # printing checkpoint for previous step
-                self.checkpoint_config,
+                patch_checkpoint_map(self.checkpoint_config, self.checkpoint_metadata),
                 self.channels,
                 self.stream_keys,
                 self.checkpoint_metadata,
                 self.checkpoint,
                 self.tasks.values(),
                 self.checkpoint_pending_writes,
-                self.prev_checkpoint_config,
+                patch_checkpoint_map(
+                    self.prev_checkpoint_config, self.prev_checkpoint_metadata
+                ),
                 self.output_keys,
             )
 
@@ -518,13 +521,15 @@ class PregelLoop:
         self.checkpoint = create_checkpoint(self.checkpoint, self.channels, self.step)
         # bail if no checkpointer
         if self._checkpointer_put_after_previous is not None:
+            self.prev_checkpoint_metadata = self.checkpoint_metadata
+            self.checkpoint_metadata = metadata
+
             self.prev_checkpoint_config = (
                 self.checkpoint_config
                 if CONFIG_KEY_CHECKPOINT_ID in self.checkpoint_config[CONF]
                 and self.checkpoint_config[CONF][CONFIG_KEY_CHECKPOINT_ID]
                 else None
             )
-            self.checkpoint_metadata = metadata
             self.checkpoint_config = {
                 **self.checkpoint_config,
                 CONF: {