From 0e8252fd3152e1ba46f8e2bfafa00c7c5d492cad Mon Sep 17 00:00:00 2001 From: Caspar Broekhuizen Date: Mon, 27 Oct 2025 18:27:57 -0700 Subject: [PATCH] refactor(langgraph): clean up code, remove rehydration logic, add test --- libs/langgraph/langgraph/pregel/_algo.py | 53 +++--------------------- libs/langgraph/langgraph/pregel/_loop.py | 22 +++++++--- libs/langgraph/tests/test_pregel.py | 53 ++++++++++++++++++++++-- 3 files changed, 73 insertions(+), 55 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/_algo.py b/libs/langgraph/langgraph/pregel/_algo.py index e54cd8ac8..644c33f66 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, 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) diff --git a/libs/langgraph/langgraph/pregel/_loop.py b/libs/langgraph/langgraph/pregel/_loop.py index 27ed2a48b..953e893f8 100644 --- a/libs/langgraph/langgraph/pregel/_loop.py +++ b/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 = ( diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 3b027af17..14ae56323 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/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"}