fix(langgraph): hydrate subgraph delta channels with the resolved saver

A subgraph has no saver of its own: it uses the parent's, passed through
`CONFIG_KEY_CHECKPOINTER`, or holds `checkpointer=True`. Every state
reader resolved that saver, then called `_prepare_state_snapshot`, which
re-derived it from `self.checkpointer` and got `None`. A `DeltaChannel`
stores nothing in `channel_values`, so without a saver it hydrated empty
while plain channels in the same snapshot read correctly.

`bulk_update_state` had the same expression, and once an update reached
the channel's snapshot cadence it persisted the empty value as a
`_DeltaSnapshot`, losing history on disk.

Pass the caller's saver in as a required keyword argument.

Fixes #8470

Co-authored-by: gururafiki <22777967+gururafiki@users.noreply.github.com>
Co-authored-by: Yuan Gao <119447586+DavidGao520@users.noreply.github.com>

libs/langgraph/langgraph/pregel/main.py

@@ -1146,6 +1146,8 @@ class Pregel(
         self,
         config: RunnableConfig,
         saved: CheckpointTuple | None,
+        *,
+        saver: BaseCheckpointSaver,
         recurse: BaseCheckpointSaver | None = None,
         apply_pending_writes: bool = False,
     ) -> StateSnapshot:
