From 239df52e74a0cf3b88db1ada7a9e5b58540cbec5 Mon Sep 17 00:00:00 2001 From: Tat Dat Duong Date: Wed, 25 Jun 2025 02:56:53 +0200 Subject: [PATCH] Avoid creating new root checkpoint to preserve behaviour --- libs/langgraph/langgraph/pregel/__init__.py | 22 +++++++++++++++------ 1 file changed, 16 insertions(+), 6 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index d6cc3e432..fd65810c5 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -1555,12 +1555,17 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou raise InvalidUpdateError("Cannot copy a non-existent checkpoint") next_checkpoint = create_checkpoint(checkpoint, None, step) + copy_and_apply_tasks = isinstance(values, list) and len(values) > 0 # copy checkpoint next_config = checkpointer.put( saved.parent_config - or patch_configurable( - saved.config, {CONFIG_KEY_CHECKPOINT_ID: None} + or ( + patch_configurable( + saved.config, {CONFIG_KEY_CHECKPOINT_ID: None} + ) + if copy_and_apply_tasks + else saved.config ), next_checkpoint, { @@ -1573,7 +1578,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou # 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: + if copy_and_apply_tasks: # figure out the task IDs for the next update checkpoint next_tasks = prepare_next_tasks( next_checkpoint, @@ -2008,12 +2013,17 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou raise InvalidUpdateError("Cannot copy a non-existent checkpoint") next_checkpoint = create_checkpoint(checkpoint, None, step) + copy_and_apply_tasks = isinstance(values, list) and len(values) > 0 # copy checkpoint next_config = await checkpointer.aput( saved.parent_config - or patch_configurable( - saved.config, {CONFIG_KEY_CHECKPOINT_ID: None} + or ( + patch_configurable( + saved.config, {CONFIG_KEY_CHECKPOINT_ID: None} + ) + if copy_and_apply_tasks + else saved.config ), next_checkpoint, { @@ -2026,7 +2036,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou # 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: + if copy_and_apply_tasks: # figure out the task IDs for the next update checkpoint next_tasks = prepare_next_tasks( next_checkpoint,