diff --git a/libs/langgraph/langgraph/pregel/main.py b/libs/langgraph/langgraph/pregel/main.py index 1cf1aa4bb..92a05864b 100644 --- a/libs/langgraph/langgraph/pregel/main.py +++ b/libs/langgraph/langgraph/pregel/main.py @@ -1788,6 +1788,25 @@ class Pregel( ) if input_writes := deque(map_input(self.input_channels, values)): + # A DeltaChannel reads the writes stored on a checkpoint's + # ancestors, not its own, so they go on the checkpoint this + # update builds on, as a node's writes do. + if ( + is_first + and saved is not None + and checkpoint_superseded(checkpointer, config, saved) + ): + fork_pending.update( + ch + for ch, _ in input_writes + if isinstance(self.channels.get(ch), DeltaChannel) + ) + elif saved is not None: + checkpointer.put_writes( + checkpoint_config, + input_writes, + str(uuid5(UUID(checkpoint["id"]), INPUT)), + ) updated_channels = apply_writes( checkpoint, channels, @@ -1833,13 +1852,6 @@ class Pregel( ), ) - # store the writes - checkpointer.put_writes( - next_config, - input_writes, - str(uuid5(UUID(checkpoint["id"]), INPUT)), - ) - return patch_checkpoint_map( next_config, saved.metadata if saved else None ) @@ -2280,6 +2292,25 @@ class Pregel( ) if input_writes := deque(map_input(self.input_channels, values)): + # A DeltaChannel reads the writes stored on a checkpoint's + # ancestors, not its own, so they go on the checkpoint this + # update builds on, as a node's writes do. + if ( + is_first + and saved is not None + and await acheckpoint_superseded(checkpointer, config, saved) + ): + fork_pending.update( + ch + for ch, _ in input_writes + if isinstance(self.channels.get(ch), DeltaChannel) + ) + elif saved is not None: + await checkpointer.aput_writes( + checkpoint_config, + input_writes, + str(uuid5(UUID(checkpoint["id"]), INPUT)), + ) updated_channels = apply_writes( checkpoint, channels, @@ -2325,13 +2356,6 @@ class Pregel( ), ) - # store the writes - await checkpointer.aput_writes( - next_config, - input_writes, - str(uuid5(UUID(checkpoint["id"]), INPUT)), - ) - return patch_checkpoint_map( next_config, saved.metadata if saved else None ) diff --git a/libs/langgraph/tests/test_delta_channel_update_state.py b/libs/langgraph/tests/test_delta_channel_update_state.py index b8424dfac..692dd3c1d 100644 --- a/libs/langgraph/tests/test_delta_channel_update_state.py +++ b/libs/langgraph/tests/test_delta_channel_update_state.py @@ -25,9 +25,12 @@ from langgraph.checkpoint.memory import InMemorySaver from langgraph.checkpoint.serde.types import _DeltaSnapshot from typing_extensions import TypedDict +from langgraph.channels.binop import BinaryOperatorAggregate from langgraph.channels.delta import DeltaChannel +from langgraph.channels.last_value import LastValue from langgraph.graph import START, StateGraph from langgraph.graph.message import _messages_delta_reducer +from langgraph.pregel import NodeBuilder, Pregel from langgraph.types import StateSnapshot, StateUpdate pytestmark = pytest.mark.anyio @@ -538,3 +541,94 @@ def test_update_state_that_snapshots_keeps_a_deferred_node_pending() -> None: assert [m.content for m in final["messages"]] == ["s", "a", "u", "b"] assert graph.get_state(config).next == () + + +def _sorted_extend(current: list, writes: list) -> list: + return sorted([*current, *(item for write in writes for item in write)]) + + +def _delta_input_graph(snapshot_frequency: int = 1000) -> Any: + node = NodeBuilder().subscribe_only("go").do(lambda _: [2]).write_to("log", "plain") + return Pregel( + nodes={"n": node}, + channels={ + "log": DeltaChannel(_sorted_extend, snapshot_frequency=snapshot_frequency), + "plain": BinaryOperatorAggregate(list, lambda a, b: sorted(a + b)), + "go": LastValue(int), + }, + input_channels=["log", "plain", "go"], + output_channels=["log", "plain"], + checkpointer=InMemorySaver(), + ) + + +@pytest.mark.parametrize("snapshot_frequency", [1, 2]) +def test_update_as_input_reads_back_on_its_checkpoint_and_after_the_next_run( + snapshot_frequency: int, +) -> None: + graph = _delta_input_graph(snapshot_frequency) + config = {"configurable": {"thread_id": "t"}} + graph.invoke({"log": [0], "plain": [0], "go": 1}, config) + + graph.update_state(config, {"log": [1], "plain": [1], "go": 1}, as_node="__input__") + after_update = graph.get_state(config).values + graph.invoke(None, config) + after_run = graph.get_state(config).values + + assert after_update["log"] == after_update["plain"] + assert after_run["log"] == after_run["plain"] + + +@pytest.mark.parametrize("snapshot_frequency", [1, 2]) +async def test_aupdate_as_input_reads_back_on_its_checkpoint_and_after_the_next_run( + snapshot_frequency: int, +) -> None: + graph = _delta_input_graph(snapshot_frequency) + config = {"configurable": {"thread_id": "t"}} + await graph.ainvoke({"log": [0], "plain": [0], "go": 1}, config) + + await graph.aupdate_state( + config, {"log": [1], "plain": [1], "go": 1}, as_node="__input__" + ) + after_update = (await graph.aget_state(config)).values + await graph.ainvoke(None, config) + after_run = (await graph.aget_state(config)).values + + assert after_update["log"] == after_update["plain"] + assert after_run["log"] == after_run["plain"] + + +def test_update_as_input_to_an_older_checkpoint_stays_out_of_its_other_branch() -> None: + graph = _delta_input_graph() + config = {"configurable": {"thread_id": "t"}} + graph.invoke({"go": 1}, config) + older = graph.get_state(config).config + graph.invoke({"go": 1}, config) + other_branch = graph.get_state(config) + + edited = graph.update_state( + older, {"log": [1], "plain": [1], "go": 1}, as_node="__input__" + ) + + values = graph.get_state(edited).values + assert values["log"] == values["plain"] + assert graph.get_state(other_branch.config).values == other_branch.values + + +async def test_aupdate_as_input_to_an_older_checkpoint_stays_out_of_its_other_branch() -> ( + None +): + graph = _delta_input_graph() + config = {"configurable": {"thread_id": "t"}} + await graph.ainvoke({"go": 1}, config) + older = (await graph.aget_state(config)).config + await graph.ainvoke({"go": 1}, config) + other_branch = await graph.aget_state(config) + + edited = await graph.aupdate_state( + older, {"log": [1], "plain": [1], "go": 1}, as_node="__input__" + ) + + values = (await graph.aget_state(edited)).values + assert values["log"] == values["plain"] + assert (await graph.aget_state(other_branch.config)).values == other_branch.values