fix(checkpoint/postgres): pass cursor to avoid deadlock in diff chain traversal

Fixes a critical deadlock that occurs when _load_diff_chains calls self._cursor()
from within _load_blobs while the outer _load_checkpoint_tuple already holds
self._cursor(). On bare (non-pool) connections, the threading.Lock is not
reentrant, causing a deadlock.

Solution: Pass the cursor as a parameter to _load_diff_chains and _load_blobs
instead of acquiring a new cursor within those methods. Updated _load_checkpoint_tuple
to acquire a cursor once at the top level and pass it through the call chain.

Changes:
- Updated _load_blobs signature to accept optional cur parameter
- Updated _load_diff_chains signature (base and implementations) to accept optional cur parameter
- Modified _load_checkpoint_tuple in PostgresSaver to acquire cursor and pass it
- Modified _load_checkpoint_tuple_async to acquire cursor only when diff_payloads exist
- Removed nested self._cursor() calls in _load_diff_chains and _load_diff_chains_async

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py

@@ -435,39 +435,41 @@ class PostgresSaver(BasePostgresSaver):
         thread_id: str,
         checkpoint_ns: str,
         diff_channel_payloads: dict[str, dict[str, Any]],
+        *,
+        cur: Any = None,
     ) -> 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
+        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
+            visited: set[str] = set()
 
-                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]
+            while version_cursor is not None:
+                if version_cursor in visited:
+                    break
+                visited.add(version_cursor)
+                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
+                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:
@@ -482,6 +484,13 @@ class PostgresSaver(BasePostgresSaver):
             including its configuration, metadata, parent checkpoint (if any),
             and pending writes.
         """
+        with self._cursor() as cur:
+            channel_values = self._load_blobs(
+                value["channel_values"],
+                thread_id=value["thread_id"],
+                checkpoint_ns=value["checkpoint_ns"],
+                cur=cur,
+            )
         return CheckpointTuple(
             {
                 "configurable": {
@@ -494,11 +503,7 @@ class PostgresSaver(BasePostgresSaver):
                 **value["checkpoint"],
                 "channel_values": {
                     **(value["checkpoint"].get("channel_values") or {}),
-                    **self._load_blobs(
-                        value["channel_values"],
-                        thread_id=value["thread_id"],
-                        checkpoint_ns=value["checkpoint_ns"],
-                    ),
+                    **channel_values,
                 },
             },
             value["metadata"],

libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py

@@ -396,38 +396,41 @@ class AsyncPostgresSaver(BasePostgresSaver):
         thread_id: str,
         checkpoint_ns: str,
         diff_channel_payloads: dict[str, dict[str, Any]],
+        *,
+        cur: Any,
     ) -> dict[str, Any]:
         from langgraph.checkpoint.base import DiffChainValue
 
         result: dict[str, Any] = {}
-        async 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
+        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
+            visited: set[str] = set()
 
-                while version_cursor is not None:
-                    await 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 = await cur.fetchone()
-                    if row is None:
-                        break
-                    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]
+            while version_cursor is not None:
+                if version_cursor in visited:
+                    break
+                visited.add(version_cursor)
+                await 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 = await cur.fetchone()
+                if row is None:
+                    break
+                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
 
     async def _load_checkpoint_tuple(self, value: DictRow) -> CheckpointTuple:
@@ -457,11 +460,13 @@ class AsyncPostgresSaver(BasePostgresSaver):
                 elif type_tag != "empty":
                     non_diff[channel] = self.serde.loads_typed((type_tag, v))
 
-        diff_values = (
-            await self._load_diff_chains_async(thread_id, checkpoint_ns, diff_payloads)
-            if diff_payloads
-            else {}
-        )
+        if diff_payloads:
+            async with self._cursor() as cur:
+                diff_values = await self._load_diff_chains_async(
+                    thread_id, checkpoint_ns, diff_payloads, cur=cur
+                )
+        else:
+            diff_values = {}
 
         return CheckpointTuple(
             {

libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py

@@ -190,6 +190,7 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
         *,
         thread_id: str = "",
         checkpoint_ns: str = "",
+        cur: Any = None,
     ) -> dict[str, Any]:
         if not blob_values:
             return {}
@@ -201,15 +202,13 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
             channel = k.decode()
             type_tag = t.decode()
             if type_tag == "diff":
-                diff_channel_payloads[channel] = self.serde.loads_typed(
-                    (type_tag, v)
-                )
+                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)
+                self._load_diff_chains(thread_id, checkpoint_ns, diff_channel_payloads, cur=cur)
             )
 
         return result
@@ -219,6 +218,8 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
         thread_id: str,
         checkpoint_ns: str,
         diff_channel_payloads: dict[str, dict[str, Any]],
+        *,
+        cur: Any = None,
     ) -> dict[str, Any]:
         """Override in sync/async subclasses. Resolves diff-chain blobs to DiffChainValue."""
         raise NotImplementedError