fun

libs/checkpoint/langgraph/checkpoint/base/__init__.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import asyncio
 import copy
 import dataclasses
 import logging
@@ -567,22 +568,34 @@ class BaseCheckpointSaver(Generic[V]):
         thread_id = str(config["configurable"]["thread_id"])
         checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
         current_checkpoint_id = checkpoint.get("id")
-        assembled: dict[str, Any] = {}
 
-        for spec in plan.channels:
-            value = checkpoint["channel_values"].get(spec.name)
-            if spec.kind != "delta" or not isinstance(value, DeltaValue):
-                continue
+        targets = [
+            (spec, checkpoint["channel_values"][spec.name])
+            for spec in plan.channels
+            if spec.kind == "delta"
+            and isinstance(checkpoint["channel_values"].get(spec.name), DeltaValue)
+        ]
+        if not targets:
+            return checkpoint
 
-            assembled_value = await self._amaterialize_delta_channel(
-                thread_id=thread_id,
-                checkpoint_ns=checkpoint_ns,
-                current_checkpoint_id=current_checkpoint_id,
-                channel=spec.name,
-                value=value,
+        # Walks for independent channels can run concurrently — each has its own chain.
+        results = await asyncio.gather(
+            *(
+                self._amaterialize_delta_channel(
+                    thread_id=thread_id,
+                    checkpoint_ns=checkpoint_ns,
+                    current_checkpoint_id=current_checkpoint_id,
+                    channel=spec.name,
+                    value=value,
+                )
+                for spec, value in targets
             )
-            if assembled_value is not None:
-                assembled[spec.name] = assembled_value
+        )
+        assembled = {
+            spec.name: result
+            for (spec, _), result in zip(targets, results, strict=True)
+            if result is not None
+        }
 
         if not assembled:
             return checkpoint
@@ -683,18 +696,6 @@ class BaseCheckpointSaver(Generic[V]):
         clone.serde = maybe_add_typed_methods(serde)
         return clone
 
-    def _get_checkpoint_tuple_for_materialization(
-        self, config: RunnableConfig
-    ) -> CheckpointTuple | None:
-        """Internal raw checkpoint lookup used by fallback materialization."""
-        return self.get_tuple(config)
-
-    async def _aget_checkpoint_tuple_for_materialization(
-        self, config: RunnableConfig
-    ) -> CheckpointTuple | None:
-        """Async internal raw checkpoint lookup used by fallback materialization."""
-        return await self.aget_tuple(config)
-
     def _materialize_delta_channel(
         self,
         *,
@@ -738,7 +739,7 @@ class BaseCheckpointSaver(Generic[V]):
                     "checkpoint_id": prev_id,
                 }
             }
-            parent_tuple = self._get_checkpoint_tuple_for_materialization(parent_config)
+            parent_tuple = self.get_tuple(parent_config)
             if parent_tuple is None:
                 logger.warning(
                     "DeltaChannel chain broken: checkpoint %r not found for channel %r",
@@ -805,9 +806,7 @@ class BaseCheckpointSaver(Generic[V]):
                     "checkpoint_id": prev_id,
                 }
             }
-            parent_tuple = await self._aget_checkpoint_tuple_for_materialization(
-                parent_config
-            )
+            parent_tuple = await self.aget_tuple(parent_config)
             if parent_tuple is None:
                 logger.warning(
                     "DeltaChannel chain broken: checkpoint %r not found for channel %r",

libs/langgraph/langgraph/channels/delta.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 import collections.abc
 from collections.abc import Callable, Sequence
-from typing import Any, Generic
+from typing import Any, Generic, Literal
 
 from langgraph.checkpoint.base import DeltaChainValue, DeltaValue
 from typing_extensions import Self
@@ -99,7 +99,7 @@ class DeltaChannel(Generic[Value], BaseChannel[list[Value], Value, DeltaValue]):
         return self.typ | list[self.typ]  # type: ignore[name-defined]
 
     @property
-    def checkpoint_hydration_kind(self) -> str:
+    def checkpoint_hydration_kind(self) -> Literal["delta"]:
         return "delta"
 
     def copy(self) -> Self:

libs/langgraph/langgraph/pregel/_loop.py

@@ -12,6 +12,7 @@ from contextlib import (
     ExitStack,
 )
 from datetime import datetime, timezone
+from functools import cached_property
 from inspect import signature
 from types import TracebackType
 from typing import (
@@ -316,6 +317,7 @@ class PregelLoop:
         )
         self.prev_checkpoint_config = None
 
+    @cached_property
     def _checkpoint_hydration_plan(self) -> CheckpointHydrationPlan | None:
         """Build the saver hydration plan from this loop's channel specs."""
         return checkpoint_hydration_plan(self.specs)
@@ -326,7 +328,7 @@ class PregelLoop:
         if saved is None or self.checkpointer is None:
             return saved
         return self.checkpointer.materialize_checkpoint_tuple(
-            saved, self._checkpoint_hydration_plan()
+            saved, self._checkpoint_hydration_plan
         )
 
     async def _amaterialize_saved_checkpoint(
@@ -335,7 +337,7 @@ class PregelLoop:
         if saved is None or self.checkpointer is None:
             return saved
         return await self.checkpointer.amaterialize_checkpoint_tuple(
-            saved, self._checkpoint_hydration_plan()
+            saved, self._checkpoint_hydration_plan
         )
 
     def _push_graph_lifecycle_event(

libs/langgraph/langgraph/pregel/main.py

@@ -17,7 +17,7 @@ from collections.abc import (
     Sequence,
 )
 from dataclasses import is_dataclass, replace
-from functools import partial
+from functools import cached_property, partial
 from inspect import isclass
 from typing import (
     Any,
@@ -730,6 +730,7 @@ class Pregel(
             return checkpointer
         return _serde.apply_checkpointer_allowlist(checkpointer, self._serde_allowlist)
 
+    @cached_property
     def _checkpoint_hydration_plan(self) -> CheckpointHydrationPlan | None:
         return checkpoint_hydration_plan(self.channels)
 
@@ -741,7 +742,7 @@ class Pregel(
         if saved is None or checkpointer is None:
             return saved
         return checkpointer.materialize_checkpoint_tuple(
-            saved, self._checkpoint_hydration_plan()
+            saved, self._checkpoint_hydration_plan
         )
 
     async def _amaterialize_saved_checkpoint(
@@ -752,7 +753,7 @@ class Pregel(
         if saved is None or checkpointer is None:
             return saved
         return await checkpointer.amaterialize_checkpoint_tuple(
-            saved, self._checkpoint_hydration_plan()
+            saved, self._checkpoint_hydration_plan
         )
 
     def _materialize_saved_checkpoints(
@@ -763,7 +764,7 @@ class Pregel(
         if checkpointer is None or not saved:
             return list(saved)
         return checkpointer.materialize_checkpoint_tuples(
-            saved, self._checkpoint_hydration_plan()
+            saved, self._checkpoint_hydration_plan
         )
 
     async def _amaterialize_saved_checkpoints(
@@ -774,7 +775,7 @@ class Pregel(
         if checkpointer is None or not saved:
             return list(saved)
         return await checkpointer.amaterialize_checkpoint_tuples(
-            saved, self._checkpoint_hydration_plan()
+            saved, self._checkpoint_hydration_plan
         )
 
     def get_graph(