fix(checkpoint): handle metadata.writes when serializing old checkpoints with Jsonb (#6236)

Issue

Support for `Checkpoint.metadata.writes` was dropped in `langgraph`
v0.5.x.

In `langgraph-checkpoint-postgres` v2.0.23, metadata was serialized with
`BasePostgresSaver._dump_metadata` -> `JsonPlusSerializer.dumps` which
handles `pydantic.BaseModel`.

In v2.0.23, metadata is serialized with `psycopg.types.json.Jsonb`,
which raises `TypeError: Object of type AIMessage is not JSON
serializable` when trying to serialize `writes`.

Solution

- Add `BaseCheckpointSaver.get_serializable_checkpoint_metadata` which
pops the `writes` key.
- Log deprecation warning when strange version combinations are used 

Solves https://github.com/langchain-ai/langgraph/issues/5769

---------

Co-authored-by: Alex Kondratev <56111142+soapun@users.noreply.github.com>

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

@@ -14,7 +14,7 @@ from langgraph.checkpoint.base import (
     CheckpointMetadata,
     CheckpointTuple,
     get_checkpoint_id,
-    get_checkpoint_metadata,
+    get_serializable_checkpoint_metadata,
 )
 from langgraph.checkpoint.serde.base import SerializerProtocol
 from psycopg import Capabilities, Connection, Cursor, Pipeline
@@ -325,7 +325,7 @@ class PostgresSaver(BasePostgresSaver):
                     checkpoint["id"],
                     checkpoint_id,
                     Jsonb(copy),
-                    Jsonb(get_checkpoint_metadata(config, metadata)),
+                    Jsonb(get_serializable_checkpoint_metadata(config, metadata)),
                 ),
             )
         return next_config

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

@@ -14,7 +14,7 @@ from langgraph.checkpoint.base import (
     CheckpointMetadata,
     CheckpointTuple,
     get_checkpoint_id,
-    get_checkpoint_metadata,
+    get_serializable_checkpoint_metadata,
 )
 from langgraph.checkpoint.serde.base import SerializerProtocol
 from psycopg import AsyncConnection, AsyncCursor, AsyncPipeline, Capabilities
@@ -283,7 +283,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
                     checkpoint["id"],
                     checkpoint_id,
                     Jsonb(copy),
-                    Jsonb(get_checkpoint_metadata(config, metadata)),
+                    Jsonb(get_serializable_checkpoint_metadata(config, metadata)),
                 ),
             )
         return next_config

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

@@ -1,7 +1,9 @@
 from __future__ import annotations
 
 import random
+import warnings
 from collections.abc import Sequence
+from importlib.metadata import version as get_version
 from typing import Any, Optional, cast
 
 from langchain_core.runnables import RunnableConfig
@@ -16,6 +18,18 @@ from psycopg.types.json import Jsonb
 
 MetadataInput = Optional[dict[str, Any]]
 
+try:
+    major, minor = get_version("langgraph").split(".")[:2]
+    if int(major) == 0 and int(minor) < 5:
+        warnings.warn(
+            "You're using incompatible versions of langgraph and checkpoint-postgres. Please upgrade langgraph to avoid unexpected behavior.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+except Exception:
+    # skip version check if running from source
+    pass
+
 """
 To add a new migration, add a new string to the MIGRATIONS list.
 The position of the migration in the list is the version number.

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

@@ -12,7 +12,7 @@ from langgraph.checkpoint.base import (
     Checkpoint,
     CheckpointMetadata,
     CheckpointTuple,
-    get_checkpoint_metadata,
+    get_serializable_checkpoint_metadata,
 )
 from langgraph.checkpoint.serde.base import SerializerProtocol
 from langgraph.checkpoint.serde.types import TASKS
@@ -441,7 +441,7 @@ class ShallowPostgresSaver(BasePostgresSaver):
                     thread_id,
                     checkpoint_ns,
                     Jsonb(copy),
-                    Jsonb(get_checkpoint_metadata(config, metadata)),
+                    Jsonb(get_serializable_checkpoint_metadata(config, metadata)),
                 ),
             )
         return next_config
@@ -774,7 +774,7 @@ class AsyncShallowPostgresSaver(BasePostgresSaver):
                     thread_id,
                     checkpoint_ns,
                     Jsonb(copy),
-                    Jsonb(get_checkpoint_metadata(config, metadata)),
+                    Jsonb(get_serializable_checkpoint_metadata(config, metadata)),
                 ),
             )
         return next_config

libs/checkpoint-postgres/tests/test_async.py

@@ -187,13 +187,11 @@ def test_data():
     metadata_1: CheckpointMetadata = {
         "source": "input",
         "step": 2,
-        "writes": {},
         "score": 1,
     }
     metadata_2: CheckpointMetadata = {
         "source": "loop",
         "step": 1,
-        "writes": {"foo": "bar"},
         "score": None,
     }
     metadata_3: CheckpointMetadata = {}
@@ -220,7 +218,6 @@ async def test_combined_metadata(saver_name: str, test_data) -> None:
         metadata: CheckpointMetadata = {
             "source": "loop",
             "step": 1,
-            "writes": {"foo": "bar"},
             "score": None,
         }
         await saver.aput(config, chkpnt, metadata, {})
@@ -246,7 +243,6 @@ async def test_asearch(saver_name: str, test_data) -> None:
         query_1 = {"source": "input"}  # search by 1 key
         query_2 = {
             "step": 1,
-            "writes": {"foo": "bar"},
         }  # search by multiple keys
         query_3: dict[str, Any] = {}  # search by no keys, return all checkpoints
         query_4 = {"source": "update", "step": 1}  # no match

libs/checkpoint-postgres/tests/test_sync.py

@@ -169,13 +169,11 @@ def test_data():
     metadata_1: CheckpointMetadata = {
         "source": "input",
         "step": 2,
-        "writes": {},
         "score": 1,
     }
     metadata_2: CheckpointMetadata = {
         "source": "loop",
         "step": 1,
-        "writes": {"foo": "bar"},
         "score": None,
     }
     metadata_3: CheckpointMetadata = {}
@@ -202,7 +200,6 @@ def test_combined_metadata(saver_name: str, test_data) -> None:
         metadata: CheckpointMetadata = {
             "source": "loop",
             "step": 1,
-            "writes": {"foo": "bar"},
             "score": None,
         }
         saver.put(config, chkpnt, metadata, {})
@@ -228,7 +225,6 @@ def test_search(saver_name: str, test_data) -> None:
         query_1 = {"source": "input"}  # search by 1 key
         query_2 = {
             "step": 1,
-            "writes": {"foo": "bar"},
         }  # search by multiple keys
         query_3: dict[str, Any] = {}  # search by no keys, return all checkpoints
         query_4 = {"source": "update", "step": 1}  # no match

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

@@ -404,6 +404,16 @@ def get_checkpoint_metadata(
     return metadata
 
 
+def get_serializable_checkpoint_metadata(
+    config: RunnableConfig, metadata: CheckpointMetadata
+) -> CheckpointMetadata:
+    """Get checkpoint metadata in a backwards-compatible manner."""
+    checkpoint_metadata = get_checkpoint_metadata(config, metadata)
+    if "writes" in checkpoint_metadata:
+        checkpoint_metadata.pop("writes")
+    return checkpoint_metadata
+
+
 """
 Mapping from error type to error index.
 Regular writes just map to their index in the list of writes being saved.