langgraph-checkpoint-postgres 3.1.3__tar.gz → 3.2.0__tar.gz

{langgraph_checkpoint_postgres-3.1.3 → langgraph_checkpoint_postgres-3.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: langgraph-checkpoint-postgres
-Version: 3.1.3
+Version: 3.2.0
 Summary: Library with a Postgres implementation of LangGraph checkpoint saver.
 Project-URL: Source, https://github.com/langchain-ai/langgraph/tree/main/libs/checkpoint-postgres
 Project-URL: Twitter, https://x.com/langchain_oss

{langgraph_checkpoint_postgres-3.1.3 → langgraph_checkpoint_postgres-3.2.0}/langgraph/checkpoint/postgres/__init__.py

@@ -27,10 +27,15 @@ from psycopg_pool import ConnectionPool
 from langgraph.checkpoint.postgres import _internal
 from langgraph.checkpoint.postgres.base import (
     _DELTA_PAGE_SIZE,
+    PRUNE_DELETE_THREADS_SQL,
+    PRUNE_STRATEGIES,
     BasePostgresSaver,
     _build_delta_stage1_sql,
     _build_delta_stage2_sql,
     _DeltaStage2Row,
+    _execute,
+    _keep_latest,
+    _refuse_delta_threads,
 )
 from langgraph.checkpoint.postgres.shallow import ShallowPostgresSaver
 
@@ -401,6 +406,73 @@ class PostgresSaver(BasePostgresSaver):
                 (str(thread_id),),
             )
 
