WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-10-20 06:51:41 +02:00
Code Issues Packages Projects Releases Wiki Activity

fix(langgraph): keep DeltaChannel counters on every update_state path (#9142)

Browse Source 
`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.
This commit is contained in:
Elior Nataf Lackritz
2026-10-05 13:04:30 -04:00
committed by GitHub
parent 8e15e0d0bc
commit 1b84cb0873
4 changed files with 195 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/langgraph/pregel/_checkpoint.py
+3 -2
View File
@@ -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,
}
libs/langgraph/langgraph/pregel/main.py
+82 -34
View File
@@ -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,
libs/langgraph/tests/test_delta_channel_fork.py
+1 -1
View File
@@ -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")
libs/langgraph/tests/test_delta_channel_supersteps_bound.py
+109
View File
@@ -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