diff --git a/libs/langgraph/langgraph/_internal/_constants.py b/libs/langgraph/langgraph/_internal/_constants.py index 1553acc89..abab56295 100644 --- a/libs/langgraph/langgraph/_internal/_constants.py +++ b/libs/langgraph/langgraph/_internal/_constants.py @@ -24,7 +24,6 @@ PREVIOUS = sys.intern("__previous__") # --- Reserved cache namespaces --- CACHE_NS_WRITES = sys.intern("__pregel_ns_writes") - # cache namespace for node writes # --- Reserved config.configurable keys --- @@ -78,8 +77,6 @@ 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 -UNTRACKED_VALUE_PLACEHOLDER = sys.intern("__pregel_untracked_value_placeholder") -# placeholder for untracked values replaced at runtime # redefined to avoid circular import with langgraph.constants _TAG_HIDDEN = sys.intern("langsmith:hidden") diff --git a/libs/langgraph/langgraph/pregel/_algo.py b/libs/langgraph/langgraph/pregel/_algo.py index 644c33f66..08b5b3cc0 100644 --- a/libs/langgraph/langgraph/pregel/_algo.py +++ b/libs/langgraph/langgraph/pregel/_algo.py @@ -58,7 +58,6 @@ from langgraph._internal._constants import ( RESUME, RETURN, TASKS, - UNTRACKED_VALUE_PLACEHOLDER, ) from langgraph._internal._scratchpad import PregelScratchpad from langgraph._internal._typing import EMPTY_SEQ, MISSING @@ -1114,9 +1113,9 @@ class LazyAtomicCounter: def sanitize_untracked_values_in_send( packet: Send, channels: Mapping[str, BaseChannel] ) -> Send: - """Replace any UntrackedValue contents in Send.arg with UNTRACKED_VALUE_PLACEHOLDER for checkpointing. + """Pop any UntrackedValue contents in Send.arg for safe checkpointing. - Send is not typed and arg may be a nested dict.""" + Send is not typed and arg may be a nested dict. We only look at the top level.""" if not isinstance(packet.arg, dict): # Command @@ -1124,8 +1123,12 @@ def sanitize_untracked_values_in_send( sanitized_arg = dict(packet.arg) + # top level keys should be the channel names + to_pop = set() for k, v in sanitized_arg.items(): if isinstance(channels.get(k), UntrackedValue): - sanitized_arg[k] = UNTRACKED_VALUE_PLACEHOLDER + to_pop.add(k) + for k in to_pop: + sanitized_arg.pop(k) return Send(node=packet.node, arg=sanitized_arg) diff --git a/libs/langgraph/langgraph/pregel/_loop.py b/libs/langgraph/langgraph/pregel/_loop.py index 953e893f8..4ff5d2ad9 100644 --- a/libs/langgraph/langgraph/pregel/_loop.py +++ b/libs/langgraph/langgraph/pregel/_loop.py @@ -325,20 +325,20 @@ class PregelLoop: ] writes_to_save = writes - # Check if any writes are to an UntrackedValue channel + # check if any writes are to an UntrackedValue channel if any( isinstance(channel, UntrackedValue) for channel in self.channels.values() ): - # We do not persist untracked values in checkpoints + # we do not persist untracked values in checkpoints writes_to_save = [ - # Sanitize UntrackedValues that are nested within Send packets + # 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 + # dont persist UntrackedValue channel writes if not isinstance(self.specs.get(c), UntrackedValue) ] @@ -757,7 +757,7 @@ class PregelLoop: id=self.checkpoint["id"] if exiting else None, updated_channels=self.updated_channels, ) - # sanitize TASK channel in the checkpoint before saving + # 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() ): diff --git a/libs/langgraph/langgraph/pregel/_write.py b/libs/langgraph/langgraph/pregel/_write.py index d99cbbd3a..8b4508257 100644 --- a/libs/langgraph/langgraph/pregel/_write.py +++ b/libs/langgraph/langgraph/pregel/_write.py @@ -176,8 +176,6 @@ def _assemble_writes( tuples: list[tuple[str, Any]] = [] for w in writes: if isinstance(w, Send): - # Send packets go to TASKS channel; sanitation for storage - # is handled centrally where channel specs are available. tuples.append((TASKS, w)) elif isinstance(w, ChannelWriteTupleEntry): if ww := w.mapper(w.value):