refactor(langgraph): simplify Send sanitization with dict comprehension

Co-authored-by: sydney-runkle <54324534+sydney-runkle@users.noreply.github.com>

libs/langgraph/langgraph/pregel/_algo.py

@@ -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
@@ -639,6 +640,7 @@ def prepare_single_task(
                     f"Ignoring invalid packet type {type(packet)} in pending sends"
                 )
                 return
+
             if packet.node not in processes:
                 logger.warning(
                     f"Ignoring unknown node name {packet.node} in pending sends"
@@ -1106,3 +1108,23 @@ 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:
+    """Pop any UntrackedValue contents in Send.arg for safe checkpointing.
+
+    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
+        return packet
+
+    sanitized_arg = {
+        k: v
+        for k, v in packet.arg.items()
+        if not isinstance(channels.get(k), UntrackedValue)
+    }
+
+    return Send(node=packet.node, arg=sanitized_arg)

libs/langgraph/langgraph/pregel/_loop.py

@@ -56,10 +56,12 @@ 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
 from langgraph.channels.base import BaseChannel
+from langgraph.channels.untracked_value import UntrackedValue
 from langgraph.constants import TAG_HIDDEN
 from langgraph.errors import (
     EmptyInputError,
@@ -78,6 +80,7 @@ from langgraph.pregel._algo import (
     increment,
     prepare_next_tasks,
     prepare_single_task,
+    sanitize_untracked_values_in_send,
     should_interrupt,
     task_path_str,
 )
@@ -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
+
+        # 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
+            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
+                # dont 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)
         if self.durability != "exit" and self.checkpointer_put_writes is not None:
@@ -735,6 +757,17 @@ class PregelLoop:
             id=self.checkpoint["id"] if exiting else None,
             updated_channels=self.updated_channels,
         )
+        # 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()
+        ):
+            sanitized_tasks = [
+                sanitize_untracked_values_in_send(value, self.channels)
+                if isinstance(value, Send)
+                else value
+                for value in self.checkpoint["channel_values"][TASKS]
+            ]
+            self.checkpoint["channel_values"][TASKS] = sanitized_tasks
         # bail if no checkpointer
         if do_checkpoint and self._checkpointer_put_after_previous is not None:
             self.prev_checkpoint_config = (

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

@@ -44,6 +44,7 @@ from langgraph.channels.binop import BinaryOperatorAggregate
 from langgraph.channels.ephemeral_value import EphemeralValue
 from langgraph.channels.last_value import LastValue
 from langgraph.channels.topic import Topic
+from langgraph.channels.untracked_value import UntrackedValue
 from langgraph.config import get_stream_writer
 from langgraph.errors import GraphRecursionError, InvalidUpdateError, ParentCommand
 from langgraph.func import entrypoint, task
@@ -8597,3 +8598,102 @@ 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)
+
+    app = graph.compile(checkpointer=sync_checkpointer)
+    config = {"configurable": {"thread_id": "1"}}
+    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"
+
+    state = app.get_state(config)
+    assert "session_resource" not in state.values
+
+
+def test_send_with_untracked_value_overlapping_keys(
+    sync_checkpointer: BaseCheckpointSaver,
+):
+    """Test that Send objects work correctly with untracked values in state."""
+
+    class State(TypedDict):
+        dictionary: dict
+        session_resource: Annotated[str, UntrackedValue]
+
+    def setup_node(state: State) -> State:
+        return {}
+
+    def send_to_tool(state: State):
+        return [
+            Send(
+                "tool_node",
+                {
+                    "dictionary": {"session_resource": "legal_value"},
+                    "session_resource": "illegal_value",
+                },
+            )
+        ]
+
+    def tool_node(state: State) -> State:
+        print(f"STATE: {state}")
+        assert state["dictionary"] == {"session_resource": "legal_value"}
+        assert state["session_resource"] == "illegal_value"
+
+        return {
+            "dictionary": state["dictionary"],
+            "session_resource": "new_illegal_value",
+        }
+
+    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)
+
+    app = graph.compile(checkpointer=sync_checkpointer)
+    config = {"configurable": {"thread_id": "1"}}
+    result = app.invoke({}, config)
+
+    assert result["session_resource"] == "new_illegal_value"
+    state = app.get_state(config)
+    assert "session_resource" not in state.values
+    assert state.values.get("dictionary") == {"session_resource": "legal_value"}