feat(checkpoint/postgres): diff chain reconstruction in _load_blobs (sync)

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Sydney Runkle
2026-04-17 15:37:18 -04:00
co-authored by Claude Sonnet 4.6
parent dba1987c9b
commit d4e1efa1f6
3 changed files with 80 additions and 7 deletions
@@ -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"],
@@ -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,
@@ -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)