diff --git a/libs/langgraph/langgraph/pregel/_algo.py b/libs/langgraph/langgraph/pregel/_algo.py index 98f25d4ac..d1d7f32da 100644 --- a/libs/langgraph/langgraph/pregel/_algo.py +++ b/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,24 @@ 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 values belonging to UntrackedValue channels in Send.arg for safe checkpointing. + + Send is often called with state to be passed to the dest node, which may contain + UntrackedValues at the top level. Send is not typed and arg may be a nested dict.""" + + if not isinstance(packet.arg, dict): + # Command + return packet + + # top level keys should be the channel names + 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) diff --git a/libs/langgraph/langgraph/pregel/_loop.py b/libs/langgraph/langgraph/pregel/_loop.py index 97a937ccb..4ff5d2ad9 100644 --- a/libs/langgraph/langgraph/pregel/_loop.py +++ b/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 = ( 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 4096a4a4f..e0b6871ad 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/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 @@ -8600,6 +8601,105 @@ def test_multiple_writes_same_channel_from_same_node( ] +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"} + + @pytest.mark.parametrize("as_json", [False, True]) def test_overwrite_sequential( sync_checkpointer: BaseCheckpointSaver, as_json: bool