@@ -1169,9 +1171,7 @@ class Pregel(
         channels, managed = channels_from_checkpoint(
             self.channels,
             saved.checkpoint,
-            saver=self.checkpointer
+            saver=saver,
-            if isinstance(self.checkpointer, BaseCheckpointSaver)
-            else None,
             config=saved.config,
         )
         # tasks for this checkpoint
@@ -1186,11 +1186,7 @@ class Pregel(
             stop,
             for_execution=True,
             store=self.store,
-            checkpointer=(
+            checkpointer=saver,
-                self.checkpointer
-                if isinstance(self.checkpointer, BaseCheckpointSaver)
-                else None
-            ),
             manager=None,
         )
         # get the subgraphs
@@ -1269,6 +1265,8 @@ class Pregel(
         self,
         config: RunnableConfig,
         saved: CheckpointTuple | None,
+        *,
+        saver: BaseCheckpointSaver,
         recurse: BaseCheckpointSaver | None = None,
         apply_pending_writes: bool = False,
     ) -> StateSnapshot:
@@ -1292,9 +1290,7 @@ class Pregel(
         channels, managed = await achannels_from_checkpoint(
             self.channels,
             saved.checkpoint,
-            saver=self.checkpointer
+            saver=saver,
-            if isinstance(self.checkpointer, BaseCheckpointSaver)
-            else None,
             config=saved.config,
         )
         # tasks for this checkpoint
@@ -1309,11 +1305,7 @@ class Pregel(
             stop,
             for_execution=True,
             store=self.store,
-            checkpointer=(
+            checkpointer=saver,
-                self.checkpointer
-                if isinstance(self.checkpointer, BaseCheckpointSaver)
-                else None
-            ),
             manager=None,
         )
         # get the subgraphs
@@ -1429,6 +1421,7 @@ class Pregel(
         return self._prepare_state_snapshot(
             config,
             saved,
+            saver=checkpointer,
             recurse=checkpointer if subgraphs else None,
             apply_pending_writes=CONFIG_KEY_CHECKPOINT_ID not in config[CONF],
         )
@@ -1473,6 +1466,7 @@ class Pregel(
         return await self._aprepare_state_snapshot(
             config,
             saved,
+            saver=checkpointer,
             recurse=checkpointer if subgraphs else None,
             apply_pending_writes=CONFIG_KEY_CHECKPOINT_ID not in config[CONF],
         )
@@ -1527,7 +1521,7 @@ class Pregel(
             checkpointer.list(config, before=before, limit=limit, filter=filter)
         ):
             yield self._prepare_state_snapshot(
-                checkpoint_tuple.config, checkpoint_tuple
+                checkpoint_tuple.config, checkpoint_tuple, saver=checkpointer
             )
 
     async def aget_state_history(
@@ -1584,7 +1578,7 @@ class Pregel(
             )
         ]:
             yield await self._aprepare_state_snapshot(
-                checkpoint_tuple.config, checkpoint_tuple
+                checkpoint_tuple.config, checkpoint_tuple, saver=checkpointer
             )
 
     def bulk_update_state(
@@ -1666,10 +1660,7 @@ class Pregel(
             channels, managed = channels_from_checkpoint(
                 self.channels,
                 checkpoint,
-                saver=self.checkpointer
+                saver=checkpointer if saved is not None else None,
-                if saved is not None
-                and isinstance(self.checkpointer, BaseCheckpointSaver)
-                else None,
                 config=saved.config if saved is not None else None,
             )
             values, as_node = updates[0][:2]
@@ -2132,10 +2123,7 @@ class Pregel(
             channels, managed = await achannels_from_checkpoint(
                 self.channels,
                 checkpoint,
-                saver=self.checkpointer
+                saver=checkpointer if saved is not None else None,
-                if saved is not None
-                and isinstance(self.checkpointer, BaseCheckpointSaver)
-                else None,
                 config=saved.config if saved is not None else None,
             )
             values, as_node = updates[0][:2]

libs/langgraph/tests/test_delta_channel_subgraph.py

@@ -0,0 +1,253 @@
+import operator
+from typing import Annotated, Any, Literal
+
+import pytest
+from langgraph.checkpoint.base import BaseCheckpointSaver
+from typing_extensions import TypedDict
+
+from langgraph.channels.delta import DeltaChannel
+from langgraph.graph import END, START, StateGraph
+
+pytestmark = pytest.mark.anyio
+
+
+def _extend(state: list | None, writes: list[Any]) -> list:
+    out = list(state or [])
+    for write in writes:
+        out.extend(write if isinstance(write, list) else [write])
+    return out
+
+
+def _state_schema(snapshot_frequency: int = 1000) -> type:
+    class State(TypedDict, total=False):
+        delta: Annotated[
+            list, DeltaChannel(_extend, snapshot_frequency=snapshot_frequency)
+        ]
+        plain: Annotated[list, operator.add]
+
+    return State
+
+
+def _both(*items: str) -> dict:
+    return {"delta": list(items), "plain": list(items)}
+
+
+def _child_builder(*, snapshot_frequency: int = 1000) -> StateGraph:
+    builder = StateGraph(_state_schema(snapshot_frequency))
+    builder.add_node("a", lambda state: _both("a1"))
+    builder.add_node("b", lambda state: _both("b1", "b2"))
+    builder.add_edge(START, "a")
+    builder.add_edge("a", "b")
+    builder.add_edge("b", END)
+    return builder
+
+
+def _wrap(
+    inner: StateGraph,
+    *,
+    checkpointer: bool | None = None,
+    interrupt_before: list[str] | None = None,
+) -> StateGraph:
+    builder = StateGraph(inner.state_schema)
+    builder.add_node(
+        "child",
+        inner.compile(checkpointer=checkpointer, interrupt_before=interrupt_before),
+    )
+    builder.add_edge(START, "child")
+    builder.add_edge("child", END)
+    return builder
+
+
+def _nested_app(
+    checkpointer: BaseCheckpointSaver,
+    *,
+    depth: int = 1,
+    snapshot_frequency: int = 1000,
+    pause_before_b: bool = False,
+    subgraph_checkpointer: bool | None = None,
+) -> Any:
+    graph = _child_builder(snapshot_frequency=snapshot_frequency)
+    for _ in range(depth):
+        graph = _wrap(
+            graph,
+            checkpointer=subgraph_checkpointer,
+            interrupt_before=["b"] if pause_before_b else None,
+        )
+    return graph.compile(checkpointer=checkpointer)
+
+
+def _scoped(config: dict, namespace: str) -> dict:
+    return {"configurable": {**config["configurable"], "checkpoint_ns": namespace}}
+
+
+def _child_namespace(app: Any, config: dict, *, depth: int = 1) -> str:
+    namespace = ""
+    for level in range(depth):
+        scoped = _scoped(config, namespace) if namespace else config
+        namespace = next(
+            (
+                task.state["configurable"]["checkpoint_ns"]
+                for snapshot in app.get_state_history(scoped)
+                for task in snapshot.tasks
+                if task.name == "child" and isinstance(task.state, dict)
+            ),
+            "",
+        )
+        assert namespace, f"no `child` subgraph task at nesting level {level}"
+    return namespace
+
+
+async def _achild_namespace(app: Any, config: dict) -> str:
+    async for snapshot in app.aget_state_history(config):
+        for task in snapshot.tasks:
+            if task.name == "child" and isinstance(task.state, dict):
+                return task.state["configurable"]["checkpoint_ns"]
+    raise AssertionError("no `child` subgraph task")
+
+
+HISTORY = [_both("a1", "b1", "b2"), _both("a1"), _both(), _both()]
+
+
+def test_subgraph_get_state(sync_checkpointer: BaseCheckpointSaver) -> None:
+    app = _nested_app(sync_checkpointer)
+    config = {"configurable": {"thread_id": "1"}}
+    app.invoke({}, config)
+
+    child = _scoped(config, _child_namespace(app, config))
+
+    assert app.get_state(config).values == _both("a1", "b1", "b2")
+    assert app.get_state(child).values == _both("a1", "b1", "b2")
+
+
+async def test_subgraph_aget_state(async_checkpointer: BaseCheckpointSaver) -> None:
+    app = _nested_app(async_checkpointer)
+    config = {"configurable": {"thread_id": "1"}}
+    await app.ainvoke({}, config)
+
+    child = _scoped(config, await _achild_namespace(app, config))
+
+    assert (await app.aget_state(child)).values == _both("a1", "b1", "b2")
+
+
+def test_subgraph_get_state_history(sync_checkpointer: BaseCheckpointSaver) -> None:
+    app = _nested_app(sync_checkpointer)
+    config = {"configurable": {"thread_id": "1"}}
+    app.invoke({}, config)
+
+    child = _scoped(config, _child_namespace(app, config))
+
+    assert [s.values for s in app.get_state_history(child)] == HISTORY
+
+
+async def test_subgraph_aget_state_history(
+    async_checkpointer: BaseCheckpointSaver,
+) -> None:
+    app = _nested_app(async_checkpointer)
+    config = {"configurable": {"thread_id": "1"}}
+    await app.ainvoke({}, config)
+
+    child = _scoped(config, await _achild_namespace(app, config))
+
+    assert [s.values async for s in app.aget_state_history(child)] == HISTORY
+
+
+def test_doubly_nested_subgraph_get_state(
+    sync_checkpointer: BaseCheckpointSaver,
+) -> None:
+    app = _nested_app(sync_checkpointer, depth=2)
+    config = {"configurable": {"thread_id": "1"}}
+    app.invoke({}, config)
+
+    child = _scoped(config, _child_namespace(app, config, depth=2))
+
+    assert app.get_state(child).values == _both("a1", "b1", "b2")
+
+
+@pytest.mark.parametrize("persistence", ["per-invocation", "per-thread"])
+def test_interrupted_subgraph_task_state(
+    sync_checkpointer: BaseCheckpointSaver,
+    persistence: Literal["per-invocation", "per-thread"],
+) -> None:
+    app = _nested_app(
+        sync_checkpointer,
+        pause_before_b=True,
+        subgraph_checkpointer=True if persistence == "per-thread" else None,
+    )
+    config = {"configurable": {"thread_id": "1"}}
+    app.invoke({}, config)
+
+    (task,) = app.get_state(config, subgraphs=True).tasks
+
+    assert task.state.values == _both("a1")
+
+
+@pytest.mark.parametrize("persistence", ["per-invocation", "per-thread"])
+async def test_interrupted_subgraph_task_state_async(
+    async_checkpointer: BaseCheckpointSaver,
+    persistence: Literal["per-invocation", "per-thread"],
+) -> None:
+    app = _nested_app(
+        async_checkpointer,
+        pause_before_b=True,
+        subgraph_checkpointer=True if persistence == "per-thread" else None,
+    )
+    config = {"configurable": {"thread_id": "1"}}
+    await app.ainvoke({}, config)
+
+    (task,) = (await app.aget_state(config, subgraphs=True)).tasks
+
+    assert task.state.values == _both("a1")
+
+
+def test_subgraph_update_state_keeps_history(
+    sync_checkpointer: BaseCheckpointSaver,
+) -> None:
+    app = _nested_app(sync_checkpointer, snapshot_frequency=2)
+    config = {"configurable": {"thread_id": "1"}}
+    app.invoke({}, config)
+
+    child = _scoped(config, _child_namespace(app, config))
+    app.update_state(child, _both("manual"))
+
+    assert app.get_state(child).values == _both("a1", "b1", "b2", "manual")
+
+
+async def test_subgraph_aupdate_state_keeps_history(
+    async_checkpointer: BaseCheckpointSaver,
+) -> None:
+    app = _nested_app(async_checkpointer, snapshot_frequency=2)
+    config = {"configurable": {"thread_id": "1"}}
+    await app.ainvoke({}, config)
+
+    child = _scoped(config, await _achild_namespace(app, config))
+    await app.aupdate_state(child, _both("manual"))
+
+    assert (await app.aget_state(child)).values == _both("a1", "b1", "b2", "manual")
+
+
+def test_stateless_subgraph_persists_nothing(
+    sync_checkpointer: BaseCheckpointSaver,
+) -> None:
+    app = _nested_app(sync_checkpointer, subgraph_checkpointer=False)
+    config = {"configurable": {"thread_id": "1"}}
+    app.invoke({}, config)
+
+    child_tasks = [
+        task
+        for snapshot in app.get_state_history(config)
+        for task in snapshot.tasks
+        if task.name == "child" and isinstance(task.state, dict)
+    ]
+
+    assert child_tasks == []
+    assert app.get_state(config).values == _both("a1", "b1", "b2")
+
+
+def test_completed_subgraph_exposes_no_task_state(
+    sync_checkpointer: BaseCheckpointSaver,
+) -> None:
+    app = _nested_app(sync_checkpointer)
+    config = {"configurable": {"thread_id": "1"}}
+    app.invoke({}, config)
+
+    assert app.get_state(config, subgraphs=True).tasks == ()