From fa83d64effc2d12bdbc3a10ae4b86bca5617e901 Mon Sep 17 00:00:00 2001 From: Sydney Runkle Date: Tue, 21 Apr 2026 09:34:44 -0400 Subject: [PATCH] feat(channels): DeltaChannel tracks checkpoint_id; emits prev_checkpoint_id in DeltaValue Co-Authored-By: Claude Sonnet 4.6 --- libs/langgraph/langgraph/channels/base.py | 2 +- libs/langgraph/langgraph/channels/delta.py | 22 ++++++++++------- libs/langgraph/tests/test_channels.py | 28 +++++++++++----------- 3 files changed, 28 insertions(+), 24 deletions(-) diff --git a/libs/langgraph/langgraph/channels/base.py b/libs/langgraph/langgraph/channels/base.py index 77ded1f9e..7a442bae0 100644 --- a/libs/langgraph/langgraph/channels/base.py +++ b/libs/langgraph/langgraph/channels/base.py @@ -120,7 +120,7 @@ class BaseChannel(Generic[Value, Update, Checkpoint], ABC): """ return False - def after_checkpoint(self, version: Any) -> None: + def after_checkpoint(self, version: Any, checkpoint_id: str | None = None) -> None: """Called after checkpoint() with the assigned version, and after from_checkpoint() with the current channel version. diff --git a/libs/langgraph/langgraph/channels/delta.py b/libs/langgraph/langgraph/channels/delta.py index abc3de12b..deb54af9e 100644 --- a/libs/langgraph/langgraph/channels/delta.py +++ b/libs/langgraph/langgraph/channels/delta.py @@ -37,6 +37,7 @@ class DeltaChannel(Generic[Value], BaseChannel[list[Value], Value, DeltaValue]): "snapshot_every", "_pending", "_base_version", + "_last_checkpoint_id", "_overwritten", "_steps_since_rehydrate", ) @@ -63,6 +64,7 @@ class DeltaChannel(Generic[Value], BaseChannel[list[Value], Value, DeltaValue]): self.value = [] self._pending: list[Any] = [] self._base_version: str | None = None + self._last_checkpoint_id: str | None = None self._overwritten: bool = False self._steps_since_rehydrate: int = 0 @@ -92,6 +94,7 @@ class DeltaChannel(Generic[Value], BaseChannel[list[Value], Value, DeltaValue]): new.value = self.value[:] new._pending = self._pending[:] new._base_version = self._base_version + new._last_checkpoint_id = self._last_checkpoint_id new._overwritten = self._overwritten new._steps_since_rehydrate = self._steps_since_rehydrate return new @@ -111,10 +114,12 @@ class DeltaChannel(Generic[Value], BaseChannel[list[Value], Value, DeltaValue]): # the right time regardless of how many prior invocations there were. new._steps_since_rehydrate = len(checkpoint.deltas) elif isinstance(checkpoint, DeltaValue): - raise ValueError( - "DeltaChannel received a raw DeltaValue from the checkpoint saver. " - "Your saver does not support incremental channel storage. " - "Use InMemorySaver or PostgresSaver." + # Should never reach here — the pregel layer assembles DeltaValues + # into DeltaChainValue before calling from_checkpoint. + raise AssertionError( + "DeltaChannel.from_checkpoint received a raw DeltaValue. " + "This is a bug in the pregel layer — chain assembly should have " + "occurred before from_checkpoint was called." ) else: # Backwards compat: plain list from old BinaryOperatorAggregate checkpoint. @@ -173,20 +178,19 @@ class DeltaChannel(Generic[Value], BaseChannel[list[Value], Value, DeltaValue]): return list(self.value) return DeltaValue( delta=self._pending[:], - prev_version=None if self._overwritten else self._base_version, + prev_checkpoint_id=None if self._overwritten else self._last_checkpoint_id, ) - def after_checkpoint(self, version: Any) -> None: + def after_checkpoint(self, version: Any, checkpoint_id: str | None = None) -> None: if version != self._base_version: if self._base_version is None: - # First call after from_checkpoint — anchor the base version - # without counting a step (the counter was seeded by from_checkpoint). - pass + pass # First call after from_checkpoint — anchor without counting a step. elif self.snapshot_every is not None: if self._steps_since_rehydrate >= self.snapshot_every: self._steps_since_rehydrate = 0 else: self._steps_since_rehydrate += 1 self._base_version = version + self._last_checkpoint_id = checkpoint_id self._pending = [] self._overwritten = False diff --git a/libs/langgraph/tests/test_channels.py b/libs/langgraph/tests/test_channels.py index 9d34890fa..fa706fbfa 100644 --- a/libs/langgraph/tests/test_channels.py +++ b/libs/langgraph/tests/test_channels.py @@ -134,13 +134,13 @@ def test_delta_channel_basic_two_steps() -> None: d1 = ch.checkpoint() assert isinstance(d1, DeltaValue) assert len(d1.delta) == 1 - assert d1.prev_version is None # first ever step - ch.after_checkpoint("v1") + assert d1.prev_checkpoint_id is None # first ever step + ch.after_checkpoint("v1", checkpoint_id="cid1") # Step 2: another message ch.update([AIMessage(content="hello", id="a1")]) d2 = ch.checkpoint() - assert d2.prev_version == "v1" + assert d2.prev_checkpoint_id == "cid1" assert len(d2.delta) == 1 ch.after_checkpoint("v2") @@ -217,11 +217,11 @@ def test_delta_channel_overwrite_resets_chain() -> None: ch.update([HumanMessage(content="old", id="h1")]) ch.after_checkpoint("v1") - # Overwrite should create a root blob (prev_version=None) + # Overwrite should create a root blob (prev_checkpoint_id=None) ch.update([Overwrite([HumanMessage(content="new", id="h2")])]) d = ch.checkpoint() assert isinstance(d, DeltaValue) - assert d.prev_version is None # chain root + assert d.prev_checkpoint_id is None # chain root assert len(d.delta) == 1 assert d.delta[0].content == "new" @@ -232,10 +232,10 @@ def test_delta_channel_unsupported_saver_raises() -> None: from langgraph.channels.delta import DeltaChannel from langgraph.graph.message import add_messages - # If a saver returns a raw DeltaValue (unsupported), from_checkpoint raises + # If a saver returns a raw DeltaValue, from_checkpoint raises AssertionError spec = DeltaChannel(add_messages) - raw_delta = DeltaValue(delta=[], prev_version=None) - with pytest.raises(ValueError, match="DeltaChannel received a raw DeltaValue"): + raw_delta = DeltaValue(delta=[], prev_checkpoint_id=None) + with pytest.raises(AssertionError, match="DeltaChannel.from_checkpoint received a raw DeltaValue"): spec.from_checkpoint(raw_delta) @@ -256,7 +256,7 @@ def test_delta_channel_remove_message_delta_and_replay() -> None: ch.update([AIMessage(content="hello", id="a1")]) d1 = ch.checkpoint() assert isinstance(d1, DeltaValue) - ch.after_checkpoint("v1") + ch.after_checkpoint("v1", checkpoint_id="cid1") assert ch.get() == [ HumanMessage(content="hi", id="h1"), AIMessage(content="hello", id="a1"), @@ -266,9 +266,9 @@ def test_delta_channel_remove_message_delta_and_replay() -> None: ch.update([RemoveMessage(id="a1")]) d2 = ch.checkpoint() assert isinstance(d2, DeltaValue) - assert d2.prev_version == "v1" + assert d2.prev_checkpoint_id == "cid1" assert any(isinstance(w, RemoveMessage) for w in d2.delta) - ch.after_checkpoint("v2") + ch.after_checkpoint("v2", checkpoint_id="cid2") assert ch.get() == [HumanMessage(content="hi", id="h1")] # Replay the full chain from scratch — must reproduce the post-remove state @@ -293,14 +293,14 @@ def test_delta_channel_update_by_id_delta_and_replay() -> None: ch.update([HumanMessage(content="original", id="h1")]) d1 = ch.checkpoint() assert isinstance(d1, DeltaValue) - ch.after_checkpoint("v1") + ch.after_checkpoint("v1", checkpoint_id="cid1") # Step 2: update the same message by ID ch.update([HumanMessage(content="updated", id="h1")]) d2 = ch.checkpoint() assert isinstance(d2, DeltaValue) - assert d2.prev_version == "v1" - ch.after_checkpoint("v2") + assert d2.prev_checkpoint_id == "cid1" + ch.after_checkpoint("v2", checkpoint_id="cid2") assert ch.get() == [HumanMessage(content="updated", id="h1")] # Replay the full chain — must produce the updated message, not the original