From 389b3c1d1f3e8eaf61cfcbf83ffe3814c75bbe2c Mon Sep 17 00:00:00 2001 From: Elior Nataf Lackritz Date: Wed, 14 Oct 2026 13:40:06 -0400 Subject: [PATCH] fix(checkpoint-postgres): delete finished subgraph calls and prune each thread on one snapshot keep_latest now deletes a subgraph call's namespace together with the root checkpoint the call started from. A checkpointer=True subgraph keeps its latest checkpoint, since its one namespace holds state across calls, and a call whose root checkpoint isn't saved yet is left alone. Each thread is pruned in one REPEATABLE READ transaction that deletes everything its snapshot sees except the kept checkpoints, so what a run commits meanwhile survives whatever its ID. Writes go by the checkpoints deleted, since a run can commit a superstep's writes before its checkpoint. The per-thread statements live in one generator shared by the sync and async savers, run unprepared: a cached generic plan checks an = ANY array one element per row. --- .../langgraph/checkpoint/postgres/__init__.py | 74 ++++------ .../langgraph/checkpoint/postgres/aio.py | 58 +++----- .../langgraph/checkpoint/postgres/base.py | 138 +++++++++++++++--- libs/checkpoint-postgres/tests/test_prune.py | 113 +++++++++++++- libs/langgraph/tests/test_prune_postgres.py | 57 +++++++- 5 files changed, 324 insertions(+), 116 deletions(-) diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py index e8b486437..d1e3cac59 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py @@ -27,20 +27,15 @@ from psycopg_pool import ConnectionPool from langgraph.checkpoint.postgres import _internal from langgraph.checkpoint.postgres.base import ( _DELTA_PAGE_SIZE, - PRUNE_BLOBS_SQL, - PRUNE_CHECKPOINTS_SQL, PRUNE_DELETE_THREADS_SQL, - PRUNE_DELTA_THREADS_SQL, - PRUNE_HEADS_SQL, - PRUNE_PARENTS_SQL, - PRUNE_REFERENCED_BLOBS_SQL, PRUNE_STRATEGIES, - PRUNE_WRITES_SQL, BasePostgresSaver, _build_delta_stage1_sql, _build_delta_stage2_sql, _DeltaStage2Row, - _garbage_blobs, + _execute, + _keep_latest, + _refuse_delta_threads, ) from langgraph.checkpoint.postgres.shallow import ShallowPostgresSaver @@ -422,13 +417,17 @@ class PostgresSaver(BasePostgresSaver): `"delete"` removes every checkpoint, write and blob of the threads, in one transaction. - `"keep_latest"` keeps the latest checkpoint of each namespace (plus its - parent when it predates v4 and has pending sends), with its pending - writes and the blobs it references, and deletes everything older. A - kept checkpoint whose parent it deletes is left with no parent. Each - thread is pruned in its own transaction. Checkpoints saved while it runs - are kept, so a run appending to the thread meanwhile is safe; forking or - replaying from a checkpoint being deleted is not. + `"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 @@ -444,6 +443,11 @@ class PostgresSaver(BasePostgresSaver): 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") @@ -458,38 +462,16 @@ class PostgresSaver(BasePostgresSaver): cur.execute(sql, (thread_ids,)) return with self._cursor() as cur: - self._refuse_delta_threads(cur, thread_ids) + _execute(cur, _refuse_delta_threads(thread_ids)) for thread_id in thread_ids: - with self._cursor() as cur, cur.connection.transaction(): - cur.execute(PRUNE_HEADS_SQL, (thread_id, thread_id, thread_id)) - heads = cur.fetchall() - # A run may have written this thread's first DeltaChannel - # checkpoint since the batch check. Checking after reading the - # heads sees it before anything it rebuilds from is deleted. - self._refuse_delta_threads(cur, [thread_id]) - for head in heads: - ns, keep_from = head["checkpoint_ns"], head["keep_from"] - cur.execute(PRUNE_CHECKPOINTS_SQL, (thread_id, ns, keep_from)) - freed = cur.fetchall() - cur.execute(PRUNE_WRITES_SQL, (thread_id, ns, keep_from)) - cur.execute( - PRUNE_PARENTS_SQL, (thread_id, ns, keep_from, keep_from) - ) - if not freed: - continue - cur.execute(PRUNE_REFERENCED_BLOBS_SQL, (thread_id, ns)) - garbage = _garbage_blobs(freed, cur.fetchall()) - for channel, versions in garbage.items(): - cur.execute(PRUNE_BLOBS_SQL, (thread_id, ns, channel, versions)) - - def _refuse_delta_threads( - self, cur: Cursor[DictRow], thread_ids: Sequence[str] - ) -> None: - cur.execute(PRUNE_DELTA_THREADS_SQL, (thread_ids, thread_ids)) - if delta := sorted(r["thread_id"] for r in cur.fetchall()): - raise NotImplementedError( - f"keep_latest can't prune threads that use a DeltaChannel: {', '.join(delta)}" - ) + 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]]: diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py index e04245fa8..716e54ed8 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py @@ -27,20 +27,15 @@ from psycopg_pool import AsyncConnectionPool from langgraph.checkpoint.postgres import _ainternal from langgraph.checkpoint.postgres.base import ( _DELTA_PAGE_SIZE, - PRUNE_BLOBS_SQL, - PRUNE_CHECKPOINTS_SQL, PRUNE_DELETE_THREADS_SQL, - PRUNE_DELTA_THREADS_SQL, - PRUNE_HEADS_SQL, - PRUNE_PARENTS_SQL, - PRUNE_REFERENCED_BLOBS_SQL, PRUNE_STRATEGIES, - PRUNE_WRITES_SQL, BasePostgresSaver, + _aexecute, _build_delta_stage1_sql, _build_delta_stage2_sql, _DeltaStage2Row, - _garbage_blobs, + _keep_latest, + _refuse_delta_threads, ) from langgraph.checkpoint.postgres.shallow import AsyncShallowPostgresSaver @@ -389,6 +384,11 @@ class AsyncPostgresSaver(BasePostgresSaver): 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") @@ -403,40 +403,16 @@ class AsyncPostgresSaver(BasePostgresSaver): await cur.execute(sql, (thread_ids,)) return async with self._cursor() as cur: - await self._refuse_delta_threads(cur, thread_ids) + await _aexecute(cur, _refuse_delta_threads(thread_ids)) for thread_id in thread_ids: - async with self._cursor() as cur, cur.connection.transaction(): - await cur.execute(PRUNE_HEADS_SQL, (thread_id, thread_id, thread_id)) - heads = await cur.fetchall() - # A run may have written this thread's first DeltaChannel - # checkpoint since the batch check. Checking after reading the - # heads sees it before anything it rebuilds from is deleted. - await self._refuse_delta_threads(cur, [thread_id]) - for head in heads: - ns, keep_from = head["checkpoint_ns"], head["keep_from"] - await cur.execute(PRUNE_CHECKPOINTS_SQL, (thread_id, ns, keep_from)) - freed = await cur.fetchall() - await cur.execute(PRUNE_WRITES_SQL, (thread_id, ns, keep_from)) - await cur.execute( - PRUNE_PARENTS_SQL, (thread_id, ns, keep_from, keep_from) - ) - if not freed: - continue - await cur.execute(PRUNE_REFERENCED_BLOBS_SQL, (thread_id, ns)) - garbage = _garbage_blobs(freed, await cur.fetchall()) - for channel, versions in garbage.items(): - await cur.execute( - PRUNE_BLOBS_SQL, (thread_id, ns, channel, versions) - ) - - async def _refuse_delta_threads( - self, cur: AsyncCursor[DictRow], thread_ids: Sequence[str] - ) -> None: - await cur.execute(PRUNE_DELTA_THREADS_SQL, (thread_ids, thread_ids)) - if delta := sorted(r["thread_id"] for r in await cur.fetchall()): - raise NotImplementedError( - f"keep_latest can't prune threads that use a DeltaChannel: {', '.join(delta)}" - ) + 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( diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py index 01a07f047..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 @@ -171,8 +173,7 @@ PRUNE_DELETE_THREADS_SQL = ( # 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`), and keep_latest works one namespace -# at a time. +# 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` @@ -199,19 +200,26 @@ CROSS JOIN LATERAL ( WHERE c.thread_id = ANY(%s) AND v.value = 'true'::jsonb """ -# Per namespace, the oldest checkpoint keep_latest keeps: the newest one, or -# its parent when the newest predates v4 and reads its pending sends from the -# parent's writes. +# 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 ELSE h.checkpoint_id END AS keep_from + ) 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 @@ -219,19 +227,31 @@ FROM ( ORDER BY checkpoint_ns DESC, checkpoint_id DESC ) AS h CROSS JOIN LATERAL ( - SELECT checkpoint, parent_checkpoint_id FROM checkpoints + 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 """ -# Deleting only below the newest checkpoint read above leaves alone whatever a -# concurrent run saves meanwhile. +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 < %s + 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 @@ -239,15 +259,10 @@ FROM deleted CROSS JOIN LATERAL jsonb_each_text(deleted.versions) AS v """ -PRUNE_WRITES_SQL = """ -DELETE FROM checkpoint_writes -WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id < %s -""" - PRUNE_PARENTS_SQL = """ UPDATE checkpoints SET parent_checkpoint_id = NULL WHERE thread_id = %s AND checkpoint_ns = %s -AND checkpoint_id >= %s AND parent_checkpoint_id < %s +AND checkpoint_id = ANY(%s) AND parent_checkpoint_id <> ALL(%s) """ PRUNE_REFERENCED_BLOBS_SQL = """ @@ -262,14 +277,17 @@ 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]]: - """Only blobs a deleted checkpoint used: a put that commits after `referenced` - was read isn't in it, so deleting every blob missing from it would take that - put's new blobs. - """ + """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: @@ -278,6 +296,84 @@ def _garbage_blobs( 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 index a6cd9dc7b..187abd21f 100644 --- a/libs/checkpoint-postgres/tests/test_prune.py +++ b/libs/checkpoint-postgres/tests/test_prune.py @@ -74,6 +74,7 @@ async def _put( 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: @@ -83,7 +84,7 @@ async def _put( checkpoint = empty_checkpoint() checkpoint.update( v=v, - id=str(uuid6(clock_seq=step)), + id=checkpoint_id or str(uuid6(clock_seq=step)), channel_values=values, channel_versions=versions, ) @@ -98,7 +99,12 @@ async def _put( ) -async def _three_checkpoints(saver: Any, thread_id: str, ns: str = "") -> dict: +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"}), @@ -106,7 +112,9 @@ async def _three_checkpoints(saver: Any, thread_id: str, ns: str = "") -> dict: ({"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) + 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 @@ -431,3 +439,102 @@ async def test_keep_latest_keeps_the_history_of_a_delta_checkpoint_saved_after_t ) 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 index 9eeefe945..d7aa00ea2 100644 --- a/libs/langgraph/tests/test_prune_postgres.py +++ b/libs/langgraph/tests/test_prune_postgres.py @@ -94,6 +94,22 @@ def _graph( 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)) @@ -113,6 +129,13 @@ 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))] @@ -256,14 +279,11 @@ def test_keep_latest_keeps_an_interrupted_subgraph_resumable( 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) - namespaces = { - c.config["configurable"]["checkpoint_ns"] - for c in checkpointer.list(_config("pruned")) - } + interrupted = before.tasks[0].state.config["configurable"]["checkpoint_ns"] checkpointer.prune(["pruned"]) - assert len(_checkpoint_ids(checkpointer, "pruned")) == len(namespaces) + 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)) @@ -273,6 +293,33 @@ def test_keep_latest_keeps_an_interrupted_subgraph_resumable( ) +@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,