+    def prune(
+        self,
+        thread_ids: Sequence[str],
+        *,
+        strategy: str = "keep_latest",
+    ) -> None:
+        """Prune checkpoints for the given threads.
+
+        `"delete"` removes every checkpoint, write and blob of the threads, in
+        one transaction.
+
+        `"keep_latest"` keeps the latest checkpoint of each namespace (the one
+        `get_tuple` loads, plus its parent when it predates v4 and has pending
+        sends), with its pending writes and the blobs it references, and
+        deletes the rest. A kept checkpoint whose parent it deletes is left with
+        no parent. A subgraph call's namespace is deleted with the root
+        checkpoint the call started from, except a `checkpointer=True`
+        subgraph's, which keeps its latest checkpoint. Each thread is pruned in
+        its own transaction, on one snapshot of it, so a run appending to the
+        thread meanwhile is safe; forking or replaying from a checkpoint being
+        deleted is not. A thread that fails because another transaction changed
+        it leaves the threads before it pruned, and retrying is safe.
+
+        `"keep_latest"` refuses threads that use a `DeltaChannel`, since it
+        rebuilds the latest value from older checkpoints: if any requested
+        thread does, nothing is pruned. A thread that starts using one before
+        its turn in the batch is refused then, after the threads before it.
+
+        Args:
+            thread_ids: The thread IDs to prune.
+            strategy: `"keep_latest"` or `"delete"`.
+
+        Raises:
+            TypeError: If `thread_ids` is a single string.
+            ValueError: If `strategy` is unknown.
+            NotImplementedError: If `strategy` is `"keep_latest"` and a thread
+                uses a `DeltaChannel`.
+            psycopg.errors.SerializationFailure: If another transaction, such
+                as a second prune, changes a thread's rows while `"keep_latest"`
+                prunes it.
+            psycopg.errors.DeadlockDetected: If a second prune of a thread runs
+                at the same time.
+        """
+        if isinstance(thread_ids, str):
+            raise TypeError("thread_ids must be a sequence of thread IDs, not a str")
+        if strategy not in PRUNE_STRATEGIES:
+            raise ValueError(f"Unknown prune strategy: {strategy!r}")
+        thread_ids = [str(t) for t in thread_ids]
+        if not thread_ids:
+            return
+        if strategy == "delete":
+            with self._cursor() as cur, cur.connection.transaction():
+                for sql in PRUNE_DELETE_THREADS_SQL:
+                    cur.execute(sql, (thread_ids,))
+            return
+        with self._cursor() as cur:
+            _execute(cur, _refuse_delta_threads(thread_ids))
+        for thread_id in thread_ids:
+            with self._cursor() as cur:
+                if self.pipe:
+                    # The transaction has to start with its isolation level, and
+                    # in pipeline mode it would take in what ran since the last
+                    # sync.
+                    self.pipe.sync()
+                with cur.connection.transaction():
+                    _execute(cur, _keep_latest(thread_id))
+
     @contextmanager
     def _cursor(self, *, pipeline: bool = False) -> Iterator[Cursor[DictRow]]:
         """Create a database cursor as a context manager.

{langgraph_checkpoint_postgres-3.1.3 → langgraph_checkpoint_postgres-3.2.0}/langgraph/checkpoint/postgres/aio.py

@@ -27,10 +27,15 @@ from psycopg_pool import AsyncConnectionPool
 from langgraph.checkpoint.postgres import _ainternal
 from langgraph.checkpoint.postgres.base import (
     _DELTA_PAGE_SIZE,
+    PRUNE_DELETE_THREADS_SQL,
+    PRUNE_STRATEGIES,
     BasePostgresSaver,
+    _aexecute,
     _build_delta_stage1_sql,
     _build_delta_stage2_sql,
     _DeltaStage2Row,
+    _keep_latest,
+    _refuse_delta_threads,
 )
 from langgraph.checkpoint.postgres.shallow import AsyncShallowPostgresSaver
 
@@ -360,6 +365,55 @@ class AsyncPostgresSaver(BasePostgresSaver):
                 (str(thread_id),),
             )
 
+    async def aprune(
+        self,
+        thread_ids: Sequence[str],
+        *,
+        strategy: str = "keep_latest",
+    ) -> None:
+        """Prune checkpoints for the given threads.
+
+        See `PostgresSaver.prune` for what each strategy keeps.
+
+        Args:
+            thread_ids: The thread IDs to prune.
+            strategy: `"keep_latest"` or `"delete"`.
+
+        Raises:
+            TypeError: If `thread_ids` is a single string.
+            ValueError: If `strategy` is unknown.
+            NotImplementedError: If `strategy` is `"keep_latest"` and a thread
+                uses a `DeltaChannel`.
+            psycopg.errors.SerializationFailure: If another transaction, such
+                as a second prune, changes a thread's rows while `"keep_latest"`
+                prunes it.
+            psycopg.errors.DeadlockDetected: If a second prune of a thread runs
+                at the same time.
+        """
+        if isinstance(thread_ids, str):
+            raise TypeError("thread_ids must be a sequence of thread IDs, not a str")
+        if strategy not in PRUNE_STRATEGIES:
+            raise ValueError(f"Unknown prune strategy: {strategy!r}")
+        thread_ids = [str(t) for t in thread_ids]
+        if not thread_ids:
+            return
+        if strategy == "delete":
+            async with self._cursor() as cur, cur.connection.transaction():
+                for sql in PRUNE_DELETE_THREADS_SQL:
+                    await cur.execute(sql, (thread_ids,))
+            return
+        async with self._cursor() as cur:
+            await _aexecute(cur, _refuse_delta_threads(thread_ids))
+        for thread_id in thread_ids:
+            async with self._cursor() as cur:
+                if self.pipe:
+                    # The transaction has to start with its isolation level, and
+                    # in pipeline mode it would take in what ran since the last
+                    # sync.
+                    await self.pipe.sync()
+                async with cur.connection.transaction():
+                    await _aexecute(cur, _keep_latest(thread_id))
+
     @asynccontextmanager
     async def _cursor(
         self, *, pipeline: bool = False
@@ -688,6 +742,33 @@ class AsyncPostgresSaver(BasePostgresSaver):
             self.adelete_thread(thread_id), self.loop
         ).result()
 
+    def prune(
+        self,
+        thread_ids: Sequence[str],
+        *,
+        strategy: str = "keep_latest",
+    ) -> None:
+        """Prune checkpoints for the given threads.
+
+        See `PostgresSaver.prune` for what each strategy keeps.
+
+        Args:
+            thread_ids: The thread IDs to prune.
+            strategy: `"keep_latest"` or `"delete"`.
+        """
+        try:
+            if asyncio.get_running_loop() is self.loop:
+                raise asyncio.InvalidStateError(
+                    "Synchronous calls to AsyncPostgresSaver are only allowed from a "
+                    "different thread. From the main thread, use the async interface. "
+                    "For example, use `await checkpointer.aprune(...)`."
+                )
+        except RuntimeError:
+            pass
+        return asyncio.run_coroutine_threadsafe(
+            self.aprune(thread_ids, strategy=strategy), self.loop
+        ).result()
+
     def get_delta_channel_history(
         self, *, config: RunnableConfig, channels: Sequence[str]
     ) -> Mapping[str, DeltaChannelHistory]:

{langgraph_checkpoint_postgres-3.1.3 → langgraph_checkpoint_postgres-3.2.0}/langgraph/checkpoint/postgres/base.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 import random
 import warnings
-from collections.abc import Mapping, Sequence
+from collections.abc import Generator, Mapping, Sequence
 from importlib.metadata import version as get_version
 from typing import Any, TypedDict, cast
 
@@ -17,6 +17,8 @@ from langgraph.checkpoint.base import (
     writes_sort_key,
 )
 from langgraph.checkpoint.serde.types import TASKS
+from psycopg import AsyncCursor, Cursor
+from psycopg.rows import DictRow
 from psycopg.types.json import Jsonb
 
 # Page size for stage-1 paged scan in `get_delta_channel_history`. Internal
@@ -160,6 +162,218 @@ INSERT_CHECKPOINT_WRITES_SQL = """
 """
 
 
+PRUNE_STRATEGIES = ("keep_latest", "delete")
+
+PRUNE_DELETE_THREADS_SQL = (
+    "DELETE FROM checkpoints WHERE thread_id = ANY(%s)",
+    "DELETE FROM checkpoint_blobs WHERE thread_id = ANY(%s)",
+    "DELETE FROM checkpoint_writes WHERE thread_id = ANY(%s)",
+)
+
+# The prune statements below never join two sets of a thread's rows: a join is
+# planned from per-thread row estimates, which are worst for the few huge
+# threads worth pruning, and can go quadratic. They look rows up by primary key
+# instead (`LATERAL ... LIMIT 1`, `= ANY`).
+
+# A `DeltaChannel` leaves update counters in checkpoint metadata, or a
+# snapshot: an inline `true` marker with a blob behind it, which a plain `True`
+# value never has. LangGraph.js inlines no values, but marks the threads its
+# DeltaChannel graphs start with `delta_writes_versioned`.
+PRUNE_DELTA_THREADS_SQL = """
+SELECT thread_id FROM checkpoints
+WHERE thread_id = ANY(%s) AND (
+    metadata ? 'counters_since_delta_snapshot'
+    OR metadata ? 'delta_writes_versioned'
+)
+UNION
+SELECT c.thread_id
+FROM checkpoints c
+CROSS JOIN LATERAL jsonb_each(c.checkpoint -> 'channel_values') AS v
+CROSS JOIN LATERAL (
+    SELECT 1 FROM checkpoint_blobs b
+    WHERE b.thread_id = c.thread_id
+    AND b.checkpoint_ns = c.checkpoint_ns
+    AND b.channel = v.key
+    AND b.version = c.checkpoint -> 'channel_versions' ->> v.key
+    LIMIT 1
+) AS snapshot
+WHERE c.thread_id = ANY(%s) AND v.value = 'true'::jsonb
+"""
+
+# Per namespace: the latest checkpoint, the parent it reads its pending sends
+# from when it predates v4, and the root graph's checkpoint its subgraph call
+# started from, with whether that one is saved.
+PRUNE_HEADS_SQL = f"""
+SELECT
+    h.checkpoint_ns,
+    h.checkpoint_id AS head,
+    CASE WHEN (c.checkpoint ->> 'v')::int < 4 AND EXISTS (
+        SELECT 1 FROM checkpoint_writes w
+        WHERE w.thread_id = %s
+        AND w.checkpoint_ns = h.checkpoint_ns
+        AND w.checkpoint_id = c.parent_checkpoint_id
+        AND w.channel = '{TASKS}'
+    ) THEN c.parent_checkpoint_id END AS sends_parent,
+    c.metadata -> 'parents' ->> '' AS root,
+    EXISTS (
+        SELECT 1 FROM checkpoints r
+        WHERE r.thread_id = %s AND r.checkpoint_ns = ''
+        AND r.checkpoint_id = c.metadata -> 'parents' ->> ''
+    ) AS root_saved
+FROM (
+    SELECT DISTINCT ON (checkpoint_ns) checkpoint_ns, checkpoint_id
+    FROM checkpoints
+    WHERE thread_id = %s
+    ORDER BY checkpoint_ns DESC, checkpoint_id DESC
+) AS h
+CROSS JOIN LATERAL (
+    SELECT checkpoint, parent_checkpoint_id, metadata FROM checkpoints
+    WHERE thread_id = %s
+    AND checkpoint_ns = h.checkpoint_ns
+    AND checkpoint_id = h.checkpoint_id
+) AS c
+"""
+
+PRUNE_SNAPSHOT_SQL = "SET TRANSACTION ISOLATION LEVEL REPEATABLE READ"
+
+PRUNE_UNKEPT_SQL = """
+SELECT checkpoint_id FROM checkpoints
+WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id <> ALL(%s)
+"""
+
+# A run saves a superstep's writes in the background, where they can commit
+# before their checkpoint does, so they go by the checkpoints deleted.
+PRUNE_WRITES_SQL = """
+DELETE FROM checkpoint_writes
+WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id = ANY(%s)
+"""
+
+PRUNE_CHECKPOINTS_SQL = """
+WITH deleted AS (
+    DELETE FROM checkpoints
+    WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id <> ALL(%s)
+    RETURNING checkpoint -> 'channel_versions' AS versions
+)
+SELECT DISTINCT v.key AS channel, v.value AS version
+FROM deleted
+CROSS JOIN LATERAL jsonb_each_text(deleted.versions) AS v
+"""
+
+PRUNE_PARENTS_SQL = """
+UPDATE checkpoints SET parent_checkpoint_id = NULL
+WHERE thread_id = %s AND checkpoint_ns = %s
+AND checkpoint_id = ANY(%s) AND parent_checkpoint_id <> ALL(%s)
+"""
+
+PRUNE_REFERENCED_BLOBS_SQL = """
+SELECT DISTINCT v.key AS channel, v.value AS version
+FROM checkpoints c
+CROSS JOIN LATERAL jsonb_each_text(c.checkpoint -> 'channel_versions') AS v
+WHERE c.thread_id = %s AND c.checkpoint_ns = %s
+"""
+
+PRUNE_BLOBS_SQL = """
+DELETE FROM checkpoint_blobs
+WHERE thread_id = %s AND checkpoint_ns = %s AND channel = %s AND version = ANY(%s)
+"""
+
+PRUNE_SUBGRAPH_CALLS_SQL = (
+    "DELETE FROM checkpoints WHERE thread_id = %s AND checkpoint_ns = ANY(%s)",
+    "DELETE FROM checkpoint_writes WHERE thread_id = %s AND checkpoint_ns = ANY(%s)",
+    "DELETE FROM checkpoint_blobs WHERE thread_id = %s AND checkpoint_ns = ANY(%s)",
+)
+
+
+def _garbage_blobs(
+    freed: Sequence[Mapping[str, Any]], referenced: Sequence[Mapping[str, Any]]
+) -> dict[str, list[str]]:
+    """Blobs a deleted checkpoint used and no kept one does, by channel."""
+    kept = {(r["channel"], r["version"]) for r in referenced}
+    garbage: dict[str, list[str]] = {}
+    for r in freed:
+        if (r["channel"], r["version"]) not in kept:
+            garbage.setdefault(r["channel"], []).append(r["version"])
+    return garbage
+
+
+_Statements = Generator[tuple[str, Sequence[Any], bool], list[DictRow], None]
+
+
+def _refuse_delta_threads(thread_ids: Sequence[str]) -> _Statements:
+    rows = yield PRUNE_DELTA_THREADS_SQL, (thread_ids, thread_ids), True
+    if delta := sorted(r["thread_id"] for r in rows):
+        raise NotImplementedError(
+            f"keep_latest can't prune threads that use a DeltaChannel: {', '.join(delta)}"
+        )
+
+
+def _keep_latest(thread_id: str) -> _Statements:
+    """The statements of one thread's keep_latest transaction, for `_execute`
+    and `_aexecute` to run. Each comes with whether its rows are sent back.
+
+    They all read one snapshot, so what a run saves meanwhile is neither seen
+    nor deleted, whatever its checkpoint ID.
+    """
+    yield PRUNE_SNAPSHOT_SQL, (), False
+    heads = yield PRUNE_HEADS_SQL, (thread_id,) * 4, True
+    # Catches a thread that started using a DeltaChannel after the batch check.
+    yield from _refuse_delta_threads([thread_id])
+    kept = {
+        h["checkpoint_ns"]: [i for i in (h["head"], h["sends_parent"]) if i]
+        for h in heads
+    }
+    # A subgraph call's namespace has its task ID after a ":", unlike a
+    # `checkpointer=True` subgraph's, whose state outlives each call. A call
+    # goes with the root checkpoint it started from, but only once that one is
+    # saved: runs save in the background, so a call can commit first.
+    calls = {
+        h["checkpoint_ns"]
+        for h in heads
+        if ":" in h["checkpoint_ns"]
+        and h["root_saved"]
+        and h["root"] not in kept.get("", ())
+    }
+    if calls:
+        for sql in PRUNE_SUBGRAPH_CALLS_SQL:
+            yield sql, (thread_id, list(calls)), False
+    for ns, ids in kept.items():
+        if ns in calls:
+            continue
+        unkept = yield PRUNE_UNKEPT_SQL, (thread_id, ns, ids), True
+        if unkept:
+            deleted = [r["checkpoint_id"] for r in unkept]
+            yield PRUNE_WRITES_SQL, (thread_id, ns, deleted), False
+        freed = yield PRUNE_CHECKPOINTS_SQL, (thread_id, ns, ids), True
+        yield PRUNE_PARENTS_SQL, (thread_id, ns, ids, ids), False
+        if freed:
+            referenced = yield PRUNE_REFERENCED_BLOBS_SQL, (thread_id, ns), True
+            for channel, versions in _garbage_blobs(freed, referenced).items():
+                yield PRUNE_BLOBS_SQL, (thread_id, ns, channel, versions), False
+
+
+# Unprepared, so Postgres plans each statement knowing its arrays: a cached
+# generic plan checks a row against an `= ANY(...)` array one element at a time.
+def _execute(cur: Cursor[DictRow], statements: _Statements) -> None:
+    try:
+        sql, params, fetch = next(statements)
+        while True:
+            cur.execute(sql, params, prepare=False)
+            sql, params, fetch = statements.send(cur.fetchall() if fetch else [])
+    except StopIteration:
+        pass
+
+
+async def _aexecute(cur: AsyncCursor[DictRow], statements: _Statements) -> None:
+    try:
+        sql, params, fetch = next(statements)
+        while True:
+            await cur.execute(sql, params, prepare=False)
+            rows = await cur.fetchall() if fetch else []
+            sql, params, fetch = statements.send(rows)
+    except StopIteration:
+        pass
+
+
 class _DeltaStage2Row(TypedDict, total=False):
     """One row from `_build_delta_stage2_sql` (a UNION ALL of writes and blobs)."""
 

{langgraph_checkpoint_postgres-3.1.3 → langgraph_checkpoint_postgres-3.2.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "langgraph-checkpoint-postgres"
-version = "3.1.3"
+version = "3.2.0"
 description = "Library with a Postgres implementation of LangGraph checkpoint saver."
 authors = []
 requires-python = ">=3.10"

langgraph_checkpoint_postgres-3.2.0/tests/test_prune.py

@@ -0,0 +1,540 @@
+import asyncio
+from collections.abc import Awaitable, Callable
+from contextlib import asynccontextmanager
+from typing import Any
+from uuid import uuid4
+
+import pytest
+from langgraph.checkpoint.base import empty_checkpoint
+from langgraph.checkpoint.base.id import uuid6
+from langgraph.checkpoint.conformance import validate
+from langgraph.checkpoint.conformance.initializer import checkpointer_test
+from langgraph.checkpoint.serde.types import TASKS, _DeltaSnapshot
+from psycopg import AsyncConnection
+
+from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
+from langgraph.checkpoint.postgres.base import (
+    PRUNE_DELTA_THREADS_SQL,
+    PRUNE_HEADS_SQL,
+    PRUNE_REFERENCED_BLOBS_SQL,
+)
+from tests.conftest import DEFAULT_POSTGRES_URI, DEFAULT_URI
+from tests.test_async import _saver as _async_saver
+from tests.test_sync import _saver as _sync_saver
+
+SAVERS = ["base", "pool", "pipe", "aio_base", "aio_pool", "aio_pipe"]
+
+
+@asynccontextmanager
+async def _saver(name: str):
+    if name.startswith("aio_"):
+        async with _async_saver(name.removeprefix("aio_")) as saver:
+            yield saver
+    else:
+        with _sync_saver(name) as saver:
+            yield saver
+
+
+async def _call(saver: Any, method: str, *args: Any, **kwargs: Any) -> Any:
+    if isinstance(saver, AsyncPostgresSaver):
+        return await getattr(saver, f"a{method}")(*args, **kwargs)
+    return getattr(saver, method)(*args, **kwargs)
+
+
+async def _rows(saver: Any, thread_id: str) -> dict[str, list[tuple]]:
+    queries = {
+        "checkpoints": "SELECT checkpoint_ns, checkpoint_id FROM checkpoints"
+        " WHERE thread_id = %s ORDER BY 1, 2",
+        "writes": "SELECT checkpoint_ns, checkpoint_id, task_id, idx"
+        " FROM checkpoint_writes WHERE thread_id = %s ORDER BY 1, 2, 3, 4",
+        "blobs": "SELECT checkpoint_ns, channel, version FROM checkpoint_blobs"
+        " WHERE thread_id = %s ORDER BY 1, 2, 3",
+    }
+    rows = {}
+    for table, sql in queries.items():
+        if isinstance(saver, AsyncPostgresSaver):
+            async with saver._cursor() as cur:
+                await cur.execute(sql, (thread_id,))
+                rows[table] = [tuple(r.values()) for r in await cur.fetchall()]
+        else:
+            with saver._cursor() as cur:
+                cur.execute(sql, (thread_id,))
+                rows[table] = [tuple(r.values()) for r in cur.fetchall()]
+    return rows
+
+
+async def _put(
+    saver: Any,
+    thread_id: str,
+    parent: dict | None,
+    step: int,
+    values: dict[str, Any],
+    versions: dict[str, str],
+    *,
+    ns: str = "",
+    metadata: dict[str, Any] | None = None,
+    v: int = 4,
+    checkpoint_id: str | None = None,
+) -> dict:
+    config = {"configurable": {"thread_id": thread_id, "checkpoint_ns": ns}}
+    if parent is not None:
+        config["configurable"]["checkpoint_id"] = parent["configurable"][
+            "checkpoint_id"
+        ]
+    checkpoint = empty_checkpoint()
+    checkpoint.update(
+        v=v,
+        id=checkpoint_id or str(uuid6(clock_seq=step)),
+        channel_values=values,
+        channel_versions=versions,
+    )
+    new_versions = {k: versions[k] for k in values}
+    return await _call(
+        saver,
+        "put",
+        config,
+        checkpoint,
+        {"source": "loop", "step": step, "parents": {}, **(metadata or {})},
+        new_versions,
+    )
+
+
+async def _three_checkpoints(
+    saver: Any,
+    thread_id: str,
+    ns: str = "",
+    metadata: dict[str, Any] | None = None,
+) -> dict:
+    parent = None
+    steps = [
+        ({"a": [1], "s": ["shared"]}, {"a": "1", "s": "1"}),
+        ({"a": [1, 2]}, {"a": "2", "s": "1"}),
+        ({"a": [1, 2, 3]}, {"a": "3", "s": "1"}),
+    ]
+    for step, (values, versions) in enumerate(steps):
+        parent = await _put(
+            saver, thread_id, parent, step, values, versions, ns=ns, metadata=metadata
+        )
+        await _call(saver, "put_writes", parent, [("a", step)], str(uuid4()))
+    return parent
+
+
+@asynccontextmanager
+async def _second_saver(saver: AsyncPostgresSaver):
+    assert isinstance(saver.conn, AsyncConnection)
+    uri = DEFAULT_POSTGRES_URI + saver.conn.info.dbname
+    async with AsyncPostgresSaver.from_conn_string(uri) as other:
+        yield other
+
+
+def _run_once_after(
+    saver: AsyncPostgresSaver,
+    sql: str,
+    action: Callable[[], Awaitable[Any]],
+    *,
+    nth: int = 1,
+) -> None:
+    seen = 0
+    cursor = saver._cursor
+
+    class Hooked:
+        def __init__(self, cur: Any) -> None:
+            self.cur = cur
+
+        def __getattr__(self, name: str) -> Any:
+            return getattr(self.cur, name)
+
+        async def execute(self, query: Any, params: Any = None, **kwargs: Any) -> None:
+            nonlocal seen
+            await self.cur.execute(query, params, **kwargs)
+            if query is sql:
+                seen += 1
+                if seen == nth:
+                    await action()
+
+    @asynccontextmanager
+    async def hooked(**kwargs: Any):
+        async with cursor(**kwargs) as cur:
+            yield Hooked(cur)
+
+    saver._cursor = hooked
+
+
+async def _delta_thread(saver: Any, thread_id: str, evidence: str) -> None:
+    if evidence == "counters":
+        first = await _put(saver, thread_id, None, 0, {"m": [1]}, {"m": "1"})
+        await _put(
+            saver,
+            thread_id,
+            first,
+            1,
+            {},
+            {"m": "2"},
+            metadata={"counters_since_delta_snapshot": {"m": [1, 1]}},
+        )
+    elif evidence == "snapshot":
+        first = await _put(
+            saver, thread_id, None, 0, {"m": _DeltaSnapshot([1])}, {"m": "1"}
+        )
+        await _put(saver, thread_id, first, 1, {}, {"m": "2"})
+    else:
+        marker = {"delta_writes_versioned": True}
+        first = await _put(
+            saver, thread_id, None, 0, {"m": [1]}, {"m": "1"}, metadata=marker
+        )
+        await _put(
+            saver, thread_id, first, 1, {"m": [1, 2]}, {"m": "2"}, metadata=marker
+        )
+
+
+@pytest.mark.parametrize("name", SAVERS)
+async def test_keep_latest_keeps_each_namespace_head_with_its_writes_and_blobs(
+    name: str,
+) -> None:
+    async with _saver(name) as saver:
+        thread_id, other_id = str(uuid4()), str(uuid4())
+        heads = {
+            ns: await _three_checkpoints(saver, thread_id, ns) for ns in ("", "sub:1")
+        }
+        await _three_checkpoints(saver, other_id)
+        before = {ns: await _call(saver, "get_tuple", h) for ns, h in heads.items()}
+        other_before = await _rows(saver, other_id)
+
+        await _call(saver, "prune", [thread_id])
+
+        rows = await _rows(saver, thread_id)
+        assert rows["checkpoints"] == [
+            (ns, heads[ns]["configurable"]["checkpoint_id"]) for ns in ("", "sub:1")
+        ]
+        assert [(ns, cid) for ns, cid, *_ in rows["writes"]] == rows["checkpoints"]
+        assert rows["blobs"] == [
+            (ns, ch, ver)
+            for ns in ("", "sub:1")
+            for ch, ver in (("a", "3"), ("s", "1"))
+        ]
+        for ns, head in heads.items():
+            assert await _call(saver, "get_tuple", head) == before[ns]._replace(
+                parent_config=None
+            )
+        assert await _rows(saver, other_id) == other_before
+
+
+@pytest.mark.parametrize("evidence", ["counters", "snapshot", "versioned"])
+@pytest.mark.parametrize("name", SAVERS)
+async def test_keep_latest_refuses_a_batch_with_a_delta_channel_thread(
+    name: str, evidence: str
+) -> None:
+    async with _saver(name) as saver:
+        ordinary_id, delta_id = str(uuid4()), str(uuid4())
+        await _three_checkpoints(saver, ordinary_id)
+        await _delta_thread(saver, delta_id, evidence)
+        before = [await _rows(saver, t) for t in (ordinary_id, delta_id)]
+
+        with pytest.raises(NotImplementedError, match=delta_id):
+            await _call(saver, "prune", [ordinary_id, delta_id])
+
+        assert [await _rows(saver, t) for t in (ordinary_id, delta_id)] == before
+
+
+@pytest.mark.parametrize("name", SAVERS)
+async def test_keep_latest_prunes_a_thread_storing_a_plain_true(name: str) -> None:
+    async with _saver(name) as saver:
+        thread_id = str(uuid4())
+        first = await _put(
+            saver, thread_id, None, 0, {"flag": True, "a": [1]}, {"flag": "1", "a": "1"}
+        )
+        head = await _put(
+            saver, thread_id, first, 1, {"flag": True}, {"flag": "2", "a": "1"}
+        )
+
+        await _call(saver, "prune", [thread_id])
+
+        assert (await _rows(saver, thread_id))["checkpoints"] == [
+            ("", head["configurable"]["checkpoint_id"])
+        ]
+
+
+@pytest.mark.parametrize("name", SAVERS)
+async def test_keep_latest_keeps_the_parent_holding_a_pre_v4_heads_sends(
+    name: str,
+) -> None:
+    async with _saver(name) as saver:
+        thread_id = str(uuid4())
+        first = await _put(saver, thread_id, None, 0, {}, {}, v=1)
+        await _call(saver, "put_writes", first, [(TASKS, "send-1")], str(uuid4()))
+        second = await _put(saver, thread_id, first, 1, {}, {}, v=1)
+        head = await _put(saver, thread_id, second, 2, {}, {}, v=1)
+        await _call(saver, "put_writes", second, [(TASKS, "send-2")], str(uuid4()))
+
+        await _call(saver, "prune", [thread_id])
+
+        assert (await _rows(saver, thread_id))["checkpoints"] == [
+            ("", second["configurable"]["checkpoint_id"]),
+            ("", head["configurable"]["checkpoint_id"]),
+        ]
+        loaded = await _call(saver, "get_tuple", head)
+        assert loaded.checkpoint["channel_values"][TASKS] == ["send-2"]
+        assert loaded.parent_config == second
+        assert (await _call(saver, "get_tuple", second)).parent_config is None
+
+
+@pytest.mark.parametrize("name", SAVERS)
+async def test_delete_removes_every_row_of_the_threads(name: str) -> None:
+    async with _saver(name) as saver:
+        ordinary_id, delta_id, other_id = str(uuid4()), str(uuid4()), str(uuid4())
+        await _three_checkpoints(saver, ordinary_id)
+        await _delta_thread(saver, delta_id, "counters")
+        await _three_checkpoints(saver, other_id)
+        other_before = await _rows(saver, other_id)
+
+        await _call(saver, "prune", [ordinary_id, delta_id], strategy="delete")
+
+        empty = {"checkpoints": [], "writes": [], "blobs": []}
+        assert await _rows(saver, ordinary_id) == empty
+        assert await _rows(saver, delta_id) == empty
+        assert await _rows(saver, other_id) == other_before
+
+
+@pytest.mark.parametrize("name", SAVERS)
+async def test_unknown_strategy_changes_nothing(name: str) -> None:
+    async with _saver(name) as saver:
+        thread_id = str(uuid4())
+        await _three_checkpoints(saver, thread_id)
+        before = await _rows(saver, thread_id)
+
+        with pytest.raises(ValueError, match="keep_all"):
+            await _call(saver, "prune", [thread_id], strategy="keep_all")
+
+        assert await _rows(saver, thread_id) == before
+
+
+async def test_async_saver_prune_runs_from_another_thread_only() -> None:
+    async with _saver("aio_base") as saver:
+        thread_id = str(uuid4())
+        head = await _three_checkpoints(saver, thread_id)
+
+        with pytest.raises(asyncio.InvalidStateError):
+            saver.prune([thread_id])
+        await asyncio.to_thread(saver.prune, [thread_id])
+
+        assert (await _rows(saver, thread_id))["checkpoints"] == [
+            ("", head["configurable"]["checkpoint_id"])
+        ]
+
+
+async def test_prune_conformance() -> None:
+    @checkpointer_test(name="AsyncPostgresSaver")
+    async def postgres_saver():
+        async with AsyncPostgresSaver.from_conn_string(DEFAULT_URI) as saver:
+            await saver.setup()
+            yield saver
+
+    report = await validate(postgres_saver, capabilities={"prune"})
+
+    prune = report.results["prune"]
+    if not prune.passed:
+        details = "\n".join(prune.failures)
+        pytest.fail(f"Capability prune failed (passed={prune.passed}):\n{details}")
+
+
+@pytest.mark.parametrize("strategy", ["keep_latest", "delete"])
+@pytest.mark.parametrize("name", ["base", "aio_base"])
+async def test_prune_rejects_a_single_thread_id_string(
+    name: str, strategy: str
+) -> None:
+    async with _saver(name) as saver:
+        for thread_id in ("ab", "a", "b"):
+            await _three_checkpoints(saver, thread_id)
+        before = [await _rows(saver, t) for t in ("ab", "a", "b")]
+
+        with pytest.raises(TypeError, match="not a str"):
+            await _call(saver, "prune", "ab", strategy=strategy)
+
+        assert [await _rows(saver, t) for t in ("ab", "a", "b")] == before
+
+
+@pytest.mark.parametrize(
+    "after", [PRUNE_HEADS_SQL, PRUNE_REFERENCED_BLOBS_SQL], ids=["heads", "blobs"]
+)
+async def test_keep_latest_keeps_a_checkpoint_saved_while_it_runs(after: str) -> None:
+    async with _saver("aio_base") as saver, _second_saver(saver) as other:
+        thread_id = str(uuid4())
+        head = await _three_checkpoints(saver, thread_id)
+        saved: list[dict] = []
+
+        async def save_a_child() -> None:
+            saved.append(
+                await _put(
+                    other, thread_id, head, 3, {"a": [1, 2, 3, 4]}, {"a": "4", "s": "1"}
+                )
+            )
+
+        _run_once_after(saver, after, save_a_child)
+        await saver.aprune([thread_id])
+
+        loaded = await saver.aget_tuple(saved[0])
+        assert loaded is not None, "the checkpoint saved during the prune was deleted"
+        assert loaded.checkpoint["channel_values"] == {
+            "a": [1, 2, 3, 4],
+            "s": ["shared"],
+        }
+        assert loaded.parent_config == head
+
+
+async def test_keep_latest_refuses_a_thread_that_turns_delta_after_the_batch_check() -> (
+    None
+):
+    async with _saver("aio_base") as saver, _second_saver(saver) as other:
+        thread_id = str(uuid4())
+        head = await _three_checkpoints(saver, thread_id)
+
+        async def save_a_delta_checkpoint() -> None:
+            await _put(
+                other,
+                thread_id,
+                head,
+                3,
+                {},
+                {"a": "4", "s": "1"},
+                metadata={"counters_since_delta_snapshot": {"a": [1, 1]}},
+            )
+
+        _run_once_after(saver, PRUNE_DELTA_THREADS_SQL, save_a_delta_checkpoint)
+        with pytest.raises(NotImplementedError, match=thread_id):
+            await saver.aprune([thread_id])
+
+        assert len((await _rows(saver, thread_id))["checkpoints"]) == 4
+
+
+async def test_keep_latest_keeps_the_history_of_a_delta_checkpoint_saved_after_the_thread_check() -> (
+    None
+):
+    async with _saver("aio_base") as saver, _second_saver(saver) as other:
+        thread_id = str(uuid4())
+        head = await _three_checkpoints(saver, thread_id)
+        saved: list[dict] = []
+
+        async def save_a_delta_checkpoint() -> None:
+            saved.append(
+                await _put(
+                    other,
+                    thread_id,
+                    head,
+                    3,
+                    {},
+                    {"a": "4", "s": "1"},
+                    metadata={"counters_since_delta_snapshot": {"a": [1, 1]}},
+                )
+            )
+
+        _run_once_after(saver, PRUNE_DELTA_THREADS_SQL, save_a_delta_checkpoint, nth=2)
+        await saver.aprune([thread_id])
+
+        assert (await _rows(saver, thread_id))["checkpoints"] == [
+            ("", head["configurable"]["checkpoint_id"]),
+            ("", saved[0]["configurable"]["checkpoint_id"]),
+        ]
+        history = await saver.aget_delta_channel_history(
+            config=saved[0], channels=["a"]
+        )
+        assert history["a"]["seed"] == [1, 2, 3]
+        assert [value for _, _, value in history["a"]["writes"]] == [2]
+
+
+@pytest.mark.parametrize("name", SAVERS)
+async def test_keep_latest_deletes_a_subgraph_call_with_the_checkpoint_it_started_from(
+    name: str,
+) -> None:
+    async with _saver(name) as saver:
+        thread_id = str(uuid4())
+        older = await _put(saver, thread_id, None, 0, {}, {})
+        head = await _put(saver, thread_id, older, 1, {}, {})
+        unsaved = str(uuid6(clock_seq=2))
+        started_from = {
+            "finished:1": older["configurable"]["checkpoint_id"],
+            "interrupted:2": head["configurable"]["checkpoint_id"],
+            "running:3": unsaved,
+            "checkpointer_true": older["configurable"]["checkpoint_id"],
+        }
+        for ns, root in started_from.items():
+            await _three_checkpoints(
+                saver, thread_id, ns, metadata={"parents": {"": root}}
+            )
+
+        await _call(saver, "prune", [thread_id])
+
+        rows = await _rows(saver, thread_id)
+        assert {ns for ns, *_ in rows["checkpoints"]} == {
+            "",
+            "interrupted:2",
+            "running:3",
+            "checkpointer_true",
+        }
+        assert {ns for ns, *_ in rows["writes"]} == {
+            "interrupted:2",
+            "running:3",
+            "checkpointer_true",
+        }
+        assert {ns for ns, *_ in rows["blobs"]} == {
+            "interrupted:2",
+            "running:3",
+            "checkpointer_true",
+        }
+
+
+async def test_keep_latest_keeps_a_checkpoint_saved_meanwhile_whatever_its_id() -> None:
+    async with _saver("aio_base") as saver, _second_saver(saver) as other:
+        thread_id = str(uuid4())
+        head = await _three_checkpoints(saver, thread_id)
+        lower_id = "0" + head["configurable"]["checkpoint_id"][1:]
+        saved: list[dict] = []
+
+        async def save_a_child_with_a_lower_id() -> None:
+            child = await _put(
+                other,
+                thread_id,
+                head,
+                3,
+                {"a": [1, 2, 3, 4]},
+                {"a": "4", "s": "1"},
+                checkpoint_id=lower_id,
+            )
+            await _call(other, "put_writes", child, [("a", 4)], str(uuid4()))
+            saved.append(child)
+
+        _run_once_after(saver, PRUNE_HEADS_SQL, save_a_child_with_a_lower_id)
+        await saver.aprune([thread_id])
+
+        loaded = await saver.aget_tuple(saved[0])
+        assert loaded is not None, "the checkpoint saved during the prune was deleted"
+        assert loaded.checkpoint["channel_values"] == {
+            "a": [1, 2, 3, 4],
+            "s": ["shared"],
+        }
+        assert [value for _, _, value in loaded.pending_writes or []] == [4]
+
+
+@pytest.mark.parametrize("name", SAVERS)
+async def test_keep_latest_keeps_writes_saved_before_their_checkpoint(
+    name: str,
+) -> None:
+    async with _saver(name) as saver:
+        thread_id = str(uuid4())
+        head = await _three_checkpoints(saver, thread_id)
+        next_id = str(uuid6(clock_seq=3))
+        next_config = {
+            "configurable": {
+                "thread_id": thread_id,
+                "checkpoint_ns": "",
+                "checkpoint_id": next_id,
+            }
+        }
+        await _call(saver, "put_writes", next_config, [("a", 4)], str(uuid4()))
+
+        await _call(saver, "prune", [thread_id])
+        await _put(
+            saver, thread_id, head, 3, {}, {"a": "3", "s": "1"}, checkpoint_id=next_id
+        )
+
+        loaded = await _call(saver, "get_tuple", next_config)
+        assert [value for _, _, value in loaded.pending_writes] == [4]

{langgraph_checkpoint_postgres-3.1.3 → langgraph_checkpoint_postgres-3.2.0}/uv.lock

@@ -353,7 +353,7 @@ test = [
 
 [[package]]
 name = "langgraph-checkpoint-postgres"
-version = "3.1.3"
+version = "3.2.0"
 source = { editable = "." }
 dependencies = [
     { name = "langgraph-checkpoint" },