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,