From 06b50627a6ebb18c10674dd32bc3d14d968ab470 Mon Sep 17 00:00:00 2001 From: Elior Nataf Lackritz Date: Fri, 9 Oct 2026 11:31:58 -0400 Subject: [PATCH] fix(langgraph): read a run's DeltaChannel input back only on its own checkpoints Under sync and async durability, a raw Pregel run saved its input to a DeltaChannel input channel as a NULL_TASK_ID write on the checkpoint it started from, which readers apply as that checkpoint's own state. A new thread lost its first input, a run's last checkpoint read the next run's input, and a run from an older checkpoint leaked into the branch that grew from it. The input is now stored there under its own task id, or, on a new thread or an addressed checkpoint, the input checkpoint snapshots the channel. --- libs/langgraph/langgraph/pregel/_loop.py | 24 +++- .../tests/test_delta_channel_run_input.py | 103 ++++++++++++++++++ 2 files changed, 123 insertions(+), 4 deletions(-) create mode 100644 libs/langgraph/tests/test_delta_channel_run_input.py diff --git a/libs/langgraph/langgraph/pregel/_loop.py b/libs/langgraph/langgraph/pregel/_loop.py index 2c46108c3..7aaec19ef 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"}) @@ -1774,6 +1788,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 +2045,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/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", [])