From bbe8346c8b20e4bf6878263ac333c337ea986144 Mon Sep 17 00:00:00 2001 From: Caspar Broekhuizen Date: Mon, 20 Oct 2025 16:53:47 -0700 Subject: [PATCH] fix(langgraph): dont persist UntrackedValue --- libs/langgraph/langgraph/pregel/_algo.py | 49 +++++++++++++++++++- libs/langgraph/langgraph/pregel/_loop.py | 22 +++++++++ libs/langgraph/langgraph/pregel/_write.py | 2 + libs/langgraph/langgraph/pregel/main.py | 10 ++++- libs/langgraph/langgraph/types.py | 3 ++ libs/langgraph/tests/test_channels.py | 30 +++++++++++++ libs/langgraph/tests/test_pregel.py | 54 +++++++++++++++++++++++ 7 files changed, 168 insertions(+), 2 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/_algo.py b/libs/langgraph/langgraph/pregel/_algo.py index 98f25d4ac..caf5bbc95 100644 --- a/libs/langgraph/langgraph/pregel/_algo.py +++ b/libs/langgraph/langgraph/pregel/_algo.py @@ -6,7 +6,7 @@ import sys import threading from collections import defaultdict, deque from collections.abc import Callable, Iterable, Mapping, Sequence -from copy import copy +from copy import copy, deepcopy from functools import partial from hashlib import sha1 from typing import ( @@ -63,6 +63,7 @@ from langgraph._internal._scratchpad import PregelScratchpad from langgraph._internal._typing import EMPTY_SEQ, MISSING from langgraph.channels.base import BaseChannel from langgraph.channels.topic import Topic +from langgraph.channels.untracked_value import UntrackedValue from langgraph.constants import TAG_HIDDEN from langgraph.managed.base import ManagedValueMapping from langgraph.pregel._call import get_runnable_for_task, identifier @@ -71,6 +72,7 @@ from langgraph.pregel._log import logger from langgraph.pregel._read import INPUT_CACHE_KEY_TYPE, PregelNode from langgraph.runtime import DEFAULT_RUNTIME, Runtime from langgraph.types import ( + RUNTIME_PLACEHOLDER, All, CacheKey, CachePolicy, @@ -639,6 +641,9 @@ 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) + if packet.node not in processes: logger.warning( f"Ignoring unknown node name {packet.node} in pending sends" @@ -1106,3 +1111,45 @@ class LazyAtomicCounter: if self._counter is None: self._counter = itertools.count(0).__next__ return self._counter() + +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. + + Send is not typed and arg may be a nested dict.""" + + if not isinstance(packet.arg, dict): + # Command + return packet + + def replace(obj: dict[str, Any]) -> dict[str, Any]: + for k, v in obj.items(): + if isinstance(v, dict): + # arg can be nested dicts + v = replace(v) + if isinstance(channels.get(k), UntrackedValue): + obj[k] = RUNTIME_PLACEHOLDER + return obj + + sanitized_arg = replace(packet.arg) + return Send(node=packet.node, arg=sanitized_arg) + +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.""" + + if not isinstance(packet.arg, dict): + # Command + return packet + + # deepcopy to avoid mutating the original packet, as it is later persisted in checkpoints + arg_deepcopy = deepcopy(packet.arg) + def replace(obj: dict[str, Any]) -> dict[str, Any]: + for k, v in obj.items(): + if isinstance(v, dict): + # arg can be nested dicts + v = replace(v) + if v is RUNTIME_PLACEHOLDER and k in channels and isinstance(channels[k], UntrackedValue): + obj[k] = channels[k].get() + return obj + + rehydrated_arg = replace(arg_deepcopy) + return Send(node=packet.node, arg=rehydrated_arg) diff --git a/libs/langgraph/langgraph/pregel/_loop.py b/libs/langgraph/langgraph/pregel/_loop.py index 97a937ccb..ea2a3ad26 100644 --- a/libs/langgraph/langgraph/pregel/_loop.py +++ b/libs/langgraph/langgraph/pregel/_loop.py @@ -20,6 +20,7 @@ from typing import ( TypeVar, cast, ) +from langgraph.channels.untracked_value import UntrackedValue from langchain_core.callbacks import AsyncParentRunManager, ParentRunManager from langchain_core.runnables import RunnableConfig @@ -56,6 +57,7 @@ from langgraph._internal._constants import ( NULL_TASK_ID, PUSH, RESUME, + TASKS, ) from langgraph._internal._scratchpad import PregelScratchpad from langgraph._internal._typing import EMPTY_SEQ, MISSING @@ -73,6 +75,7 @@ from langgraph.pregel._algo import ( Call, GetNextVersion, PregelTaskWrites, + sanitize_untracked_values_in_send, apply_writes, checkpoint_null_version, increment, @@ -114,6 +117,7 @@ from langgraph.types import ( Durability, PregelExecutableTask, RetryPolicy, + Send, StreamMode, ) @@ -320,6 +324,24 @@ class PregelLoop: w for w in self.checkpoint_pending_writes if w[0] != task_id ] writes_to_save = writes + + # We never want to persist untracked values in checkpoints + # because there is no guarantee that they are serializable + def _sanitize(group: WritesT) -> WritesT: + out: WritesT = [] + for c, v in group: + # Do not persist UntrackedValue channel writes + if isinstance(self.specs.get(c), UntrackedValue): + continue + # Sanitize UntrackedValues that are nested within Send packets + if c == TASKS and isinstance(v, Send): + out.append((c, sanitize_untracked_values_in_send(v, self.channels))) + else: + out.append((c, v)) + return out + + writes_to_save = _sanitize(writes_to_save) + # save writes self.checkpoint_pending_writes.extend((task_id, c, v) for c, v in writes) if self.durability != "exit" and self.checkpointer_put_writes is not None: diff --git a/libs/langgraph/langgraph/pregel/_write.py b/libs/langgraph/langgraph/pregel/_write.py index 8b4508257..d99cbbd3a 100644 --- a/libs/langgraph/langgraph/pregel/_write.py +++ b/libs/langgraph/langgraph/pregel/_write.py @@ -176,6 +176,8 @@ 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): diff --git a/libs/langgraph/langgraph/pregel/main.py b/libs/langgraph/langgraph/pregel/main.py index e0243c41c..ea3309587 100644 --- a/libs/langgraph/langgraph/pregel/main.py +++ b/libs/langgraph/langgraph/pregel/main.py @@ -106,6 +106,7 @@ from langgraph.errors import ( from langgraph.managed.base import ManagedValueSpec from langgraph.pregel._algo import ( PregelTaskWrites, + sanitize_untracked_values_in_send, _scratchpad, apply_writes, local_read, @@ -2307,8 +2308,15 @@ class Pregel( # channel writes are saved to current checkpoint channel_writes = [w for w in task.writes if w[0] != PUSH] if saved and channel_writes: + # sanitize TASKS writes for storage + sanitized = [ + (c, sanitize_untracked_values_in_send(v, channels)) + if c == TASKS and isinstance(v, Send) + else (c, v) + for c, v in channel_writes + ] await checkpointer.aput_writes( - checkpoint_config, channel_writes, task_id + checkpoint_config, sanitized, task_id ) # apply to checkpoint and save apply_writes( diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index a060e1446..a5b80c662 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -27,6 +27,9 @@ from langgraph._internal._retry import default_retry_on from langgraph._internal._typing import MISSING, DeprecatedKwargs from langgraph.warnings import LangGraphDeprecatedSinceV10 +# placeholder for untracked values replaced at runtime +RUNTIME_PLACEHOLDER = "__pregel_runtime_placeholder__" + if TYPE_CHECKING: from langgraph.pregel.protocol import PregelProtocol diff --git a/libs/langgraph/tests/test_channels.py b/libs/langgraph/tests/test_channels.py index 36037a4ee..37a5eb432 100644 --- a/libs/langgraph/tests/test_channels.py +++ b/libs/langgraph/tests/test_channels.py @@ -7,6 +7,7 @@ from langgraph._internal._typing import MISSING from langgraph.channels.binop import BinaryOperatorAggregate from langgraph.channels.last_value import LastValue from langgraph.channels.topic import Topic +from langgraph.channels.untracked_value import UntrackedValue from langgraph.errors import EmptyChannelError, InvalidUpdateError pytestmark = pytest.mark.anyio @@ -87,3 +88,32 @@ def test_binop() -> None: checkpoint = channel.checkpoint() channel = BinaryOperatorAggregate(int, operator.add).from_checkpoint(checkpoint) assert channel.get() == 10 + + +def test_untracked_value() -> None: + channel = UntrackedValue(dict).from_checkpoint(MISSING) + assert channel.ValueType is dict + assert channel.UpdateType is dict + + # UntrackedValue should start empty + with pytest.raises(EmptyChannelError): + channel.get() + + # Should be able to update with a value + test_data = {"session": "test", "temp": "dir"} + channel.update([test_data]) + assert channel.get() == test_data + + # Update with new value + new_data = {"session": "updated", "temp": "newdir"} + channel.update([new_data]) + assert channel.get() == new_data + + # On checkpoint, UntrackedValue should return MISSING + checkpoint = channel.checkpoint() + assert checkpoint is MISSING + + # Creating from checkpoint with MISSING should start empty + new_channel = UntrackedValue(dict).from_checkpoint(checkpoint) + with pytest.raises(EmptyChannelError): + new_channel.get() diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 93c3e0d0c..8ee148d04 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -13,6 +13,7 @@ from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass, field from random import randrange from typing import Annotated, Any, Literal, get_type_hints +from langgraph.channels.untracked_value import UntrackedValue import pytest from langchain_core.language_models import GenericFakeChatModel @@ -8597,3 +8598,56 @@ def test_multiple_writes_same_channel_from_same_node( "values": {"foo": ""}, }, ] + +def test_send_with_untracked_value(sync_checkpointer: BaseCheckpointSaver): + """Test that Send objects work correctly with untracked values in state.""" + + class UnserializableResource: + def __init__(self, name: str): + self.name = name + self.lock = threading.Lock() + + class State(TypedDict): + messages: Annotated[list[str], operator.add] + session_resource: Annotated[UnserializableResource, UntrackedValue] + + def setup_node(state: State) -> State: + resource = UnserializableResource("test_session") + return {"messages": ["setup complete"], "session_resource": resource} + + def send_to_tool(state: State): + return [Send("tool_node", state)] + + def tool_node(state: State) -> State: + resource = state["session_resource"] + assert isinstance(resource, UnserializableResource) + assert resource.name == "test_session" + + new_resource = UnserializableResource("new_session") + + return { + "messages": [f"tool used resource: {resource.name}"], + "session_resource": new_resource, + } + + graph = StateGraph(State) + graph.add_node("setup", setup_node) + graph.add_node("tool_node", tool_node) + graph.add_edge(START, "setup") + graph.add_conditional_edges("setup", send_to_tool) + graph.add_edge("tool_node", END) + + app = graph.compile(checkpointer=sync_checkpointer) + config = {"configurable": {"thread_id": "test_thread"}} + result = app.invoke({}, config) + + assert len(result["messages"]) == 2 + assert result["messages"][0] == "setup complete" + assert result["messages"][1] == "tool used resource: test_session" + assert result["session_resource"].name == "new_session" + + # Check that the untracked resource is NOT in the final state checkpoint + state = app.get_state(config) + assert ( + "session_resource" not in state.values + )