refactor(langgraph): clean up code, remove rehydration logic, add test

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, deepcopy
+from copy import copy
 from functools import partial
 from hashlib import sha1
 from typing import (
@@ -642,14 +642,6 @@ def prepare_single_task(
                 )
                 return
 
-            # 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(
                     f"Ignoring unknown node name {packet.node} in pending sends"
@@ -1130,43 +1122,10 @@ def sanitize_untracked_values_in_send(
         # 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] = UNTRACKED_VALUE_PLACEHOLDER
-        return obj
+    sanitized_arg = dict(packet.arg)
+
+    for k, v in sanitized_arg.items():
+        if isinstance(channels.get(k), UntrackedValue):
+            sanitized_arg[k] = UNTRACKED_VALUE_PLACEHOLDER
 
-    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 UNTRACKED_VALUE_PLACEHOLDER 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 == UNTRACKED_VALUE_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

@@ -329,13 +329,14 @@ class PregelLoop:
         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
+            # 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)
+                (
+                    (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)
@@ -756,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
+        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_pregel.py

@@ -8636,10 +8636,9 @@ def test_send_with_untracked_value(sync_checkpointer: BaseCheckpointSaver):
     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"}}
+    config = {"configurable": {"thread_id": "1"}}
     result = app.invoke({}, config)
 
     assert len(result["messages"]) == 2
@@ -8647,6 +8646,54 @@ def test_send_with_untracked_value(sync_checkpointer: BaseCheckpointSaver):
     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
+
+
+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"}