Remove Checkpoint.pending_sends

- Instead store sends in a Topic channel, removing the need to fetch sends as writes against the parent checkpoint
- Remove deprecated/unused functions in langgraph-checkpoint (will require bumping min range for langgraph-checkpoint in langgraph lib)
- Implement migration of old pending sends in langgraph-checkpoint-postgres
- Ensure parent config of `checkpoint_during=False` checkpoints always points to checkpoints that were also saved

libs/checkpoint-postgres/tests/checkpoint_utils.py

@@ -0,0 +1,51 @@
+from collections.abc import Mapping
+from datetime import datetime, timezone
+from typing import Any, Optional, Protocol
+
+from langgraph.checkpoint.base import Checkpoint, EmptyChannelError
+from langgraph.checkpoint.base.id import uuid6
+
+
+class ChannelProtocol(Protocol):
+    def checkpoint(self) -> Optional[Any]: ...
+
+
+def empty_checkpoint() -> Checkpoint:
+    return Checkpoint(
+        v=1,
+        id=str(uuid6(clock_seq=-2)),
+        ts=datetime.now(timezone.utc).isoformat(),
+        channel_values={},
+        channel_versions={},
+        versions_seen={},
+    )
+
+
+def create_checkpoint(
+    checkpoint: Checkpoint,
+    channels: Optional[Mapping[str, ChannelProtocol]],
+    step: int,
+    *,
+    id: Optional[str] = None,
+) -> Checkpoint:
+    """Create a checkpoint for the given channels."""
+    ts = datetime.now(timezone.utc).isoformat()
+    if channels is None:
+        values = checkpoint["channel_values"]
+    else:
+        values = {}
+        for k, v in channels.items():
+            if k not in checkpoint["channel_versions"]:
+                continue
+            try:
+                values[k] = v.checkpoint()
+            except EmptyChannelError:
+                pass
+    return Checkpoint(
+        v=1,
+        ts=ts,
+        id=id or str(uuid6(clock_seq=step)),
+        channel_values=values,
+        channel_versions=checkpoint["channel_versions"],
+        versions_seen=checkpoint["versions_seen"],
+    )

libs/checkpoint-postgres/tests/test_async.py

@@ -14,10 +14,10 @@ from langgraph.checkpoint.base import (
     EXCLUDED_METADATA_KEYS,
     Checkpoint,
     CheckpointMetadata,
-    create_checkpoint,
-    empty_checkpoint,
 )
 from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
+from langgraph.checkpoint.serde.types import TASKS
+from tests.checkpoint_utils import create_checkpoint, empty_checkpoint
 from tests.conftest import DEFAULT_POSTGRES_URI
 
 
@@ -263,3 +263,52 @@ async def test_null_chars(saver_name: str, test_data) -> None:
         assert [c async for c in saver.alist(None, filter={"my_key": "abc"})][
             0
         ].metadata["my_key"] == "abc"
