From 1ba96f49bfcbd41f4c93a94b61f6a31fe9f48c22 Mon Sep 17 00:00:00 2001 From: Caspar Broekhuizen Date: Mon, 6 Oct 2025 11:27:34 -0700 Subject: [PATCH] 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> --- .../langgraph/checkpoint/postgres/__init__.py | 4 ++-- .../langgraph/checkpoint/postgres/aio.py | 4 ++-- .../langgraph/checkpoint/postgres/base.py | 14 ++++++++++++++ .../langgraph/checkpoint/postgres/shallow.py | 6 +++--- libs/checkpoint-postgres/tests/test_async.py | 4 ---- libs/checkpoint-postgres/tests/test_sync.py | 4 ---- .../langgraph/checkpoint/base/__init__.py | 10 ++++++++++ 7 files changed, 31 insertions(+), 15 deletions(-) diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py index b689df5d5..0943085cc 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py +++ b/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 diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py index ed36daa4a..f4ce2e4f1 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py +++ b/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 diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py index b1ee5c9ce..2f5f8817f 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py +++ b/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. diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/shallow.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/shallow.py index 1abcaf96c..05b816e12 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/shallow.py +++ b/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 diff --git a/libs/checkpoint-postgres/tests/test_async.py b/libs/checkpoint-postgres/tests/test_async.py index cca805901..a46163395 100644 --- a/libs/checkpoint-postgres/tests/test_async.py +++ b/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 diff --git a/libs/checkpoint-postgres/tests/test_sync.py b/libs/checkpoint-postgres/tests/test_sync.py index c8564c4e4..63b93715f 100644 --- a/libs/checkpoint-postgres/tests/test_sync.py +++ b/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 diff --git a/libs/checkpoint/langgraph/checkpoint/base/__init__.py b/libs/checkpoint/langgraph/checkpoint/base/__init__.py index a2a768d0c..2b547f1cd 100644 --- a/libs/checkpoint/langgraph/checkpoint/base/__init__.py +++ b/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.