fix(langgraph): dont persist UntrackedValue

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)

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:

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):

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(

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
 

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()

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
+    )