From 36bf3887f074f2b05f42409c8a67e3a33ed1d1ba Mon Sep 17 00:00:00 2001 From: Elior Nataf Lackritz Date: Wed, 14 Oct 2026 15:09:51 -0400 Subject: [PATCH] feat(checkpoint-postgres): add prune to PostgresSaver and AsyncPostgresSaver (#9290) Closes #8531. `PostgresSaver` and `AsyncPostgresSaver` inherited `prune` / `aprune` from the base class, which raises `NotImplementedError`, so long-lived threads had no supported way to drop old checkpoints. This adds both strategies. Thanks to @wangjiawei-vegetable for proposing this scope in #8531 and #9247. - **`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), with its pending writes and the blobs it references, and deletes the rest, including blobs only the deleted checkpoints used. A kept checkpoint whose parent is deleted is left with no parent. A subgraph call's namespace goes with the root checkpoint the call started from; a `checkpointer=True` subgraph keeps its latest checkpoint, since its state outlives each call. - **Each thread is pruned on one snapshot.** Its transaction runs at `REPEATABLE READ` and deletes everything that snapshot sees except the kept checkpoints, so a run appending to the thread at the same time is safe, whatever its checkpoint IDs. - **`keep_latest` refuses threads that use a `DeltaChannel`** (`NotImplementedError` naming them, nothing pruned). A `DeltaChannel` rebuilds its latest value from older checkpoints, so deleting them would empty it. The saver can't see channel types, so it detects one by what only a `DeltaChannel` leaves: `counters_since_delta_snapshot` in metadata, a snapshot (an inline `true` with a blob behind it), or `delta_writes_versioned`, which LangGraph.js sets on threads its `DeltaChannel` graphs start. The check runs for the whole batch first and again in each thread's transaction, so a thread that starts using a `DeltaChannel` mid-batch is refused at its turn. Unlike the `prune` in #6883, `keep_latest` here deletes the blobs older checkpoints no longer share with the kept ones (most of the storage on a long thread), and guards `DeltaChannel` threads. Measured locally, with or without planner statistics: a 6,000-checkpoint thread prunes in 100 to 200 ms, and one with 1,500 finished subgraph calls in about 110 ms. ## Known limitations - Two prunes of the same thread at once can fail one of them with a serialization failure or a deadlock; retrying is safe. - Forking or replaying from a checkpoint while it's being pruned is not safe. - Threads written by older releases can escape the `DeltaChannel` check: before #9142 (Python 1.2.13) and langchain-ai/langgraphjs#2973 (JS 1.5.0), some `update_state` paths saved checkpoints with neither counters nor a snapshot. - Postgres reuses the freed space but doesn't shrink the table files; that takes `VACUUM FULL` or `pg_repack`. - `DeltaChannel`-aware pruning (keeping ancestors back to the last snapshot) is a follow-up. ## Test plan - [x] `libs/checkpoint-postgres/tests/test_prune.py`: 71 tests over the sync and async savers (single connection, pool, pipeline), including the conformance `prune` capability. Passes on Postgres 15 and 16. - [x] `libs/langgraph/tests/test_prune_postgres.py`: 39 graph tests. A pruned thread loads the same state and keeps running like an unpruned twin (sync, async and exit durability, interrupts, a failed parallel step, interrupted and `checkpointer=True` subgraphs), plus deterministic races: a turn saved mid-prune, and threads switching to a `DeltaChannel` before and after their own check. - [x] Each guard (the snapshot, the writes and blob scopes, the subgraph call rules, the parent reset, the `DeltaChannel` checks, the single-string check, the pre-v4 pending sends, the pipeline sync) broken one at a time: every break fails at least one test. - [x] A thread pruned in a loop while it runs, in each durability mode, ends identical to its twin; threads written by LangGraph.js 1.5.0 and pruned from Python behave like their twins. --- .../langgraph/checkpoint/postgres/__init__.py | 72 +++ .../langgraph/checkpoint/postgres/aio.py | 81 +++ .../langgraph/checkpoint/postgres/base.py | 216 ++++++- libs/checkpoint-postgres/tests/test_prune.py | 540 ++++++++++++++++++ libs/langgraph/tests/test_prune_postgres.py | 467 +++++++++++++++ 5 files changed, 1375 insertions(+), 1 deletion(-) create mode 100644 libs/checkpoint-postgres/tests/test_prune.py create mode 100644 libs/langgraph/tests/test_prune_postgres.py diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py index 5e61b50ed..d1e3cac59 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py +++ b/libs/checkpoint-postgres/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. diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py index 5a00c99b1..716e54ed8 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py +++ b/libs/checkpoint-postgres/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]: diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py index 81b618c03..1f048d47e 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py +++ b/libs/checkpoint-postgres/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).""" diff --git a/libs/checkpoint-postgres/tests/test_prune.py b/libs/checkpoint-postgres/tests/test_prune.py new file mode 100644 index 000000000..187abd21f --- /dev/null +++ b/libs/checkpoint-postgres/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] diff --git a/libs/langgraph/tests/test_prune_postgres.py b/libs/langgraph/tests/test_prune_postgres.py new file mode 100644 index 000000000..d7aa00ea2 --- /dev/null +++ b/libs/langgraph/tests/test_prune_postgres.py @@ -0,0 +1,467 @@ +import operator +import sys +from collections.abc import Callable +from contextlib import contextmanager +from typing import Annotated, Any + +import pytest +from langchain_core.messages import AIMessage, AnyMessage, HumanMessage +from langchain_core.runnables import RunnableConfig +from langgraph.checkpoint.postgres import PostgresSaver +from langgraph.checkpoint.postgres.base import PRUNE_DELTA_THREADS_SQL, PRUNE_HEADS_SQL +from typing_extensions import TypedDict + +from langgraph.channels.delta import DeltaChannel +from langgraph.graph import START, StateGraph +from langgraph.graph.message import _messages_delta_reducer, add_messages +from langgraph.types import Command, StateSnapshot, interrupt +from tests.conftest import NO_DOCKER +from tests.conftest_checkpointer import ( + DEFAULT_POSTGRES_URI, + _checkpointer_postgres, + _checkpointer_postgres_aio, + _checkpointer_postgres_aio_pipe, + _checkpointer_postgres_aio_pool, + _checkpointer_postgres_pipe, + _checkpointer_postgres_pool, +) + +pytestmark = [ + pytest.mark.anyio, + pytest.mark.skipif(NO_DOCKER, reason="needs Postgres"), +] + +NEEDS_CONTEXTVARS = pytest.mark.skipif( + sys.version_info < (3, 11), + reason="Python 3.11+ is required for async contextvars support", +) + +SYNC = [ + _checkpointer_postgres, + _checkpointer_postgres_pipe, + _checkpointer_postgres_pool, +] +ASYNC = [ + _checkpointer_postgres_aio, + _checkpointer_postgres_aio_pipe, + _checkpointer_postgres_aio_pool, +] + + +class State(TypedDict): + messages: Annotated[list[AnyMessage], add_messages] + notes: Annotated[list[str], operator.add] + + +class DeltaState(TypedDict): + messages: Annotated[list[AnyMessage], DeltaChannel(_messages_delta_reducer)] + + +def _reply(ask: bool): + def reply(state: dict) -> dict: + n = len(state["messages"]) + answer = interrupt("approve?") if ask else "ok" + return {"messages": [AIMessage(f"{answer} {n}", id=f"reply-{n}")]} + + return reply + + +def _graph( + checkpointer: Any, + *, + ask: bool = False, + fail: set[str] | None = None, + calls: dict[str, int] | None = None, +) -> Any: + def note(name: str): + def node(state: State) -> dict: + if calls is not None: + calls[name] = calls.get(name, 0) + 1 + if fail is not None and name in fail: + fail.discard(name) + raise ValueError(f"{name} failed") + return {"notes": [name]} + + return node + + builder = StateGraph(State) + builder.add_node("left", note("left")) + builder.add_node("right", note("right")) + builder.add_node("reply", _reply(ask)) + builder.add_edge(START, "left") + builder.add_edge(START, "right") + builder.add_edge(["left", "right"], "reply") + return builder.compile(checkpointer=checkpointer) + + +def _counting_subgraph(name: str, *, checkpointer: bool | None) -> Any: + class Calls(TypedDict): + messages: Annotated[list[AnyMessage], add_messages] + calls: Annotated[list[str], operator.add] + + def count(state: Calls) -> dict: + n = len(state["messages"]) + reply = AIMessage(f"{name} call {len(state['calls']) + 1}", id=f"{name}-{n}") + return {"calls": [name], "messages": [reply]} + + builder = StateGraph(Calls) + builder.add_node("count", count) + builder.add_edge(START, "count") + return builder.compile(checkpointer=checkpointer) + + +def _delta_graph(checkpointer: Any) -> Any: + builder = StateGraph(DeltaState) + builder.add_node("reply", _reply(ask=False)) + builder.add_edge(START, "reply") + return builder.compile(checkpointer=checkpointer) + + +def _config(thread_id: str) -> RunnableConfig: + return {"configurable": {"thread_id": thread_id}} + + +def _say(i: int) -> dict: + return {"messages": [HumanMessage(f"hi {i}", id=f"hi-{i}")]} + + +def _checkpoint_ids(checkpointer: Any, thread_id: str) -> list[str]: + return [c.checkpoint["id"] for c in checkpointer.list(_config(thread_id))] + + +def _namespaces(checkpointer: Any, thread_id: str) -> list[str]: + return sorted( + c.config["configurable"]["checkpoint_ns"] + for c in checkpointer.list(_config(thread_id)) + ) + + +async def _acheckpoint_ids(checkpointer: Any, thread_id: str) -> list[str]: + return [c.checkpoint["id"] async for c in checkpointer.alist(_config(thread_id))] + + +def _parentless(snapshot: StateSnapshot) -> StateSnapshot: + return snapshot._replace( + parent_config=None, + tasks=tuple( + t._replace(state=_parentless(t.state)) + if isinstance(t.state, StateSnapshot) + else t + for t in snapshot.tasks + ), + ) + + +def _run_once_after( + saver: PostgresSaver, sql: str, action: Callable[[], 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) + + def execute(self, query: Any, params: Any = None, **kwargs: Any) -> None: + nonlocal seen + self.cur.execute(query, params, **kwargs) + if query is sql: + seen += 1 + if seen == nth: + action() + + @contextmanager + def hooked(**kwargs: Any): + with cursor(**kwargs) as cur: + yield Hooked(cur) + + saver._cursor = hooked + + +@pytest.mark.parametrize("durability", ["sync", "async", "exit"]) +@pytest.mark.parametrize("make_checkpointer", SYNC) +def test_keep_latest_keeps_the_state_and_the_thread_runs_on_like_its_twin( + make_checkpointer: Any, durability: str +) -> None: + with make_checkpointer() as checkpointer: + graph = _graph(checkpointer) + for i in range(3): + for thread_id in ("pruned", "twin"): + graph.invoke(_say(i), _config(thread_id), durability=durability) + before = graph.get_state(_config("pruned")) + + checkpointer.prune(["pruned"]) + + assert _checkpoint_ids(checkpointer, "pruned") == [ + before.config["configurable"]["checkpoint_id"] + ] + assert graph.get_state(_config("pruned")) == _parentless(before) + for thread_id in ("pruned", "twin"): + graph.invoke(_say(3), _config(thread_id), durability=durability) + assert ( + graph.get_state(_config("pruned")).values + == graph.get_state(_config("twin")).values + ) + + +@pytest.mark.parametrize("make_checkpointer", SYNC) +def test_keep_latest_keeps_a_pending_interrupt_resumable( + make_checkpointer: Any, +) -> None: + with make_checkpointer() as checkpointer: + graph = _graph(checkpointer, ask=True) + for thread_id in ("pruned", "twin"): + for i in range(2): + graph.invoke(_say(i), _config(thread_id)) + graph.invoke(Command(resume=f"yes {i}"), _config(thread_id)) + graph.invoke(_say(2), _config(thread_id)) + before = graph.get_state(_config("pruned")) + assert before.interrupts + + checkpointer.prune(["pruned"]) + + assert graph.get_state(_config("pruned")) == _parentless(before) + for thread_id in ("pruned", "twin"): + graph.invoke(Command(resume="yes 2"), _config(thread_id)) + assert ( + graph.get_state(_config("pruned")).values + == graph.get_state(_config("twin")).values + ) + + +@pytest.mark.parametrize("make_checkpointer", SYNC) +def test_keep_latest_keeps_writes_of_tasks_that_finished_before_a_failure( + make_checkpointer: Any, +) -> None: + with make_checkpointer() as checkpointer: + calls: dict[str, int] = {} + fail: set[str] = set() + graph = _graph(checkpointer, fail=fail, calls=calls) + for thread_id in ("pruned", "twin"): + for i in range(2): + graph.invoke(_say(i), _config(thread_id)) + fail.add("right") + with pytest.raises(ValueError, match="right failed"): + graph.invoke(_say(2), _config(thread_id)) + before = graph.get_state(_config("pruned")) + + checkpointer.prune(["pruned"]) + + assert graph.get_state(_config("pruned")) == _parentless(before) + left_runs = calls["left"] + graph.invoke(None, _config("pruned")) + assert calls["left"] == left_runs, "a task that had finished ran again" + graph.invoke(None, _config("twin")) + assert ( + graph.get_state(_config("pruned")).values + == graph.get_state(_config("twin")).values + ) + + +@pytest.mark.parametrize("make_checkpointer", SYNC) +def test_keep_latest_keeps_an_interrupted_subgraph_resumable( + make_checkpointer: Any, +) -> None: + with make_checkpointer() as checkpointer: + sub = StateGraph(State) + sub.add_node("reply", _reply(ask=True)) + sub.add_edge(START, "reply") + parent = StateGraph(State) + parent.add_node("sub", sub.compile()) + parent.add_edge(START, "sub") + graph = parent.compile(checkpointer=checkpointer) + for thread_id in ("pruned", "twin"): + for i in range(2): + graph.invoke(_say(i), _config(thread_id)) + graph.invoke(Command(resume=f"yes {i}"), _config(thread_id)) + graph.invoke(_say(2), _config(thread_id)) + before = graph.get_state(_config("pruned"), subgraphs=True) + interrupted = before.tasks[0].state.config["configurable"]["checkpoint_ns"] + + checkpointer.prune(["pruned"]) + + assert _namespaces(checkpointer, "pruned") == ["", interrupted] + assert graph.get_state(_config("pruned"), subgraphs=True) == _parentless(before) + for thread_id in ("pruned", "twin"): + graph.invoke(Command(resume="yes 2"), _config(thread_id)) + assert ( + graph.get_state(_config("pruned")).values + == graph.get_state(_config("twin")).values + ) + + +@pytest.mark.parametrize("durability", ["sync", "async", "exit"]) +@pytest.mark.parametrize("make_checkpointer", SYNC) +def test_keep_latest_keeps_a_checkpointer_true_subgraph_but_not_finished_calls( + make_checkpointer: Any, durability: str +) -> None: + with make_checkpointer() as checkpointer: + parent = StateGraph(State) + parent.add_node("memory", _counting_subgraph("memory", checkpointer=True)) + parent.add_node("call", _counting_subgraph("call", checkpointer=None)) + parent.add_edge(START, "memory") + parent.add_edge("memory", "call") + graph = parent.compile(checkpointer=checkpointer) + for i in range(3): + for thread_id in ("pruned", "twin"): + graph.invoke(_say(i), _config(thread_id), durability=durability) + + checkpointer.prune(["pruned"]) + + assert _namespaces(checkpointer, "pruned") == ["", "memory"] + for thread_id in ("pruned", "twin"): + graph.invoke(_say(3), _config(thread_id), durability=durability) + assert ( + graph.get_state(_config("pruned")).values + == graph.get_state(_config("twin")).values + ) + + +def test_keep_latest_keeps_a_turn_saved_after_it_read_the_heads() -> None: + with ( + _checkpointer_postgres() as checkpointer, + PostgresSaver.from_conn_string( + DEFAULT_POSTGRES_URI + checkpointer.conn.info.dbname + ) as pruner, + ): + graph = _graph(checkpointer) + for i in range(3): + for thread_id in ("pruned", "twin"): + graph.invoke(_say(i), _config(thread_id)) + _run_once_after( + pruner, PRUNE_HEADS_SQL, lambda: graph.invoke(_say(3), _config("pruned")) + ) + + pruner.prune(["pruned"]) + + graph.invoke(_say(3), _config("twin")) + assert ( + graph.get_state(_config("pruned")).values + == graph.get_state(_config("twin")).values + ) + + +def test_keep_latest_refuses_a_thread_that_switches_to_a_delta_channel_mid_batch() -> ( + None +): + with ( + _checkpointer_postgres() as checkpointer, + PostgresSaver.from_conn_string( + DEFAULT_POSTGRES_URI + checkpointer.conn.info.dbname + ) as pruner, + ): + graph, delta_graph = _graph(checkpointer), _delta_graph(checkpointer) + for i in range(3): + for thread_id in ("switched", "twin"): + graph.invoke(_say(i), _config(thread_id)) + _run_once_after( + pruner, + PRUNE_DELTA_THREADS_SQL, + lambda: delta_graph.invoke(_say(3), _config("switched")), + ) + + with pytest.raises(NotImplementedError, match="switched"): + pruner.prune(["switched"]) + + delta_graph.invoke(_say(3), _config("twin")) + assert ( + delta_graph.get_state(_config("switched")).values + == delta_graph.get_state(_config("twin")).values + ) + + +def test_keep_latest_keeps_a_delta_channel_turn_saved_after_the_thread_check() -> None: + with ( + _checkpointer_postgres() as checkpointer, + PostgresSaver.from_conn_string( + DEFAULT_POSTGRES_URI + checkpointer.conn.info.dbname + ) as pruner, + ): + graph, delta_graph = _graph(checkpointer), _delta_graph(checkpointer) + for i in range(3): + for thread_id in ("switched", "twin"): + graph.invoke(_say(i), _config(thread_id)) + head = _checkpoint_ids(checkpointer, "switched")[0] + _run_once_after( + pruner, + PRUNE_DELTA_THREADS_SQL, + lambda: delta_graph.invoke(_say(3), _config("switched")), + nth=2, + ) + + pruner.prune(["switched"]) + + assert _checkpoint_ids(checkpointer, "switched")[-1] == head + delta_graph.invoke(_say(3), _config("twin")) + assert ( + delta_graph.get_state(_config("switched")).values + == delta_graph.get_state(_config("twin")).values + ) + + +@pytest.mark.parametrize("make_checkpointer", SYNC) +def test_keep_latest_refuses_a_batch_with_a_delta_channel_thread( + make_checkpointer: Any, +) -> None: + with make_checkpointer() as checkpointer: + graph, delta_graph = _graph(checkpointer), _delta_graph(checkpointer) + for i in range(3): + graph.invoke(_say(i), _config("ordinary")) + delta_graph.invoke(_say(i), _config("delta")) + before = {t: _checkpoint_ids(checkpointer, t) for t in ("ordinary", "delta")} + delta_before = delta_graph.get_state(_config("delta")) + + with pytest.raises(NotImplementedError, match="delta"): + checkpointer.prune(["ordinary", "delta"]) + + assert {t: _checkpoint_ids(checkpointer, t) for t in before} == before + assert delta_graph.get_state(_config("delta")) == delta_before + + +@NEEDS_CONTEXTVARS +@pytest.mark.parametrize("make_checkpointer", ASYNC) +async def test_akeep_latest_keeps_a_pending_interrupt_resumable( + make_checkpointer: Any, +) -> None: + async with make_checkpointer() as checkpointer: + graph = _graph(checkpointer, ask=True) + for thread_id in ("pruned", "twin"): + for i in range(2): + await graph.ainvoke(_say(i), _config(thread_id)) + await graph.ainvoke(Command(resume=f"yes {i}"), _config(thread_id)) + await graph.ainvoke(_say(2), _config(thread_id)) + before = await graph.aget_state(_config("pruned")) + + await checkpointer.aprune(["pruned"]) + + assert await _acheckpoint_ids(checkpointer, "pruned") == [ + before.config["configurable"]["checkpoint_id"] + ] + assert await graph.aget_state(_config("pruned")) == _parentless(before) + for thread_id in ("pruned", "twin"): + await graph.ainvoke(Command(resume="yes 2"), _config(thread_id)) + assert (await graph.aget_state(_config("pruned"))).values == ( + await graph.aget_state(_config("twin")) + ).values + + +@pytest.mark.parametrize("make_checkpointer", ASYNC) +async def test_akeep_latest_refuses_a_batch_with_a_delta_channel_thread( + make_checkpointer: Any, +) -> None: + async with make_checkpointer() as checkpointer: + graph, delta_graph = _graph(checkpointer), _delta_graph(checkpointer) + for i in range(3): + await graph.ainvoke(_say(i), _config("ordinary")) + await delta_graph.ainvoke(_say(i), _config("delta")) + before = { + t: await _acheckpoint_ids(checkpointer, t) for t in ("ordinary", "delta") + } + + with pytest.raises(NotImplementedError, match="delta"): + await checkpointer.aprune(["ordinary", "delta"]) + + assert {t: await _acheckpoint_ids(checkpointer, t) for t in before} == before