diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py index 54844c824..bf349856e 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py @@ -8,6 +8,7 @@ from typing import Any from langchain_core.runnables import RunnableConfig from langgraph.checkpoint.base import ( + DELTA_SENTINEL, WRITES_IDX_MAP, ChannelVersions, Checkpoint, @@ -448,7 +449,15 @@ class PostgresSaver(BasePostgresSaver): """ thread_id = config["configurable"]["thread_id"] checkpoint_ns = config["configurable"].get("checkpoint_ns", "") - checkpoint_id = config["configurable"]["checkpoint_id"] + checkpoint_id = get_checkpoint_id(config) + if checkpoint_id is None: + # Caller didn't specify a target — resolve to the latest + # checkpoint on the thread. `get_tuple` without `checkpoint_id` + # returns the newest; its config carries the resolved id. + target = self.get_tuple(config) + if target is None: + return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=[]) + checkpoint_id = target.config["configurable"]["checkpoint_id"] with self._cursor() as cur: cur.execute(SELECT_DELTA_PARENTS_SQL, (channel, thread_id, checkpoint_ns)) parents_rows = cur.fetchall() diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py index dc5a655e1..d76c10bf5 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py @@ -8,6 +8,7 @@ from typing import Any from langchain_core.runnables import RunnableConfig from langgraph.checkpoint.base import ( + DELTA_SENTINEL, WRITES_IDX_MAP, ChannelVersions, Checkpoint, @@ -409,7 +410,12 @@ class AsyncPostgresSaver(BasePostgresSaver): """ thread_id = config["configurable"]["thread_id"] checkpoint_ns = config["configurable"].get("checkpoint_ns", "") - checkpoint_id = config["configurable"]["checkpoint_id"] + checkpoint_id = get_checkpoint_id(config) + if checkpoint_id is None: + target = await self.aget_tuple(config) + if target is None: + return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=[]) + checkpoint_id = target.config["configurable"]["checkpoint_id"] async with self._cursor() as cur: await cur.execute( SELECT_DELTA_PARENTS_SQL, (channel, thread_id, checkpoint_ns) diff --git a/libs/checkpoint-postgres/tests/test_async.py b/libs/checkpoint-postgres/tests/test_async.py index 9dc5c78cd..e342548a7 100644 --- a/libs/checkpoint-postgres/tests/test_async.py +++ b/libs/checkpoint-postgres/tests/test_async.py @@ -361,9 +361,9 @@ async def test_get_checkpoint_no_channel_values( load_checkpoint_tuple = saver._load_checkpoint_tuple - async def patched_load_checkpoint_tuple(value, cur): + async def patched_load_checkpoint_tuple(value): value["checkpoint"].pop("channel_values", None) - return await load_checkpoint_tuple(value, cur) + return await load_checkpoint_tuple(value) monkeypatch.setattr( saver, "_load_checkpoint_tuple", patched_load_checkpoint_tuple diff --git a/libs/checkpoint-postgres/tests/test_sync.py b/libs/checkpoint-postgres/tests/test_sync.py index cc78c615e..63b93715f 100644 --- a/libs/checkpoint-postgres/tests/test_sync.py +++ b/libs/checkpoint-postgres/tests/test_sync.py @@ -348,9 +348,9 @@ def test_get_checkpoint_no_channel_values( load_checkpoint_tuple = saver._load_checkpoint_tuple - def patched_load_checkpoint_tuple(value, cur): + def patched_load_checkpoint_tuple(value): value["checkpoint"].pop("channel_values", None) - return load_checkpoint_tuple(value, cur) + return load_checkpoint_tuple(value) monkeypatch.setattr( saver, "_load_checkpoint_tuple", patched_load_checkpoint_tuple