From 9fb0493ac5da55e11ed912f3ccf3c5ce95b0ed67 Mon Sep 17 00:00:00 2001 From: Sydney Runkle Date: Fri, 17 Apr 2026 14:56:42 -0400 Subject: [PATCH] feat(pregel): call after_checkpoint hook when loading and saving channels Co-Authored-By: Claude Sonnet 4.6 --- .../langgraph/langgraph/pregel/_checkpoint.py | 13 ++-- libs/langgraph/langgraph/pregel/_loop.py | 3 + libs/langgraph/tests/test_pregel.py | 76 +++++++++++++++++++ 3 files changed, 85 insertions(+), 7 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/_checkpoint.py b/libs/langgraph/langgraph/pregel/_checkpoint.py index 3d510ac7d..1b327203d 100644 --- a/libs/langgraph/langgraph/pregel/_checkpoint.py +++ b/libs/langgraph/langgraph/pregel/_checkpoint.py @@ -67,13 +67,12 @@ def channels_from_checkpoint( channel_specs[k] = v else: managed_specs[k] = v - return ( - { - k: v.from_checkpoint(checkpoint["channel_values"].get(k, MISSING)) - for k, v in channel_specs.items() - }, - managed_specs, - ) + channels: dict[str, BaseChannel] = {} + for k, v in channel_specs.items(): + ch = v.from_checkpoint(checkpoint["channel_values"].get(k, MISSING)) + ch.after_checkpoint(checkpoint["channel_versions"].get(k)) + channels[k] = ch + return channels, managed_specs def copy_checkpoint(checkpoint: Checkpoint) -> Checkpoint: diff --git a/libs/langgraph/langgraph/pregel/_loop.py b/libs/langgraph/langgraph/pregel/_loop.py index fa26d2d9c..a47deb612 100644 --- a/libs/langgraph/langgraph/pregel/_loop.py +++ b/libs/langgraph/langgraph/pregel/_loop.py @@ -881,6 +881,9 @@ class PregelLoop: id=self.checkpoint["id"] if exiting else None, updated_channels=self.updated_channels, ) + if do_checkpoint and self.channels: + for k, ch in self.channels.items(): + ch.after_checkpoint(self.checkpoint["channel_versions"].get(k)) # 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() diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index e4d49c46d..cc5cb0eb2 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -9400,3 +9400,79 @@ def test_fork_does_not_apply_pending_writes( # Should be: 1 (input) + 20 (forked node_a) + 100 (node_b) = 121 assert result == {"value": 121} + + +async def test_diff_channel_end_to_end_inmemory() -> None: + """Full graph run: DiffChannel accumulates correctly across multiple turns.""" + from langchain_core.messages import AIMessage, HumanMessage + from langgraph.channels.diff import DiffChannel + from langgraph.checkpoint.memory import InMemorySaver + from langgraph.graph import START, StateGraph + from langgraph.graph.message import add_messages + + class State(TypedDict): + messages: Annotated[list, DiffChannel(add_messages)] + + def respond(state: State) -> dict: + n = len(state["messages"]) + return {"messages": [AIMessage(content=f"reply-{n}", id=f"ai-{n}")]} + + builder = StateGraph(State) + builder.add_node("respond", respond) + builder.add_edge(START, "respond") + graph = builder.compile(checkpointer=InMemorySaver()) + + config = {"configurable": {"thread_id": "diff-test-1"}} + + # Turn 1 + graph.invoke({"messages": [HumanMessage(content="hello", id="h1")]}, config) + # Turn 2 + graph.invoke({"messages": [HumanMessage(content="world", id="h2")]}, config) + # Turn 3 + graph.invoke({"messages": [HumanMessage(content="bye", id="h3")]}, config) + + state = graph.get_state(config) + msgs = state.values["messages"] + # 3 human + 3 AI = 6 total + assert len(msgs) == 6, f"expected 6 messages, got {len(msgs)}: {msgs}" + assert msgs[0].content == "hello" + assert msgs[2].content == "world" + assert msgs[4].content == "bye" + + +async def test_diff_channel_time_travel() -> None: + """Time-travel to an earlier checkpoint reconstructs the correct partial history.""" + from langchain_core.messages import AIMessage, HumanMessage + from langgraph.channels.diff import DiffChannel + from langgraph.checkpoint.memory import InMemorySaver + from langgraph.graph import START, StateGraph + from langgraph.graph.message import add_messages + + class State(TypedDict): + messages: Annotated[list, DiffChannel(add_messages)] + + counter = {"n": 0} + + def respond(state: State) -> dict: + counter["n"] += 1 + return {"messages": [AIMessage(content=f"ai-{counter['n']}", id=f"ai-{counter['n']}")]} + + builder = StateGraph(State) + builder.add_node("respond", respond) + builder.add_edge(START, "respond") + saver = InMemorySaver() + graph = builder.compile(checkpointer=saver) + + config = {"configurable": {"thread_id": "diff-time-travel"}} + + # Run 2 turns + graph.invoke({"messages": [HumanMessage(content="h1", id="h1")]}, config) + graph.invoke({"messages": [HumanMessage(content="h2", id="h2")]}, config) + + # Collect checkpoint history + history = list(graph.get_state_history(config)) + # Find the checkpoint after the first complete turn (should have 2 messages: h1 + ai-1) + after_turn1 = next(h for h in history if len(h.values.get("messages", [])) == 2) + + assert len(after_turn1.values["messages"]) == 2 + assert after_turn1.values["messages"][0].content == "h1"