+
+
+@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe"])
+async def test_pending_sends_migration(saver_name: str) -> None:
+    async with _saver(saver_name) as saver:
+        config = {
+            "configurable": {
+                "thread_id": "thread-1",
+                "checkpoint_ns": "",
+            }
+        }
+
+        # create the first checkpoint
+        # and put some pending sends
+        checkpoint_0 = empty_checkpoint()
+        config = await saver.aput(config, checkpoint_0, {}, {})
+        await saver.aput_writes(
+            config, [(TASKS, "send-1"), (TASKS, "send-2")], task_id="task-1"
+        )
+        await saver.aput_writes(config, [(TASKS, "send-3")], task_id="task-2")
+
+        # check that fetching checkpoint_0 doesn't attach pending sends
+        # (they should be attached to the next checkpoint)
+        tuple_0 = await saver.aget_tuple(config)
+        assert tuple_0.checkpoint["channel_values"] == {}
+        assert tuple_0.checkpoint["channel_versions"] == {}
+
+        # create the second checkpoint
+        checkpoint_1 = create_checkpoint(checkpoint_0, {}, 1)
+        config = await saver.aput(config, checkpoint_1, {}, {})
+
+        # check that pending sends are attached to checkpoint_1
+        tuple_1 = await saver.aget_tuple(config)
+        assert tuple_1.checkpoint["channel_values"] == {
+            TASKS: ["send-1", "send-2", "send-3"]
+        }
+        assert TASKS in tuple_1.checkpoint["channel_versions"]
+
+        # check that list also applies the migration
+        search_results = [
+            c async for c in saver.alist({"configurable": {"thread_id": "thread-1"}})
+        ]
+        assert len(search_results) == 2
+        assert search_results[-1].checkpoint["channel_values"] == {}
+        assert search_results[-1].checkpoint["channel_versions"] == {}
+        assert search_results[0].checkpoint["channel_values"] == {
+            TASKS: ["send-1", "send-2", "send-3"]
+        }
+        assert TASKS in search_results[0].checkpoint["channel_versions"]

libs/checkpoint-postgres/tests/test_sync.py

@@ -15,10 +15,10 @@ from langgraph.checkpoint.base import (
     EXCLUDED_METADATA_KEYS,
     Checkpoint,
     CheckpointMetadata,
-    create_checkpoint,
-    empty_checkpoint,
 )
 from langgraph.checkpoint.postgres import PostgresSaver
+from langgraph.checkpoint.serde.types import TASKS
+from tests.checkpoint_utils import create_checkpoint, empty_checkpoint
 from tests.conftest import DEFAULT_POSTGRES_URI
 
 
@@ -258,3 +258,52 @@ def test_nonnull_migrations() -> None:
     for migration in PostgresSaver.MIGRATIONS:
         statement = _leading_comment_remover.sub("", migration).split()[0]
         assert statement.strip()
+
+
+@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe"])
+def test_pending_sends_migration(saver_name: str) -> None:
+    with _saver(saver_name) as saver:
+        config = {
+            "configurable": {
+                "thread_id": "thread-1",
+                "checkpoint_ns": "",
+            }
+        }
+
+        # create the first checkpoint
+        # and put some pending sends
+        checkpoint_0 = empty_checkpoint()
+        config = saver.put(config, checkpoint_0, {}, {})
+        saver.put_writes(
+            config, [(TASKS, "send-1"), (TASKS, "send-2")], task_id="task-1"
+        )
+        saver.put_writes(config, [(TASKS, "send-3")], task_id="task-2")
+
+        # check that fetching checkpoint_0 doesn't attach pending sends
+        # (they should be attached to the next checkpoint)
+        tuple_0 = saver.get_tuple(config)
+        assert tuple_0.checkpoint["channel_values"] == {}
+        assert tuple_0.checkpoint["channel_versions"] == {}
+
+        # create the second checkpoint
+        checkpoint_1 = create_checkpoint(checkpoint_0, {}, 1)
+        config = saver.put(config, checkpoint_1, {}, {})
+
+        # check that pending sends are attached to checkpoint_1
+        checkpoint_1 = saver.get_tuple(config)
+        assert checkpoint_1.checkpoint["channel_values"] == {
+            TASKS: ["send-1", "send-2", "send-3"]
+        }
+        assert TASKS in checkpoint_1.checkpoint["channel_versions"]
+
+        # check that list also applies the migration
+        search_results = [
+            c for c in saver.list({"configurable": {"thread_id": "thread-1"}})
+        ]
+        assert len(search_results) == 2
+        assert search_results[-1].checkpoint["channel_values"] == {}
+        assert search_results[-1].checkpoint["channel_versions"] == {}
+        assert search_results[0].checkpoint["channel_values"] == {
+            TASKS: ["send-1", "send-2", "send-3"]
+        }
+        assert TASKS in search_results[0].checkpoint["channel_versions"]