From d4e1efa1f66381e056800efa10eac0aa5fcf32c0 Mon Sep 17 00:00:00 2001
From: Sydney Runkle <sydneymarierunkle@gmail.com>
Date: Fri, 17 Apr 2026 15:37:18 -0400
Subject: [PATCH] feat(checkpoint/postgres): diff chain reconstruction in
 _load_blobs (sync)

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 .../langgraph/checkpoint/postgres/__init__.py | 46 ++++++++++++++++++-
 .../langgraph/checkpoint/postgres/base.py     | 40 +++++++++++++---
 .../langgraph/checkpoint/serde/jsonplus.py    |  1 +
 3 files changed, 80 insertions(+), 7 deletions(-)

diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py
index 600127333..a771cce47 100644
--- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py
+++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py
@@ -430,6 +430,46 @@ class PostgresSaver(BasePostgresSaver):
                 with conn.cursor(binary=True, row_factory=dict_row) as cur:
                     yield cur
 
+    def _load_diff_chains(
+        self,
+        thread_id: str,
+        checkpoint_ns: str,
+        diff_channel_payloads: dict[str, dict[str, Any]],
+    ) -> dict[str, Any]:
+        from langgraph.checkpoint.base import DiffChainValue
+
+        result: dict[str, Any] = {}
+        with self._cursor() as cur:
+            for channel, current_payload in diff_channel_payloads.items():
+                payloads: list[dict[str, Any]] = [current_payload]
+                version_cursor: str | None = current_payload["p"]
+                base: list[Any] | None = None
+
+                while version_cursor is not None:
+                    cur.execute(
+                        "SELECT type, blob FROM checkpoint_blobs "
+                        "WHERE thread_id = %s AND checkpoint_ns = %s "
+                        "AND channel = %s AND version = %s",
+                        (thread_id, checkpoint_ns, channel, version_cursor),
+                    )
+                    row = cur.fetchone()
+                    if row is None:
+                        break
+                    # row is a dict (dict_row factory): {"type": str, "blob": bytes}
+                    if row["type"] == "diff":
+                        payload = self.serde.loads_typed(("diff", row["blob"]))
+                        payloads.append(payload)
+                        version_cursor = payload["p"]
+                    else:
+                        base = self.serde.loads_typed((row["type"], row["blob"]))
+                        break
+
+                payloads.reverse()
+                result[channel] = DiffChainValue(
+                    base=base, deltas=[p["d"] for p in payloads]
+                )
+        return result
+
     def _load_checkpoint_tuple(self, value: DictRow) -> CheckpointTuple:
         """
         Convert a database row into a CheckpointTuple object.
@@ -454,7 +494,11 @@ class PostgresSaver(BasePostgresSaver):
                 **value["checkpoint"],
                 "channel_values": {
                     **(value["checkpoint"].get("channel_values") or {}),
-                    **self._load_blobs(value["channel_values"]),
+                    **self._load_blobs(
+                        value["channel_values"],
+                        thread_id=value["thread_id"],
+                        checkpoint_ns=value["checkpoint_ns"],
+                    ),
                 },
             },
             value["metadata"],
diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py
index 8f27fee5b..68b61d874 100644
--- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py
+++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py
@@ -185,15 +185,43 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
         )
 
     def _load_blobs(
-        self, blob_values: list[tuple[bytes, bytes, bytes]]
+        self,
+        blob_values: list[tuple[bytes, bytes, bytes]],
+        *,
+        thread_id: str = "",
+        checkpoint_ns: str = "",
     ) -> dict[str, Any]:
         if not blob_values:
             return {}
-        return {
-            k.decode(): self.serde.loads_typed((t.decode(), v))
-            for k, t, v in blob_values
-            if t.decode() != "empty"
-        }
+
+        result: dict[str, Any] = {}
+        diff_channel_payloads: dict[str, dict[str, Any]] = {}
+
+        for k, t, v in blob_values:
+            channel = k.decode()
+            type_tag = t.decode()
+            if type_tag == "diff":
+                diff_channel_payloads[channel] = self.serde.loads_typed(
+                    (type_tag, v)
+                )
+            elif type_tag != "empty":
+                result[channel] = self.serde.loads_typed((type_tag, v))
+
+        if diff_channel_payloads:
+            result.update(
+                self._load_diff_chains(thread_id, checkpoint_ns, diff_channel_payloads)
+            )
+
+        return result
+
+    def _load_diff_chains(
+        self,
+        thread_id: str,
+        checkpoint_ns: str,
+        diff_channel_payloads: dict[str, dict[str, Any]],
+    ) -> dict[str, Any]:
+        """Override in sync/async subclasses. Resolves diff-chain blobs to DiffChainValue."""
+        raise NotImplementedError
 
     def _dump_blobs(
         self,
diff --git a/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py b/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py
index f9362089b..181f195b6 100644
--- a/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py
+++ b/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py
@@ -49,6 +49,7 @@ logger = logging.getLogger(__name__)
 
 def _is_diff_delta(obj: Any) -> bool:
     from langgraph.checkpoint.base import DiffDelta  # lazy import avoids circular dep
+
     return isinstance(obj, DiffDelta)