diff --git a/libs/langgraph/langgraph/pregel/_loop.py b/libs/langgraph/langgraph/pregel/_loop.py index 371037768..d1082d2d5 100644 --- a/libs/langgraph/langgraph/pregel/_loop.py +++ b/libs/langgraph/langgraph/pregel/_loop.py @@ -437,9 +437,7 @@ class PregelLoop: writes_to_save = writes # check if any writes are to an UntrackedValue channel - if any( - isinstance(channel, UntrackedValue) for channel in self.channels.values() - ): + if self._has_untracked_channels: # we do not persist untracked values in checkpoints writes_to_save = [ # sanitize UntrackedValues that are nested within Send packets @@ -1161,9 +1159,7 @@ class PregelLoop: elif "counters_since_delta_snapshot" in self.checkpoint_metadata: del self.checkpoint_metadata["counters_since_delta_snapshot"] # sanitize TASK channel in the checkpoint before saving (durability=="exit") - if TASKS in self.checkpoint["channel_values"] and any( - isinstance(channel, UntrackedValue) for channel in self.channels.values() - ): + if TASKS in self.checkpoint["channel_values"] and self._has_untracked_channels: sanitized_tasks = [ sanitize_untracked_values_in_send(value, self.channels) if isinstance(value, Send) @@ -1695,6 +1691,9 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager): saver=self.checkpointer, config=self.checkpoint_config, ) + self._has_untracked_channels = any( + isinstance(ch, UntrackedValue) for ch in self.channels.values() + ) self.stack.push(self._suppress_interrupt) self.status = "input" self.step = self.checkpoint_metadata["step"] + 1 @@ -1955,6 +1954,9 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager): saver=self.checkpointer, config=self.checkpoint_config, ) + self._has_untracked_channels = any( + isinstance(ch, UntrackedValue) for ch in self.channels.values() + ) self.stack.push(self._suppress_interrupt) self.status = "input" self.step = self.checkpoint_metadata["step"] + 1