From 1b84cb08733a2a0a102ed19845bf20511dcc6535 Mon Sep 17 00:00:00 2001 From: Elior Nataf Lackritz Date: Mon, 5 Oct 2026 13:04:30 -0400 Subject: [PATCH] fix(langgraph): keep DeltaChannel counters on every update_state path (#9142) `update_state` saves `counters_since_delta_snapshot` on the checkpoint it creates, except on three paths: `update_state(None, as_node=END)`, `as_node="__input__"` and `as_node="__copy__"`. Those save their checkpoint without it, so the next checkpoint restarts every delta channel's counters from zero: ``` before update_state(None, END): [6, 9] after it: (none) after the next run: [2, 3] (with this change: [6, 10], then [8, 13]) ``` The counters decide when a channel snapshots, so after any of these updates the next snapshot comes later than `snapshot_frequency` and reads walk further back. A thread driven only by these updates never snapshots, so its walk keeps growing. The field's docstring says it is absent only on threads that don't use delta channels. ## Fix `END` and `__input__` build their checkpoint with `create_checkpoint_plan_for_update_state_api`, the plan an update as a node already uses, which now takes the checkpoint's `source`. The counters advance by one superstep (`END` counts the channels its absorbed writes touched as updated), and a channel that reaches its bound snapshots right there. `__copy__` saves a checkpoint with the same values on the same parent, so it keeps the copied checkpoint's counters. ## Tests `test_delta_channel_supersteps_bound.py`, sync and async: - each of the three paths saves the base's counters advanced by one superstep (or unchanged for `__copy__`). - `END` and `__input__` snapshot a channel that reaches its supersteps bound. All 10 fail on main. --- .../langgraph/langgraph/pregel/_checkpoint.py | 5 +- libs/langgraph/langgraph/pregel/main.py | 116 +++++++++++++----- .../tests/test_delta_channel_fork.py | 2 +- .../test_delta_channel_supersteps_bound.py | 109 ++++++++++++++++ 4 files changed, 195 insertions(+), 37 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/_checkpoint.py b/libs/langgraph/langgraph/pregel/_checkpoint.py index 45e581157..c4e98d09b 100644 --- a/libs/langgraph/langgraph/pregel/_checkpoint.py +++ b/libs/langgraph/langgraph/pregel/_checkpoint.py @@ -3,7 +3,7 @@ from __future__ import annotations import uuid from collections.abc import Callable, Iterable, Mapping from datetime import datetime, timezone -from typing import Any, cast +from typing import Any, Literal, cast from langchain_core.runnables import RunnableConfig from langgraph.checkpoint.base import ( @@ -145,6 +145,7 @@ def create_checkpoint_plan_for_update_state_api( channels: Mapping[str, BaseChannel], updated_channels: set[str], *, + source: Literal["update", "input"], step: int, parents: dict[str, Any], saved_metadata: Mapping[str, Any] | None, @@ -154,7 +155,7 @@ def create_checkpoint_plan_for_update_state_api( ) -> tuple[set[str], dict[str, Any]]: """Return ``(channels_to_snapshot, metadata)`` for an update_state head.""" metadata: dict[str, Any] = { - "source": "update", + "source": source, "step": step, "parents": parents, } diff --git a/libs/langgraph/langgraph/pregel/main.py b/libs/langgraph/langgraph/pregel/main.py index ffb828cbc..6ee4f0625 100644 --- a/libs/langgraph/langgraph/pregel/main.py +++ b/libs/langgraph/langgraph/pregel/main.py @@ -1696,6 +1696,7 @@ class Pregel( "Cannot apply multiple updates when clearing state" ) + updated_channels: set[str] = set() if saved is not None: # tasks for this checkpoint next_tasks = prepare_next_tasks( @@ -1718,7 +1719,7 @@ class Pregel( for w in saved.pending_writes or [] if w[0] == NULL_TASK_ID ]: - apply_writes( + updated_channels |= apply_writes( checkpoint, channels, [PregelTaskWrites((), INPUT, null_writes, [])], @@ -1732,7 +1733,7 @@ class Pregel( if tid in next_tasks: next_tasks[tid].writes.extend(status.output) # clear all current tasks - apply_writes( + updated_channels |= apply_writes( checkpoint, channels, next_tasks.values(), @@ -1740,22 +1741,31 @@ class Pregel( self.trigger_to_nodes, ) # save checkpoint + channels_to_snapshot, checkpoint_metadata = ( + create_checkpoint_plan_for_update_state_api( + channels, + updated_channels, + source="update", + step=step + 1, + parents=saved.metadata.get("parents", {}) if saved else {}, + saved_metadata=saved.metadata if saved else None, + is_fresh_thread=saved is None, + fork_channels=fork_pending, + channel_versions=checkpoint["channel_versions"], + ) + ) next_checkpoint = create_checkpoint( checkpoint, channels, step, get_next_version=checkpointer.get_next_version, - channels_to_snapshot=fork_pending, + channels_to_snapshot=channels_to_snapshot, stored_versions=checkpoint_previous_versions, ) next_config = checkpointer.put( checkpoint_config, next_checkpoint, - { - "source": "update", - "step": step + 1, - "parents": saved.metadata.get("parents", {}) if saved else {}, - }, + checkpoint_metadata, get_new_channel_versions( checkpoint_previous_versions, next_checkpoint["channel_versions"], @@ -1773,7 +1783,7 @@ class Pregel( ) if input_writes := deque(map_input(self.input_channels, values)): - apply_writes( + updated_channels = apply_writes( checkpoint, channels, [PregelTaskWrites((), INPUT, input_writes, [])], @@ -1787,24 +1797,31 @@ class Pregel( if saved and saved.metadata.get("step") is not None else -1 ) + channels_to_snapshot, checkpoint_metadata = ( + create_checkpoint_plan_for_update_state_api( + channels, + updated_channels, + source="input", + step=next_step, + parents=saved.metadata.get("parents", {}) if saved else {}, + saved_metadata=saved.metadata if saved else None, + is_fresh_thread=saved is None, + fork_channels=fork_pending, + channel_versions=checkpoint["channel_versions"], + ) + ) next_checkpoint = create_checkpoint( checkpoint, channels, next_step, get_next_version=checkpointer.get_next_version, - channels_to_snapshot=fork_pending, + channels_to_snapshot=channels_to_snapshot, stored_versions=checkpoint_previous_versions, ) next_config = checkpointer.put( checkpoint_config, next_checkpoint, - { - "source": "input", - "step": next_step, - "parents": saved.metadata.get("parents", {}) - if saved - else {}, - }, + checkpoint_metadata, get_new_channel_versions( checkpoint_previous_versions, next_checkpoint["channel_versions"], @@ -1849,6 +1866,12 @@ class Pregel( "source": "fork", "step": step + 1, "parents": saved.metadata.get("parents", {}), + # The copy has the same values and the same parent. + **{ + k: v + for k, v in saved.metadata.items() + if k == "counters_since_delta_snapshot" + }, }, {}, ) @@ -2023,6 +2046,7 @@ class Pregel( create_checkpoint_plan_for_update_state_api( channels, updated_channels, + source="update", step=step + 1, parents=saved.metadata.get("parents", {}) if saved else {}, saved_metadata=saved.metadata if saved else None, @@ -2147,6 +2171,7 @@ class Pregel( raise InvalidUpdateError( "Cannot apply multiple updates when clearing state" ) + updated_channels: set[str] = set() if saved is not None: # tasks for this checkpoint next_tasks = prepare_next_tasks( @@ -2169,7 +2194,7 @@ class Pregel( for w in saved.pending_writes or [] if w[0] == NULL_TASK_ID ]: - apply_writes( + updated_channels |= apply_writes( checkpoint, channels, [PregelTaskWrites((), INPUT, null_writes, [])], @@ -2183,7 +2208,7 @@ class Pregel( if tid in next_tasks: next_tasks[tid].writes.extend(status.output) # clear all current tasks - apply_writes( + updated_channels |= apply_writes( checkpoint, channels, next_tasks.values(), @@ -2191,22 +2216,31 @@ class Pregel( self.trigger_to_nodes, ) # save checkpoint + channels_to_snapshot, checkpoint_metadata = ( + create_checkpoint_plan_for_update_state_api( + channels, + updated_channels, + source="update", + step=step + 1, + parents=saved.metadata.get("parents", {}) if saved else {}, + saved_metadata=saved.metadata if saved else None, + is_fresh_thread=saved is None, + fork_channels=fork_pending, + channel_versions=checkpoint["channel_versions"], + ) + ) next_checkpoint = create_checkpoint( checkpoint, channels, step, get_next_version=checkpointer.get_next_version, - channels_to_snapshot=fork_pending, + channels_to_snapshot=channels_to_snapshot, stored_versions=checkpoint_previous_versions, ) next_config = await checkpointer.aput( checkpoint_config, next_checkpoint, - { - "source": "update", - "step": step + 1, - "parents": saved.metadata.get("parents", {}) if saved else {}, - }, + checkpoint_metadata, get_new_channel_versions( checkpoint_previous_versions, next_checkpoint["channel_versions"], @@ -2224,7 +2258,7 @@ class Pregel( ) if input_writes := deque(map_input(self.input_channels, values)): - apply_writes( + updated_channels = apply_writes( checkpoint, channels, [PregelTaskWrites((), INPUT, input_writes, [])], @@ -2238,24 +2272,31 @@ class Pregel( if saved and saved.metadata.get("step") is not None else -1 ) + channels_to_snapshot, checkpoint_metadata = ( + create_checkpoint_plan_for_update_state_api( + channels, + updated_channels, + source="input", + step=next_step, + parents=saved.metadata.get("parents", {}) if saved else {}, + saved_metadata=saved.metadata if saved else None, + is_fresh_thread=saved is None, + fork_channels=fork_pending, + channel_versions=checkpoint["channel_versions"], + ) + ) next_checkpoint = create_checkpoint( checkpoint, channels, next_step, get_next_version=checkpointer.get_next_version, - channels_to_snapshot=fork_pending, + channels_to_snapshot=channels_to_snapshot, stored_versions=checkpoint_previous_versions, ) next_config = await checkpointer.aput( checkpoint_config, next_checkpoint, - { - "source": "input", - "step": next_step, - "parents": saved.metadata.get("parents", {}) - if saved - else {}, - }, + checkpoint_metadata, get_new_channel_versions( checkpoint_previous_versions, next_checkpoint["channel_versions"], @@ -2300,6 +2341,12 @@ class Pregel( "source": "fork", "step": step + 1, "parents": saved.metadata.get("parents", {}), + # The copy has the same values and the same parent. + **{ + k: v + for k, v in saved.metadata.items() + if k == "counters_since_delta_snapshot" + }, }, {}, ) @@ -2473,6 +2520,7 @@ class Pregel( create_checkpoint_plan_for_update_state_api( channels, updated_channels, + source="update", step=step + 1, parents=saved.metadata.get("parents", {}) if saved else {}, saved_metadata=saved.metadata if saved else None, diff --git a/libs/langgraph/tests/test_delta_channel_fork.py b/libs/langgraph/tests/test_delta_channel_fork.py index 5736dc22e..6d7ae1415 100644 --- a/libs/langgraph/tests/test_delta_channel_fork.py +++ b/libs/langgraph/tests/test_delta_channel_fork.py @@ -315,7 +315,7 @@ def test_fork_before_first_value_by_bulk_update( @pytest.mark.parametrize("first_as_node", [INPUT, END, "__copy__"]) -def test_fork_by_bulk_update_whose_first_superstep_skips_the_plan( +def test_fork_by_bulk_update_whose_first_superstep_is_not_a_node( sync_checkpointer: BaseCheckpointSaver, first_as_node: str ) -> None: config = _thread("t") diff --git a/libs/langgraph/tests/test_delta_channel_supersteps_bound.py b/libs/langgraph/tests/test_delta_channel_supersteps_bound.py index a73dfc51d..b6281eb60 100644 --- a/libs/langgraph/tests/test_delta_channel_supersteps_bound.py +++ b/libs/langgraph/tests/test_delta_channel_supersteps_bound.py @@ -174,6 +174,115 @@ async def test_predicate_fires_on_supersteps_overflow() -> None: ) +def _delta_counters(saver: InMemorySaver, config: Any) -> dict[str, list[int]]: + tup = saver.get_tuple(config) + assert tup is not None + counters = tup.metadata.get("counters_since_delta_snapshot") or {} + return {ch: list(c) for ch, c in counters.items()} + + +_CLEAR_COPY_AND_INPUT_UPDATES = pytest.mark.parametrize( + ("values", "as_node", "supersteps"), + [ + (None, END, 1), + (None, "__copy__", 0), + ({"a": []}, "__input__", 1), + ], + ids=["clear as END", "copy", "update as input"], +) + + +@_CLEAR_COPY_AND_INPUT_UPDATES +def test_update_state_path_keeps_delta_counters( + values: Any, as_node: str, supersteps: int +) -> None: + saver = InMemorySaver() + graph = _build_two_channel_graph(saver) + config = {"configurable": {"thread_id": "counters"}} + graph.invoke({"a": ["seed-a"], "b": ["seed-b"]}, config) + before = _delta_counters(saver, config) + assert set(before) == {"a", "b"}, f"both channels need live counters: {before}" + + updated = graph.update_state(config, values, as_node=as_node) + + assert _delta_counters(saver, updated) == { + ch: [u, s + supersteps] for ch, (u, s) in before.items() + } + + +@_CLEAR_COPY_AND_INPUT_UPDATES +async def test_aupdate_state_path_keeps_delta_counters( + values: Any, as_node: str, supersteps: int +) -> None: + saver = InMemorySaver() + graph = _build_two_channel_graph(saver) + config = {"configurable": {"thread_id": "counters"}} + await graph.ainvoke({"a": ["seed-a"], "b": ["seed-b"]}, config) + before = _delta_counters(saver, config) + assert set(before) == {"a", "b"}, f"both channels need live counters: {before}" + + updated = await graph.aupdate_state(config, values, as_node=as_node) + + assert _delta_counters(saver, updated) == { + ch: [u, s + supersteps] for ch, (u, s) in before.items() + } + + +_UPDATES_THAT_ADD_A_SUPERSTEP = pytest.mark.parametrize( + ("values", "as_node"), + [(None, END), ({"a": []}, "__input__")], + ids=["clear as END", "update as input"], +) + + +@_UPDATES_THAT_ADD_A_SUPERSTEP +def test_update_state_path_snapshots_at_the_supersteps_bound( + values: Any, as_node: str +) -> None: + saver = InMemorySaver() + graph = _build_two_channel_graph(saver) + config = {"configurable": {"thread_id": "bound"}} + graph.invoke({"a": ["seed-a"], "b": ["seed-b"]}, config) + expected = graph.get_state(config).values + supersteps = _delta_counters(saver, config)["b"][1] + + with patch( + "langgraph.pregel._checkpoint.DELTA_MAX_SUPERSTEPS_SINCE_SNAPSHOT", + supersteps + 1, + ): + updated = graph.update_state(config, values, as_node=as_node) + + head = saver.get_tuple(updated) + assert head is not None + assert isinstance(head.checkpoint["channel_values"].get("b"), _DeltaSnapshot) + assert "b" not in _delta_counters(saver, updated) + assert graph.get_state(updated).values == expected + + +@_UPDATES_THAT_ADD_A_SUPERSTEP +async def test_aupdate_state_path_snapshots_at_the_supersteps_bound( + values: Any, as_node: str +) -> None: + saver = InMemorySaver() + graph = _build_two_channel_graph(saver) + config = {"configurable": {"thread_id": "bound"}} + await graph.ainvoke({"a": ["seed-a"], "b": ["seed-b"]}, config) + expected = (await graph.aget_state(config)).values + supersteps = _delta_counters(saver, config)["b"][1] + + with patch( + "langgraph.pregel._checkpoint.DELTA_MAX_SUPERSTEPS_SINCE_SNAPSHOT", + supersteps + 1, + ): + updated = await graph.aupdate_state(config, values, as_node=as_node) + + head = saver.get_tuple(updated) + assert head is not None + assert isinstance(head.checkpoint["channel_values"].get("b"), _DeltaSnapshot) + assert "b" not in _delta_counters(saver, updated) + assert (await graph.aget_state(updated)).values == expected + + async def test_counter_reset_after_supersteps_snapshot() -> None: """After the supersteps bound triggers a snapshot, the counters for that channel reset. Verify by using a bound higher than one run's