diff --git a/libs/checkpoint/langgraph/checkpoint/base/__init__.py b/libs/checkpoint/langgraph/checkpoint/base/__init__.py index 7b727bc1e..6e4206119 100644 --- a/libs/checkpoint/langgraph/checkpoint/base/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/base/__init__.py @@ -60,20 +60,29 @@ class CheckpointMetadata(TypedDict, total=False): """ run_id: str """The ID of the run that created this checkpoint.""" - delta_updates_since_snapshot: dict[str, int] - """Per-channel update count since the last `_DeltaSnapshot` was written. + counters_since_delta_snapshot: dict[str, tuple[int, int]] + """Per-channel counters since the last `_DeltaSnapshot` was written. !!! warning "Beta" This metadata field backs `DeltaChannel` (beta). The key name and contents may change while the delta-channel design stabilizes. - Maps channel name → number of supersteps that wrote to this channel - since its last snapshot blob. Used by `pregel.create_checkpoint` to - decide when to write the next snapshot (when the count reaches the - channel's `snapshot_frequency`, snapshot fires and the count resets - to 0). Absent on threads that don't use delta channels. Version-format - independent — works for int, float, and string version schemes. + Maps channel name -> `(updates, supersteps)`: + + - index 0 (`updates`): number of supersteps that wrote to this channel + since its last snapshot blob. + - index 1 (`supersteps`): total supersteps elapsed since this channel's + last snapshot, regardless of whether the channel was written. + + A snapshot fires when EITHER `updates >= ch.snapshot_frequency` OR + `supersteps >= DELTA_MAX_SUPERSTEPS_SINCE_SNAPSHOT` (system-wide bound, + default 5000, env `LANGGRAPH_DELTA_MAX_SUPERSTEPS_SINCE_SNAPSHOT`). + The supersteps bound prevents unbounded ancestor walks on threads where + a delta channel exists but is no longer being updated. + + Absent on threads that don't use delta channels. Persisted as a + 2-element list in JSON (no native tuple). """ diff --git a/libs/langgraph/langgraph/_internal/_config.py b/libs/langgraph/langgraph/_internal/_config.py index d6fae9b06..1bb9128da 100644 --- a/libs/langgraph/langgraph/_internal/_config.py +++ b/libs/langgraph/langgraph/_internal/_config.py @@ -29,6 +29,9 @@ from langgraph._internal._constants import ( ) DEFAULT_RECURSION_LIMIT = int(getenv("LANGGRAPH_DEFAULT_RECURSION_LIMIT", "10007")) +DELTA_MAX_SUPERSTEPS_SINCE_SNAPSHOT = int( + getenv("LANGGRAPH_DELTA_MAX_SUPERSTEPS_SINCE_SNAPSHOT", "5000") +) def recast_checkpoint_ns(ns: str) -> str: diff --git a/libs/langgraph/langgraph/channels/delta.py b/libs/langgraph/langgraph/channels/delta.py index 00ecd4212..052ced50d 100644 --- a/libs/langgraph/langgraph/channels/delta.py +++ b/libs/langgraph/langgraph/channels/delta.py @@ -32,7 +32,7 @@ class DeltaChannel(Generic[Value], BaseChannel[Any, Any, Any]): change in future releases. Threads written with `DeltaChannel` today are expected to remain readable, but the surrounding contract (`BaseCheckpointSaver.get_delta_channel_history`, the - `_DeltaSnapshot` blob shape, the `delta_updates_since_snapshot` + `_DeltaSnapshot` blob shape, the `counters_since_delta_snapshot` metadata field) is not yet stable. The reducer receives the current accumulated value and a batch of writes @@ -47,9 +47,12 @@ class DeltaChannel(Generic[Value], BaseChannel[Any, Any, Any]): This lets LangGraph replay checkpointed writes in larger batches than they were originally produced without changing reconstructed state. - Snapshot cadence is driven by per-channel update count. `create_checkpoint` - writes a full `_DeltaSnapshot` blob every `snapshot_frequency` updates to - this channel, bounding replay depth. + Snapshot cadence is driven by two counters: per-channel update count and + total supersteps since last snapshot. `create_checkpoint` writes a full + `_DeltaSnapshot` blob when EITHER the update count reaches + `snapshot_frequency` OR the supersteps count reaches the system-wide + `DELTA_MAX_SUPERSTEPS_SINCE_SNAPSHOT` bound (default 5000), bounding + replay depth even for channels that stop receiving writes. Parameters: reducer: `(state, list[writes]) -> new_state`. Must be deterministic diff --git a/libs/langgraph/langgraph/pregel/_checkpoint.py b/libs/langgraph/langgraph/pregel/_checkpoint.py index fe990ce96..081bc69bf 100644 --- a/libs/langgraph/langgraph/pregel/_checkpoint.py +++ b/libs/langgraph/langgraph/pregel/_checkpoint.py @@ -8,11 +8,11 @@ from langchain_core.runnables import RunnableConfig from langgraph.checkpoint.base import ( BaseCheckpointSaver, Checkpoint, - CheckpointMetadata, ) from langgraph.checkpoint.base.id import uuid6 from langgraph.checkpoint.serde.types import _DeltaSnapshot +from langgraph._internal._config import DELTA_MAX_SUPERSTEPS_SINCE_SNAPSHOT from langgraph._internal._typing import MISSING from langgraph.channels.base import BaseChannel from langgraph.channels.delta import DeltaChannel @@ -36,21 +36,26 @@ def empty_checkpoint() -> Checkpoint: def delta_channels_to_snapshot( channels: Mapping[str, BaseChannel], - counts: Mapping[str, int], + counters_since_delta_snapshot: Mapping[str, tuple[int, int]], ) -> set[str]: """Return the set of DeltaChannel names that should snapshot now. - A channel snapshots when its accumulated update count (since the last - snapshot) reaches or exceeds `snapshot_frequency`. This is a pure + A channel snapshots when EITHER its accumulated update count reaches + `snapshot_frequency` OR the total supersteps since its last snapshot + reaches `DELTA_MAX_SUPERSTEPS_SINCE_SNAPSHOT`. This is a pure predicate — no mutation. """ - return { - name - for name, ch in channels.items() - if isinstance(ch, DeltaChannel) - and ch.is_available() - and counts.get(name, 0) >= ch.snapshot_frequency - } + result: set[str] = set() + for name, ch in channels.items(): + if not isinstance(ch, DeltaChannel) or not ch.is_available(): + continue + updates, supersteps = counters_since_delta_snapshot.get(name, (0, 0)) + if ( + updates >= ch.snapshot_frequency + or supersteps >= DELTA_MAX_SUPERSTEPS_SINCE_SNAPSHOT + ): + result.add(name) + return result def create_checkpoint( @@ -69,7 +74,7 @@ def create_checkpoint( is written into `channel_values[k]`. Other delta channels are omitted from `channel_values` — the ancestor walk reconstructs their state from `checkpoint_writes`. Callers compute the set via - `delta_channels_to_snapshot(channels, counts)`; defaults to empty + `delta_channels_to_snapshot(channels, counters)`; defaults to empty (no snapshots) when not provided. """ ts = datetime.now(timezone.utc).isoformat() @@ -231,17 +236,3 @@ def copy_checkpoint(checkpoint: Checkpoint) -> Checkpoint: versions_seen={k: v.copy() for k, v in checkpoint["versions_seen"].items()}, updated_channels=checkpoint.get("updated_channels", None), ) - - -def read_delta_updates_since_snapshot( - metadata: CheckpointMetadata | None, -) -> dict[str, int]: - """Read the per-channel update counter from checkpoint metadata. - - Returns an empty dict for missing/None metadata; the dict is - `total=False` on `CheckpointMetadata`, so absence means "no prior - delta-channel activity tracked." - """ - if not metadata: - return {} - return dict(metadata.get("delta_updates_since_snapshot", {}) or {}) diff --git a/libs/langgraph/langgraph/pregel/_loop.py b/libs/langgraph/langgraph/pregel/_loop.py index f94062a5a..9949a62d9 100644 --- a/libs/langgraph/langgraph/pregel/_loop.py +++ b/libs/langgraph/langgraph/pregel/_loop.py @@ -978,35 +978,41 @@ class PregelLoop: if exiting and self.checkpoint["id"] == self.checkpoint_id_saved: # checkpoint already saved return - # Per-delta-channel update bookkeeping. + # Per-delta-channel counter bookkeeping. + # + # Each delta channel tracks a (updates, supersteps) tuple: + # - `updates` increments only when the channel is written this step. + # - `supersteps` increments every superstep regardless. # # `_put_checkpoint` is called once per superstep with a fresh # metadata dict (source="input"|"loop"|"fork") — those are the - # intermediate calls that bump the count by +1 for each delta - # channel touched that step. In exit mode, + # intermediate calls that bump counters. In exit mode, # `_suppress_interrupt`(will rename to _on_loop_exit soon) # additionally calls `_put_checkpoint(self.checkpoint_metadata)` AT # EXIT to commit the final checkpoint — this runs *after* the last # intermediate call already counted the last superstep. So the # exit call must NOT bump again or it would double-count the last - # superstep. (Sync/async durability does not call `_put_checkpoint` - # at exit, so the issue only surfaces in exit mode. force_delta_snapshot - # used to mask this latent bug by resetting every count to 0.) + # superstep. if not exiting: - prev_counts = dict( - self.checkpoint_metadata.get("delta_updates_since_snapshot", {}) or {} + prev_counters = dict( + self.checkpoint_metadata.get("counters_since_delta_snapshot") or {} ) - new_counts = dict(prev_counts) - if self.updated_channels: - for ch_name in self.updated_channels: - if isinstance(self.channels.get(ch_name), DeltaChannel): - new_counts[ch_name] = new_counts.get(ch_name, 0) + 1 + new_counters: dict[str, tuple[int, int]] = {} + updated = self.updated_channels or set() + for ch_name, ch in self.channels.items(): + if not isinstance(ch, DeltaChannel): + continue + u, s = prev_counters.get(ch_name, (0, 0)) + s += 1 + if ch_name in updated: + u += 1 + new_counters[ch_name] = (u, s) metadata["step"] = self.step metadata["parents"] = self.config[CONF].get(CONFIG_KEY_CHECKPOINT_MAP, {}) self.checkpoint_metadata = metadata else: - new_counts = dict( - self.checkpoint_metadata.get("delta_updates_since_snapshot", {}) or {} + new_counters = dict( + self.checkpoint_metadata.get("counters_since_delta_snapshot") or {} ) # do checkpoint? do_checkpoint = self._checkpointer_put_after_previous is not None and ( @@ -1014,7 +1020,7 @@ class PregelLoop: ) # create new checkpoint channels_to_snapshot = ( - delta_channels_to_snapshot(self.channels, new_counts) + delta_channels_to_snapshot(self.channels, new_counters) if do_checkpoint else set() ) @@ -1030,11 +1036,12 @@ class PregelLoop: channels_to_snapshot=channels_to_snapshot, ) for k in channels_to_snapshot: - new_counts[k] = 0 - if new_counts: - self.checkpoint_metadata["delta_updates_since_snapshot"] = new_counts - elif "delta_updates_since_snapshot" in self.checkpoint_metadata: - del self.checkpoint_metadata["delta_updates_since_snapshot"] + new_counters[k] = (0, 0) + non_zero = {k: v for k, v in new_counters.items() if v != (0, 0)} + if non_zero: + self.checkpoint_metadata["counters_since_delta_snapshot"] = non_zero + elif "counters_since_delta_snapshot" in self.checkpoint_metadata: + del self.checkpoint_metadata["counters_since_delta_snapshot"] # 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() @@ -1109,8 +1116,10 @@ class PregelLoop: ): return - counts = self.checkpoint_metadata.get("delta_updates_since_snapshot", {}) or {} - channels_to_snapshot = delta_channels_to_snapshot(self.channels, counts) + counters = dict( + self.checkpoint_metadata.get("counters_since_delta_snapshot") or {} + ) + channels_to_snapshot = delta_channels_to_snapshot(self.channels, counters) pending = [ (step, tid, ch, v) diff --git a/libs/langgraph/tests/test_delta_channel_exit_mode.py b/libs/langgraph/tests/test_delta_channel_exit_mode.py index 2ab09b521..ae53dbbf7 100644 --- a/libs/langgraph/tests/test_delta_channel_exit_mode.py +++ b/libs/langgraph/tests/test_delta_channel_exit_mode.py @@ -160,8 +160,8 @@ async def test_exit_resumed_run_sub_freq() -> None: async def test_exit_count_parity_sync_vs_exit() -> None: - """Sync and exit durability produce the same delta_updates_since_snapshot - after an equivalent run.""" + """Sync and exit durability produce the same update count in + counters_since_delta_snapshot after an equivalent run.""" for durability in ("sync", "exit"): saver = InMemorySaver() graph = _build_graph(saver) @@ -175,9 +175,13 @@ async def test_exit_count_parity_sync_vs_exit() -> None: head = saver.get_tuple(config) assert head is not None - counts = head.metadata.get("delta_updates_since_snapshot", {}) - assert counts.get("messages") == 2, ( - f"durability={durability}: expected count=2, got {counts}" + counters = head.metadata.get("counters_since_delta_snapshot", {}) + updates, supersteps = counters.get("messages", (0, 0)) + assert updates == 2, ( + f"durability={durability}: expected updates=2, got {updates}" + ) + assert supersteps >= 2, ( + f"durability={durability}: expected supersteps>=2, got {supersteps}" ) @@ -196,8 +200,9 @@ async def test_exit_snapshot_fires_at_frequency() -> None: ) head = saver.get_tuple(config) assert head is not None - count1 = head.metadata.get("delta_updates_since_snapshot", {}).get("messages", 0) - assert count1 == 2 + counters1 = head.metadata.get("counters_since_delta_snapshot", {}) + updates1 = counters1.get("messages", (0, 0))[0] + assert updates1 == 2 graph.invoke( {"messages": [HumanMessage(content="m2", id="h2")]}, @@ -206,8 +211,9 @@ async def test_exit_snapshot_fires_at_frequency() -> None: ) head = saver.get_tuple(config) assert head is not None - count2 = head.metadata.get("delta_updates_since_snapshot", {}).get("messages", 0) - assert count2 == 0, f"Expected reset to 0 after snapshot, got {count2}" + counters2 = head.metadata.get("counters_since_delta_snapshot", {}) + updates2 = counters2.get("messages", (0, 0))[0] + assert updates2 == 0, f"Expected reset to 0 after snapshot, got {updates2}" assert isinstance(head.checkpoint["channel_values"].get("messages"), _DeltaSnapshot) @@ -284,7 +290,7 @@ async def test_exit_multi_run_replay_chain() -> None: async def test_exit_metadata_round_trip() -> None: """K=5 consecutive exit runs with snapshot_frequency=5. Verify metadata - delta_updates_since_snapshot increments correctly across runs.""" + counters_since_delta_snapshot increments correctly across runs.""" freq = 5 saver = InMemorySaver() graph = _build_graph(saver, freq=freq) @@ -298,15 +304,16 @@ async def test_exit_metadata_round_trip() -> None: ) head = saver.get_tuple(config) assert head is not None - count = head.metadata.get("delta_updates_since_snapshot", {}).get("messages", 0) + counters = head.metadata.get("counters_since_delta_snapshot", {}) + updates = counters.get("messages", (0, 0))[0] cumulative = i * 2 if cumulative >= freq: - assert count == 0 or count == cumulative % freq or count < freq, ( - f"After run {i}: count={count} should have reset or be partial" + assert updates == 0 or updates == cumulative % freq or updates < freq, ( + f"After run {i}: updates={updates} should have reset or be partial" ) else: - assert count == cumulative, ( - f"After run {i}: expected {cumulative}, got {count}" + assert updates == cumulative, ( + f"After run {i}: expected {cumulative}, got {updates}" ) diff --git a/libs/langgraph/tests/test_delta_channel_supersteps_bound.py b/libs/langgraph/tests/test_delta_channel_supersteps_bound.py new file mode 100644 index 000000000..a2c3a2776 --- /dev/null +++ b/libs/langgraph/tests/test_delta_channel_supersteps_bound.py @@ -0,0 +1,195 @@ +"""Tests for the supersteps-since-last-snapshot bound on DeltaChannel. + +Validates that a delta channel which stops receiving writes is still +force-snapshotted after DELTA_MAX_SUPERSTEPS_SINCE_SNAPSHOT supersteps, +preventing unbounded ancestor walks. +""" + +from typing import Annotated, Any +from unittest.mock import patch + +import pytest +from langgraph.checkpoint.memory import InMemorySaver +from langgraph.checkpoint.serde.types import _DeltaSnapshot +from typing_extensions import TypedDict + +from langgraph.channels.delta import DeltaChannel +from langgraph.graph import END, START, StateGraph +from langgraph.pregel._checkpoint import delta_channels_to_snapshot + +pytestmark = pytest.mark.anyio + + +def _simple_reducer(current: list, updates: list) -> list: + """Flatten updates into current list (each update is itself a list).""" + result = list(current) + for u in updates: + if isinstance(u, list): + result.extend(u) + else: + result.append(u) + return result + + +def _build_two_channel_graph( + checkpointer: InMemorySaver, + *, + freq_a: int = 10_000, + freq_b: int = 10_000, + n_loops: int = 1, +) -> Any: + """Graph with two delta channels A and B. + + The node only writes to channel A; B is never written by the node. + `n_loops` controls how many supersteps the graph runs (via chained nodes). + """ + ch_a = DeltaChannel(_simple_reducer, list, snapshot_frequency=freq_a) + ch_b = DeltaChannel(_simple_reducer, list, snapshot_frequency=freq_b) + State = TypedDict( # noqa: UP013 + "State", + {"a": Annotated[list, ch_a], "b": Annotated[list, ch_b]}, + ) # type: ignore[call-overload] + + builder = StateGraph(State) + + for i in range(n_loops): + name = f"step_{i}" + + def node_fn(state: dict, _i: int = i) -> dict: + return {"a": [f"a-val-{_i}"]} + + builder.add_node(name, node_fn) + if i == 0: + builder.add_edge(START, name) + else: + builder.add_edge(f"step_{i - 1}", name) + if i == n_loops - 1: + builder.add_edge(name, END) + + return builder.compile(checkpointer=checkpointer) + + +async def test_forced_snapshot_single_run() -> None: + """A single invoke with enough supersteps triggers snapshot on the + unwritten channel B via the supersteps bound.""" + max_ss = 3 + with patch( + "langgraph.pregel._checkpoint.DELTA_MAX_SUPERSTEPS_SINCE_SNAPSHOT", + max_ss, + ): + saver = InMemorySaver() + graph = _build_two_channel_graph(saver, n_loops=4) + config = {"configurable": {"thread_id": "single-run-ss"}} + + graph.invoke({"a": ["seed-a"], "b": ["seed-b"]}, config) + + head = saver.get_tuple(config) + assert head is not None + assert isinstance(head.checkpoint["channel_values"].get("b"), _DeltaSnapshot), ( + "Channel B should have been force-snapshotted via supersteps bound" + ) + + state = graph.get_state(config) + assert state.values["b"] == ["seed-b"] + assert "seed-a" in state.values["a"] + + +async def test_forced_snapshot_accumulates_across_runs() -> None: + """Supersteps counter for an unwritten channel persists across separate + invoke() calls. After enough runs, the channel is force-snapshotted.""" + max_ss = 5 + with patch( + "langgraph.pregel._checkpoint.DELTA_MAX_SUPERSTEPS_SINCE_SNAPSHOT", + max_ss, + ): + saver = InMemorySaver() + graph = _build_two_channel_graph(saver, n_loops=1) + config = {"configurable": {"thread_id": "multi-run-ss"}} + + graph.invoke({"a": ["init-a"], "b": ["init-b"]}, config) + + for i in range(1, 6): + graph.invoke({"a": [f"run-{i}"]}, config) + + head = saver.get_tuple(config) + assert head is not None + counters = head.metadata.get("counters_since_delta_snapshot", {}) + b_counters = counters.get("b", (0, 0)) + + if b_counters == (0, 0): + assert isinstance( + head.checkpoint["channel_values"].get("b"), _DeltaSnapshot + ), f"Run {i}: counter reset but no snapshot blob for B" + break + else: + pytest.fail("Channel B was never force-snapshotted after multiple runs") + + state = graph.get_state(config) + assert state.values["b"] == ["init-b"] + assert "init-a" in state.values["a"] + + +async def test_predicate_fires_on_supersteps_overflow() -> None: + """Unit test: delta_channels_to_snapshot fires when supersteps >= MAX + even when updates == 0.""" + ch = DeltaChannel(_simple_reducer, list, snapshot_frequency=10_000) + ch.key = "x" + ch_instance = ch.from_checkpoint(None) + + channels = {"x": ch_instance} + counters: dict[str, tuple[int, int]] = {"x": (0, 5000)} + + result = delta_channels_to_snapshot(channels, counters) + assert "x" in result + + counters_below: dict[str, tuple[int, int]] = {"x": (0, 4999)} + result2 = delta_channels_to_snapshot(channels, counters_below) + assert "x" not in result2 + + +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 + supersteps so we can see the counter in an intermediate state.""" + max_ss = 15 + with patch( + "langgraph.pregel._checkpoint.DELTA_MAX_SUPERSTEPS_SINCE_SNAPSHOT", + max_ss, + ): + saver = InMemorySaver() + graph = _build_two_channel_graph(saver, n_loops=4) + config = {"configurable": {"thread_id": "counter-reset"}} + + graph.invoke({"a": ["seed-a"], "b": ["seed-b"]}, config) + + head = saver.get_tuple(config) + assert head is not None + counters = head.metadata.get("counters_since_delta_snapshot", {}) + b_counters = counters.get("b", (0, 0)) + run1_supersteps = b_counters[1] + assert run1_supersteps > 0, "Should have some supersteps" + assert b_counters[0] == 1, "B written once (input step)" + + graph.invoke({"a": ["more-a"]}, config) + head2 = saver.get_tuple(config) + assert head2 is not None + counters2 = head2.metadata.get("counters_since_delta_snapshot", {}) + b_counters2 = counters2.get("b", (0, 0)) + run2_supersteps = b_counters2[1] + assert run2_supersteps > run1_supersteps, "Supersteps should accumulate" + assert b_counters2[0] == 1, "B written once total (only original input)" + + graph.invoke({"a": ["even-more"]}, config) + head3 = saver.get_tuple(config) + assert head3 is not None + assert isinstance( + head3.checkpoint["channel_values"].get("b"), _DeltaSnapshot + ), "B should have snapshotted at supersteps >= max_ss" + counters3 = head3.metadata.get("counters_since_delta_snapshot", {}) + b_counters3 = counters3.get("b", (0, 0)) + assert b_counters3[1] < max_ss, ( + f"After snapshot, supersteps should have reset, got {b_counters3}" + ) + + state = graph.get_state(config) + assert state.values["b"] == ["seed-b"]