From d18da208b9df41aac65c699d4b18bc2431255611 Mon Sep 17 00:00:00 2001 From: Caspar Broekhuizen Date: Fri, 24 Oct 2025 18:00:28 -0700 Subject: [PATCH] refactor(langgraph): add guards and rename constant --- .../langgraph/_internal/_constants.py | 2 +- libs/langgraph/langgraph/pregel/_algo.py | 20 +++++++++----- libs/langgraph/langgraph/pregel/_loop.py | 26 +++++++++++-------- 3 files changed, 29 insertions(+), 19 deletions(-) diff --git a/libs/langgraph/langgraph/_internal/_constants.py b/libs/langgraph/langgraph/_internal/_constants.py index aebeb95e4..1553acc89 100644 --- a/libs/langgraph/langgraph/_internal/_constants.py +++ b/libs/langgraph/langgraph/_internal/_constants.py @@ -78,7 +78,7 @@ CONF = cast(Literal["configurable"], sys.intern("configurable")) # key for the configurable dict in RunnableConfig NULL_TASK_ID = sys.intern("00000000-0000-0000-0000-000000000000") # the task_id to use for writes that are not associated with a task -RUNTIME_PLACEHOLDER = "__pregel_runtime_placeholder__" +UNTRACKED_VALUE_PLACEHOLDER = sys.intern("__pregel_untracked_value_placeholder") # placeholder for untracked values replaced at runtime # redefined to avoid circular import with langgraph.constants diff --git a/libs/langgraph/langgraph/pregel/_algo.py b/libs/langgraph/langgraph/pregel/_algo.py index ef7d31898..e54cd8ac8 100644 --- a/libs/langgraph/langgraph/pregel/_algo.py +++ b/libs/langgraph/langgraph/pregel/_algo.py @@ -57,8 +57,8 @@ from langgraph._internal._constants import ( RESERVED, RESUME, RETURN, - RUNTIME_PLACEHOLDER, TASKS, + UNTRACKED_VALUE_PLACEHOLDER, ) from langgraph._internal._scratchpad import PregelScratchpad from langgraph._internal._typing import EMPTY_SEQ, MISSING @@ -641,8 +641,14 @@ def prepare_single_task( f"Ignoring invalid packet type {type(packet)} in pending sends" ) return - # Replace runtime placeholders with untracked values - packet = rehydrate_untracked_values_in_send(packet, channels) + + # Check if any channels are UntrackedValue - if true, some + # untracked values may have been replaced with runtime placeholders + if any( + isinstance(channel, UntrackedValue) for channel in channels.values() + ): + # Replace runtime placeholders with untracked values + packet = rehydrate_untracked_values_in_send(packet, channels) if packet.node not in processes: logger.warning( @@ -1116,7 +1122,7 @@ class LazyAtomicCounter: def sanitize_untracked_values_in_send( packet: Send, channels: Mapping[str, BaseChannel] ) -> Send: - """Replace any UntrackedValue contents in Send.arg with RUNTIME_PLACEHOLDER for checkpointing. + """Replace any UntrackedValue contents in Send.arg with UNTRACKED_VALUE_PLACEHOLDER for checkpointing. Send is not typed and arg may be a nested dict.""" @@ -1130,7 +1136,7 @@ def sanitize_untracked_values_in_send( # arg can be nested dicts v = replace(v) if isinstance(channels.get(k), UntrackedValue): - obj[k] = RUNTIME_PLACEHOLDER + obj[k] = UNTRACKED_VALUE_PLACEHOLDER return obj sanitized_arg = replace(packet.arg) @@ -1140,7 +1146,7 @@ def sanitize_untracked_values_in_send( def rehydrate_untracked_values_in_send( packet: Send, channels: Mapping[str, BaseChannel] ) -> Send: - """Replace RUNTIME_PLACEHOLDERs in Send.arg with actual untracked values from UntrackedValue channels.""" + """Replace UNTRACKED_VALUE_PLACEHOLDER in Send.arg with actual untracked values from UntrackedValue channels.""" if not isinstance(packet.arg, dict): # Command @@ -1155,7 +1161,7 @@ def rehydrate_untracked_values_in_send( # arg can be nested dicts v = replace(v) if ( - v is RUNTIME_PLACEHOLDER + v == UNTRACKED_VALUE_PLACEHOLDER and k in channels and isinstance(channels[k], UntrackedValue) ): diff --git a/libs/langgraph/langgraph/pregel/_loop.py b/libs/langgraph/langgraph/pregel/_loop.py index 5d8491489..27ed2a48b 100644 --- a/libs/langgraph/langgraph/pregel/_loop.py +++ b/libs/langgraph/langgraph/pregel/_loop.py @@ -325,17 +325,21 @@ class PregelLoop: ] writes_to_save = writes - # We never want to persist untracked values in checkpoints - # because there is no guarantee that they are serializable - writes_to_save = [ - # Sanitize UntrackedValues that are nested within Send packets - (c, sanitize_untracked_values_in_send(v, self.channels)) - if c == TASKS and isinstance(v, Send) - else (c, v) - for c, v in writes_to_save - # Do not persist UntrackedValue channel writes - if not isinstance(self.specs.get(c), UntrackedValue) - ] + # Check if any writes are to an UntrackedValue channel + if any( + isinstance(channel, UntrackedValue) for channel in self.channels.values() + ): + # We never want to persist untracked values in checkpoints + # because there is no guarantee that they are serializable + writes_to_save = [ + # Sanitize UntrackedValues that are nested within Send packets + (c, sanitize_untracked_values_in_send(v, self.channels)) + if c == TASKS and isinstance(v, Send) + else (c, v) + for c, v in writes_to_save + # Do not persist UntrackedValue channel writes + if not isinstance(self.specs.get(c), UntrackedValue) + ] # save writes self.checkpoint_pending_writes.extend((task_id, c, v) for c, v in writes)