diff --git a/libs/checkpoint-conformance/tests/test_validate_memory.py b/libs/checkpoint-conformance/tests/test_validate_memory.py index f458c82a2..bf50f73e1 100644 --- a/libs/checkpoint-conformance/tests/test_validate_memory.py +++ b/libs/checkpoint-conformance/tests/test_validate_memory.py @@ -14,8 +14,8 @@ async def memory_checkpointer(): @pytest.mark.asyncio -async def test_validate_memory_base(): - """InMemorySaver passes all base capability tests.""" +async def test_validate_memory(): + """InMemorySaver passes the tests of every capability it implements.""" report = await validate(memory_checkpointer) report.print_report() - assert report.passed_all_base(), f"Base tests failed: {report.to_dict()}" + assert report.passed_all(), f"Capability tests failed: {report.to_dict()}" diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py index 78222205f..5e61b50ed 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py @@ -469,7 +469,7 @@ class PostgresSaver(BasePostgresSaver): thread_id = config["configurable"]["thread_id"] checkpoint_ns = config["configurable"].get("checkpoint_ns", "") checkpoint_id = get_checkpoint_id(config) - if checkpoint_id is None: + if not checkpoint_id: target = self.get_tuple(config) if target is None: return {ch: {"writes": []} for ch in channels} diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py index 2b813b391..5a00c99b1 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py @@ -417,7 +417,7 @@ class AsyncPostgresSaver(BasePostgresSaver): thread_id = config["configurable"]["thread_id"] checkpoint_ns = config["configurable"].get("checkpoint_ns", "") checkpoint_id = get_checkpoint_id(config) - if checkpoint_id is None: + if not checkpoint_id: target = await self.aget_tuple(config) if target is None: return {ch: {"writes": []} for ch in channels} diff --git a/libs/checkpoint-postgres/tests/test_delta_pagination.py b/libs/checkpoint-postgres/tests/test_delta_pagination.py index 0c7407068..f9b483fc3 100644 --- a/libs/checkpoint-postgres/tests/test_delta_pagination.py +++ b/libs/checkpoint-postgres/tests/test_delta_pagination.py @@ -164,6 +164,40 @@ def test_sync_walk_reads_nothing_newer_than_the_target( assert read == _ids_from_target_down(configs) +def _empty_checkpoint_id(config: dict) -> dict: + return {"configurable": {**config["configurable"], "checkpoint_id": ""}} + + +async def test_async_empty_checkpoint_id_reads_from_the_latest_checkpoint() -> None: + async with AsyncPostgresSaver.from_conn_string(DEFAULT_URI) as saver: + await saver.setup() + configs = await _abuild_chain(saver) + latest = await saver.aget_delta_channel_history( + config=configs[-1], channels=[CHANNEL] + ) + + result = await saver.aget_delta_channel_history( + config=_empty_checkpoint_id(configs[-1]), channels=[CHANNEL] + ) + + assert latest[CHANNEL]["writes"] + assert result == latest + + +def test_sync_empty_checkpoint_id_reads_from_the_latest_checkpoint() -> None: + with PostgresSaver.from_conn_string(DEFAULT_URI) as saver: + saver.setup() + configs = _build_chain(saver) + latest = saver.get_delta_channel_history(config=configs[-1], channels=[CHANNEL]) + + result = saver.get_delta_channel_history( + config=_empty_checkpoint_id(configs[-1]), channels=[CHANNEL] + ) + + assert latest[CHANNEL]["writes"] + assert result == latest + + async def test_root_target_has_no_history_and_still_terminates( monkeypatch: pytest.MonkeyPatch, ) -> None: diff --git a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py index d7369921d..ba77deb98 100644 --- a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py +++ b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py @@ -539,7 +539,7 @@ class SqliteSaver(BaseCheckpointSaver[str]): thread_id = str(config["configurable"]["thread_id"]) checkpoint_ns = config["configurable"].get("checkpoint_ns", "") checkpoint_id = get_checkpoint_id(config) - if checkpoint_id is None: + if not checkpoint_id: target = self.get_tuple(config) if target is None: return {ch: {"writes": []} for ch in channels} diff --git a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py index b26e6ac9c..a62a35473 100644 --- a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py +++ b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py @@ -653,7 +653,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver[str]): thread_id = str(config["configurable"]["thread_id"]) checkpoint_ns = config["configurable"].get("checkpoint_ns", "") checkpoint_id = get_checkpoint_id(config) - if checkpoint_id is None: + if not checkpoint_id: target = await self.aget_tuple(config) if target is None: return {ch: {"writes": []} for ch in channels} diff --git a/libs/checkpoint-sqlite/tests/test_delta_parent_walk.py b/libs/checkpoint-sqlite/tests/test_delta_parent_walk.py index 849752a5d..d8fc82698 100644 --- a/libs/checkpoint-sqlite/tests/test_delta_parent_walk.py +++ b/libs/checkpoint-sqlite/tests/test_delta_parent_walk.py @@ -65,6 +65,39 @@ async def test_async_walk_reaches_parent_whatever_the_id_order( assert got[CHANNEL] == EXPECTED +EMPTY_CHECKPOINT_ID: dict[str, Any] = { + "configurable": {**CONFIG["configurable"], "checkpoint_id": ""} +} + + +def test_sync_empty_checkpoint_id_reads_from_the_latest_checkpoint() -> None: + with SqliteSaver.from_conn_string(":memory:") as saver: + root = saver.put(CONFIG, _checkpoint("a-older", {CHANNEL: "seed"}), {}, {}) + saver.put_writes(root, [(CHANNEL, "write-root")], "task") + saver.put(root, _checkpoint("z-newer", {}), {}, {}) + + got = saver.get_delta_channel_history( + config=EMPTY_CHECKPOINT_ID, channels=[CHANNEL] + ) + + assert got[CHANNEL] == EXPECTED + + +async def test_async_empty_checkpoint_id_reads_from_the_latest_checkpoint() -> None: + async with AsyncSqliteSaver.from_conn_string(":memory:") as saver: + root = await saver.aput( + CONFIG, _checkpoint("a-older", {CHANNEL: "seed"}), {}, {} + ) + await saver.aput_writes(root, [(CHANNEL, "write-root")], "task") + await saver.aput(root, _checkpoint("z-newer", {}), {}, {}) + + got = await saver.aget_delta_channel_history( + config=EMPTY_CHECKPOINT_ID, channels=[CHANNEL] + ) + + assert got[CHANNEL] == EXPECTED + + def test_walk_reaches_root_of_long_chain_with_descending_ids() -> None: steps = 40 with SqliteSaver.from_conn_string(":memory:") as saver: diff --git a/libs/checkpoint/tests/test_conformance_delta.py b/libs/checkpoint/tests/test_conformance_delta.py deleted file mode 100644 index 5443601cf..000000000 --- a/libs/checkpoint/tests/test_conformance_delta.py +++ /dev/null @@ -1,37 +0,0 @@ -"""Run delta-channel conformance capabilities against InMemorySaver.""" - -from __future__ import annotations - -import pytest - -conformance = pytest.importorskip( - "langgraph.checkpoint.conformance", - reason="langgraph-checkpoint-conformance not installed", -) - - -@pytest.mark.asyncio -async def test_delta_channel_conformance(): - # Imported inside the test: the module-level importorskip above is what - # makes these safe, so they cannot move to the top of the file. - from langgraph.checkpoint.conformance import validate # noqa: PLC0415 - from langgraph.checkpoint.conformance.initializer import ( # noqa: PLC0415 - checkpointer_test, - ) - - from langgraph.checkpoint.memory import InMemorySaver # noqa: PLC0415 - - @checkpointer_test(name="InMemorySaver") - async def mem_saver(): - yield InMemorySaver() - - report = await validate( - mem_saver, - capabilities={ - "delta_channel_history", - }, - ) - for cap, result in report.results.items(): - if result.passed is False: - details = "\n".join(result.failures or []) - pytest.fail(f"Capability {cap} failed:\n{details}") diff --git a/libs/langgraph/langgraph/pregel/_checkpoint.py b/libs/langgraph/langgraph/pregel/_checkpoint.py index b2cddc19d..1583d5a96 100644 --- a/libs/langgraph/langgraph/pregel/_checkpoint.py +++ b/libs/langgraph/langgraph/pregel/_checkpoint.py @@ -1,7 +1,7 @@ from __future__ import annotations import uuid -from collections.abc import Callable, Iterable, Mapping +from collections.abc import Callable, Iterable, Mapping, Sequence from datetime import datetime, timezone from inspect import signature from typing import Any, Literal, cast @@ -32,6 +32,7 @@ from langgraph._internal._constants import ( ) from langgraph._internal._typing import MISSING from langgraph.channels.base import BaseChannel +from langgraph.channels.binop import _get_overwrite from langgraph.channels.delta import DeltaChannel from langgraph.managed.base import ManagedValueMapping, ManagedValueSpec @@ -152,6 +153,23 @@ def delta_channels_with_pending_writes( } +def delta_channels_overwritten( + specs: Mapping[str, Any], writes: Iterable[tuple[str, Any]] +) -> set[str]: + """Return the names of the DeltaChannels that `writes` set with an `Overwrite`. + + `update_state` saves a full snapshot of these channels in the checkpoint it + creates, like the loop does when a node returns an `Overwrite`. Otherwise, + reading the channel later starts from an older snapshot and replays the + writes the `Overwrite` threw away. + """ + return { + ch + for ch, value in writes + if isinstance(specs.get(ch), DeltaChannel) and _get_overwrite(value)[0] + } + + def checkpoint_superseded( saver: BaseCheckpointSaver, config: RunnableConfig, saved: CheckpointTuple ) -> bool: @@ -256,6 +274,7 @@ def create_checkpoint( get_next_version: GetNextVersion | None = None, channels_to_snapshot: set[str] | None = None, stored_versions: ChannelVersions | None = None, + trigger_to_nodes: Mapping[str, Sequence[str]] | None = None, ) -> Checkpoint: """Build a new Checkpoint from the previous one and live channel state. @@ -314,7 +333,9 @@ def create_checkpoint( id=id or str(uuid6(clock_seq=step)), channel_values=values, channel_versions=channel_versions, - versions_seen=_mark_bumps_seen(checkpoint["versions_seen"], bumped), + versions_seen=_mark_bumps_seen( + checkpoint["versions_seen"], bumped, trigger_to_nodes or {} + ), updated_channels=None if updated_channels is None else sorted(updated_channels), ) @@ -322,19 +343,27 @@ def create_checkpoint( def _mark_bumps_seen( versions_seen: dict[str, ChannelVersions], bumped: Mapping[str, tuple[Any, Any]], + trigger_to_nodes: Mapping[str, Sequence[str]], ) -> dict[str, ChannelVersions]: """Advance whoever had seen a bumped channel's old version to the new one. A bump that only stores a snapshot is not a write. Left unseen, it would - re-fire `interrupt_before` and rerun the channel's subscribers. For each - entry it advances, `SNAPSHOT_BUMPS` keeps the new version and the one the - node really read, so `versions_seen_without_bumps` can put the read back. + re-fire `interrupt_before` and rerun the channel's subscribers. A channel + bumped from no version was never written, so it also goes to the + subscribers that never ran: they have no entry, and would start on the bump. + For each entry it advances, `SNAPSHOT_BUMPS` keeps the new version and the + one the node really read, so `versions_seen_without_bumps` can put the read + back. """ if not bumped: return versions_seen out = dict(versions_seen) + for k, (old, _) in bumped.items(): + if old is None: + for node in trigger_to_nodes.get(k, ()): + out.setdefault(node, {}) marks = dict(versions_seen.get(SNAPSHOT_BUMPS, {})) - for node, seen in versions_seen.items(): + for node, seen in out.items(): if node == SNAPSHOT_BUMPS: continue for k, (old, new) in bumped.items(): diff --git a/libs/langgraph/langgraph/pregel/_loop.py b/libs/langgraph/langgraph/pregel/_loop.py index 2c46108c3..8e17b8741 100644 --- a/libs/langgraph/langgraph/pregel/_loop.py +++ b/libs/langgraph/langgraph/pregel/_loop.py @@ -19,6 +19,7 @@ from typing import ( TypeVar, cast, ) +from uuid import UUID, uuid5 from langchain_core.callbacks import AsyncParentRunManager, ParentRunManager from langchain_core.runnables import RunnableConfig @@ -269,6 +270,9 @@ class PregelLoop: # `_put_exit_delta_writes` uses this to decide between anchoring on # the existing parent (True) or creating a lazy stub (False). _has_persisted_parent: bool = False + # True iff `__enter__` loaded the thread's latest checkpoint, not one a + # `checkpoint_id` addressed, so nothing has been built on it yet. + _loaded_latest: bool = False managed: ManagedValueMapping checkpoint: Checkpoint @@ -1117,16 +1121,26 @@ class PregelLoop: self._exit_delta_writes.append( (self.step, NULL_TASK_ID, "", c, v) ) - # Persist delta-channel input writes so sub-freq inputs are - # recoverable via ancestor walk (mirrors the Command input path). + # A DeltaChannel reads its input from the writes stored on the + # checkpoint this run starts from, under a task id of their own: + # readers apply a checkpoint's NULL_TASK_ID writes as its own state. + # A new thread has no checkpoint to store them on, and one a + # `checkpoint_id` addressed may have children that would read them, + # so then the input checkpoint snapshots the channel instead. if self.durability != "exit": delta_input = [ (c, v) for c, v in input_writes if isinstance(self.specs.get(c), DeltaChannel) ] - if delta_input: - self.put_writes(NULL_TASK_ID, delta_input) + if delta_input and self._has_persisted_parent and self._loaded_latest: + self.put_writes( + str(uuid5(UUID(self.checkpoint["id"]), INPUT)), delta_input + ) + else: + self._delta_channels_forced_snapshot.update( + c for c, _ in delta_input + ) # save input checkpoint self.updated_channels = updated_channels self._put_checkpoint({"source": "input"}) @@ -1256,6 +1270,7 @@ class PregelLoop: else None, channels_to_snapshot=channels_to_snapshot, stored_versions=self.checkpoint_previous_versions, + trigger_to_nodes=self.trigger_to_nodes, ) for k in channels_to_snapshot: new_counters[k] = (0, 0) @@ -1774,6 +1789,7 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager): # Normal case: fetch the most recent checkpoint for this # graph/thread. Returns None on first invocation. saved = self.checkpointer.get_tuple(self.checkpoint_config) + self._loaded_latest = True # Capture before the synthetic-empty fallback below overwrites `saved`. # `_put_exit_delta_writes` uses this on first run (no persisted parent) @@ -2030,6 +2046,7 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager): # Normal case: fetch the most recent checkpoint for this # graph/thread. Returns None on first invocation. saved = await self.checkpointer.aget_tuple(self.checkpoint_config) + self._loaded_latest = True # Capture before the synthetic-empty fallback below overwrites `saved`. # `_put_exit_delta_writes` uses this on first run (no persisted parent) diff --git a/libs/langgraph/langgraph/pregel/main.py b/libs/langgraph/langgraph/pregel/main.py index 1cf1aa4bb..237ecf7be 100644 --- a/libs/langgraph/langgraph/pregel/main.py +++ b/libs/langgraph/langgraph/pregel/main.py @@ -137,6 +137,7 @@ from langgraph.pregel._checkpoint import ( copy_checkpoint, create_checkpoint, create_checkpoint_plan_for_update_state_api, + delta_channels_overwritten, delta_channels_with_pending_writes, empty_checkpoint, get_updated_channels_from_tasks, @@ -1766,6 +1767,7 @@ class Pregel( get_next_version=checkpointer.get_next_version, channels_to_snapshot=channels_to_snapshot, stored_versions=checkpoint_previous_versions, + trigger_to_nodes=self.trigger_to_nodes, ) next_config = checkpointer.put( checkpoint_config, @@ -1795,6 +1797,9 @@ class Pregel( checkpointer.get_next_version, self.trigger_to_nodes, ) + fork_pending |= delta_channels_overwritten( + self.channels, input_writes + ) # apply input write to channels next_step = ( @@ -1822,6 +1827,7 @@ class Pregel( get_next_version=checkpointer.get_next_version, channels_to_snapshot=channels_to_snapshot, stored_versions=checkpoint_previous_versions, + trigger_to_nodes=self.trigger_to_nodes, ) next_config = checkpointer.put( checkpoint_config, @@ -2033,6 +2039,9 @@ class Pregel( ), ) updated_channels = get_updated_channels_from_tasks(run_tasks) + fork_pending |= delta_channels_overwritten( + self.channels, (w for t in run_tasks for w in t.writes) + ) # The base's other children replay whatever is stored on it, so an # edit of an older checkpoint stores none of its writes there: the # checkpoint written here carries them, its delta channels @@ -2086,6 +2095,7 @@ class Pregel( else None, channels_to_snapshot=channels_to_snapshot, stored_versions=checkpoint_previous_versions, + trigger_to_nodes=self.trigger_to_nodes, ) next_config = checkpointer.put( checkpoint_config, @@ -2258,6 +2268,7 @@ class Pregel( get_next_version=checkpointer.get_next_version, channels_to_snapshot=channels_to_snapshot, stored_versions=checkpoint_previous_versions, + trigger_to_nodes=self.trigger_to_nodes, ) next_config = await checkpointer.aput( checkpoint_config, @@ -2287,6 +2298,9 @@ class Pregel( checkpointer.get_next_version, self.trigger_to_nodes, ) + fork_pending |= delta_channels_overwritten( + self.channels, input_writes + ) # apply input write to channels next_step = ( @@ -2314,6 +2328,7 @@ class Pregel( get_next_version=checkpointer.get_next_version, channels_to_snapshot=channels_to_snapshot, stored_versions=checkpoint_previous_versions, + trigger_to_nodes=self.trigger_to_nodes, ) next_config = await checkpointer.aput( checkpoint_config, @@ -2524,6 +2539,9 @@ class Pregel( ), ) updated_channels = get_updated_channels_from_tasks(run_tasks) + fork_pending |= delta_channels_overwritten( + self.channels, (w for t in run_tasks for w in t.writes) + ) # The base's other children replay whatever is stored on it, so an # edit of an older checkpoint stores none of its writes there: the # checkpoint written here carries them, its delta channels @@ -2577,6 +2595,7 @@ class Pregel( else None, channels_to_snapshot=channels_to_snapshot, stored_versions=checkpoint_previous_versions, + trigger_to_nodes=self.trigger_to_nodes, ) next_config = await checkpointer.aput( checkpoint_config, diff --git a/libs/langgraph/tests/test_delta_channel_run_input.py b/libs/langgraph/tests/test_delta_channel_run_input.py new file mode 100644 index 000000000..45af2a51d --- /dev/null +++ b/libs/langgraph/tests/test_delta_channel_run_input.py @@ -0,0 +1,103 @@ +"""A run's input to a DeltaChannel input channel reads back on the checkpoints +built from it, and on no others.""" + +import pytest +from langgraph.checkpoint.memory import InMemorySaver + +from langgraph.channels.binop import BinaryOperatorAggregate +from langgraph.channels.delta import DeltaChannel +from langgraph.channels.last_value import LastValue +from langgraph.pregel import NodeBuilder, Pregel +from langgraph.types import Durability + +pytestmark = pytest.mark.anyio + + +def _sorted_extend(current: list, writes: list) -> list: + return sorted([*current, *(item for write in writes for item in write)]) + + +def _delta_input_graph() -> Pregel: + node = NodeBuilder().subscribe_only("go").do(lambda _: [2]).write_to("log", "plain") + return Pregel( + nodes={"n": node}, + channels={ + "log": DeltaChannel(_sorted_extend), + "plain": BinaryOperatorAggregate(list, lambda a, b: sorted(a + b)), + "go": LastValue(int), + }, + input_channels=["log", "plain", "go"], + output_channels=["log", "plain"], + checkpointer=InMemorySaver(), + ) + + +def test_each_run_input_reads_back_on_its_own_checkpoints( + durability: Durability, +) -> None: + graph = _delta_input_graph() + config = {"configurable": {"thread_id": "t"}} + + graph.invoke({"log": [0], "plain": [0], "go": 1}, config, durability=durability) + graph.invoke({"log": [5], "plain": [5], "go": 1}, config, durability=durability) + + for state in graph.get_state_history(config): + assert state.values.get("log", []) == state.values.get("plain", []) + + +async def test_each_run_input_reads_back_on_its_own_checkpoints_async( + durability: Durability, +) -> None: + graph = _delta_input_graph() + config = {"configurable": {"thread_id": "t"}} + + await graph.ainvoke( + {"log": [0], "plain": [0], "go": 1}, config, durability=durability + ) + await graph.ainvoke( + {"log": [5], "plain": [5], "go": 1}, config, durability=durability + ) + + async for state in graph.aget_state_history(config): + assert state.values.get("log", []) == state.values.get("plain", []) + + +OTHER_BRANCH_INPUTS = pytest.mark.parametrize( + "other_branch_input", + [{"go": 1}, {"log": [5], "plain": [5], "go": 1}], + ids=["other-branch-without-delta-input", "other-branch-with-delta-input"], +) + + +@OTHER_BRANCH_INPUTS +def test_run_input_from_an_older_checkpoint_stays_out_of_its_other_branch( + durability: Durability, other_branch_input: dict +) -> None: + graph = _delta_input_graph() + config = {"configurable": {"thread_id": "t"}} + graph.invoke({"go": 1}, config, durability=durability) + older = graph.get_state(config).config + graph.invoke(other_branch_input, config, durability=durability) + + graph.invoke({"log": [7], "plain": [7], "go": 1}, older, durability=durability) + + for state in graph.get_state_history(config): + assert state.values.get("log", []) == state.values.get("plain", []) + + +@OTHER_BRANCH_INPUTS +async def test_run_input_from_an_older_checkpoint_stays_out_of_its_other_branch_async( + durability: Durability, other_branch_input: dict +) -> None: + graph = _delta_input_graph() + config = {"configurable": {"thread_id": "t"}} + await graph.ainvoke({"go": 1}, config, durability=durability) + older = (await graph.aget_state(config)).config + await graph.ainvoke(other_branch_input, config, durability=durability) + + await graph.ainvoke( + {"log": [7], "plain": [7], "go": 1}, older, durability=durability + ) + + async for state in graph.aget_state_history(config): + assert state.values.get("log", []) == state.values.get("plain", []) diff --git a/libs/langgraph/tests/test_delta_channel_seal_subscribers.py b/libs/langgraph/tests/test_delta_channel_seal_subscribers.py new file mode 100644 index 000000000..8a46e6be8 --- /dev/null +++ b/libs/langgraph/tests/test_delta_channel_seal_subscribers.py @@ -0,0 +1,78 @@ +import pytest +from langgraph.checkpoint.memory import InMemorySaver + +from langgraph.channels.delta import DeltaChannel +from langgraph.channels.last_value import LastValue +from langgraph.pregel import NodeBuilder, Pregel + +pytestmark = pytest.mark.anyio + +CONFIG = {"configurable": {"thread_id": "t"}} + + +def _extend(current: list, writes: list) -> list: + return [*current, *(item for write in writes for item in write)] + + +def _graph(reads: list) -> Pregel: + writer = NodeBuilder().subscribe_only("a").do(lambda _: [1]).write_to("d") + reader = NodeBuilder().subscribe_only("d").do(lambda d: reads.append(list(d))) + return Pregel( + nodes={"writer": writer, "reader": reader}, + channels={"a": LastValue(str), "d": DeltaChannel(_extend)}, + input_channels="a", + output_channels=["d"], + checkpointer=InMemorySaver(), + ) + + +def test_an_input_update_from_before_the_first_write_starts_only_the_writer() -> None: + reads: list = [] + graph = _graph(reads) + graph.invoke("go", CONFIG) + first = next( + s.config for s in graph.get_state_history(CONFIG) if s.metadata["step"] == -1 + ) + + fork = graph.update_state(first, {"a": "go"}, as_node="__input__") + + assert graph.get_state(fork).next == ("writer",) + graph.invoke(None, fork) + assert reads == [[1], [1]] + + +def test_a_replay_from_before_the_first_write_forks_with_only_the_writer_next() -> None: + reads: list = [] + graph = _graph(reads) + graph.invoke("go", CONFIG) + first = next( + s.config for s in graph.get_state_history(CONFIG) if s.metadata["step"] == -1 + ) + + graph.invoke(None, first, durability="sync") + + fork = next( + s for s in graph.get_state_history(CONFIG) if s.metadata["source"] == "fork" + ) + assert fork.next == ("writer",) + graph.invoke(None, fork.config) + assert reads == [[1], [1], [1]] + + +async def test_an_ainput_update_from_before_the_first_write_starts_only_the_writer() -> ( + None +): + reads: list = [] + graph = _graph(reads) + await graph.ainvoke("go", CONFIG) + first = [ + s.config + async for s in graph.aget_state_history(CONFIG) + if s.metadata["step"] == -1 + ][0] + + fork = await graph.aupdate_state(first, {"a": "go"}, as_node="__input__") + + assert (await graph.aget_state(fork)).next == ("writer",) + await graph.ainvoke(None, fork) + assert reads == [[1], [1]] diff --git a/libs/langgraph/tests/test_delta_channel_update_overwrite.py b/libs/langgraph/tests/test_delta_channel_update_overwrite.py new file mode 100644 index 000000000..88687326a --- /dev/null +++ b/libs/langgraph/tests/test_delta_channel_update_overwrite.py @@ -0,0 +1,110 @@ +"""An `Overwrite` through `update_state` snapshots its DeltaChannel on the +checkpoint the update saves, as a node's `Overwrite` does on the loop's.""" + +from typing import Annotated, Any + +import pytest +from langchain_core.messages import HumanMessage +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.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 Overwrite + +pytestmark = pytest.mark.anyio + + +class _State(TypedDict): + messages: Annotated[list, DeltaChannel(_messages_delta_reducer)] + + +def _messages_graph(saver: InMemorySaver) -> Any: + builder = StateGraph(_State) + builder.add_node("model", lambda state: {}) + builder.add_edge(START, "model") + return builder.compile(checkpointer=saver) + + +def _extend(current: list, writes: list) -> list: + return [*current, *(item for write in writes for item in write)] + + +def _delta_input_graph(saver: InMemorySaver) -> Pregel: + node = NodeBuilder().subscribe_only("go").do(lambda _: [2]).write_to("log") + return Pregel( + nodes={"n": node}, + channels={"log": DeltaChannel(_extend), "go": LastValue(int)}, + input_channels=["log", "go"], + output_channels=["log"], + checkpointer=saver, + ) + + +def test_update_state_with_an_overwrite_snapshots_the_channel() -> None: + saver = InMemorySaver() + graph = _messages_graph(saver) + config = {"configurable": {"thread_id": "t"}} + graph.invoke({"messages": [HumanMessage(content="a", id="1")]}, config) + + graph.update_state( + config, + {"messages": Overwrite([HumanMessage(content="b", id="2")])}, + as_node="model", + ) + + head = saver.get_tuple(config) + assert head is not None + assert isinstance(head.checkpoint["channel_values"].get("messages"), _DeltaSnapshot) + assert [m.content for m in graph.get_state(config).values["messages"]] == ["b"] + + +async def test_aupdate_state_with_an_overwrite_snapshots_the_channel() -> None: + saver = InMemorySaver() + graph = _messages_graph(saver) + config = {"configurable": {"thread_id": "t"}} + await graph.ainvoke({"messages": [HumanMessage(content="a", id="1")]}, config) + + await graph.aupdate_state( + config, + {"messages": Overwrite([HumanMessage(content="b", id="2")])}, + as_node="model", + ) + + head = await saver.aget_tuple(config) + assert head is not None + assert isinstance(head.checkpoint["channel_values"].get("messages"), _DeltaSnapshot) + values = (await graph.aget_state(config)).values + assert [m.content for m in values["messages"]] == ["b"] + + +def test_update_state_as_input_with_an_overwrite_snapshots_the_channel() -> None: + saver = InMemorySaver() + graph = _delta_input_graph(saver) + config = {"configurable": {"thread_id": "t"}} + graph.invoke({"log": [0], "go": 1}, config) + + graph.update_state(config, {"log": Overwrite([1])}, as_node="__input__") + + head = saver.get_tuple(config) + assert head is not None + assert isinstance(head.checkpoint["channel_values"].get("log"), _DeltaSnapshot) + assert graph.get_state(config).values["log"] == [1] + + +async def test_aupdate_state_as_input_with_an_overwrite_snapshots_the_channel() -> None: + saver = InMemorySaver() + graph = _delta_input_graph(saver) + config = {"configurable": {"thread_id": "t"}} + await graph.ainvoke({"log": [0], "go": 1}, config) + + await graph.aupdate_state(config, {"log": Overwrite([1])}, as_node="__input__") + + head = await saver.aget_tuple(config) + assert head is not None + assert isinstance(head.checkpoint["channel_values"].get("log"), _DeltaSnapshot) + assert (await graph.aget_state(config)).values["log"] == [1]