diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py index 5e61b50ed..b177feec4 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py @@ -27,10 +27,19 @@ 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_REFERENCED_BLOBS_SQL, + PRUNE_STRATEGIES, + PRUNE_WRITES_SQL, BasePostgresSaver, _build_delta_stage1_sql, _build_delta_stage2_sql, _DeltaStage2Row, + _garbage_blobs, ) from langgraph.checkpoint.postgres.shallow import ShallowPostgresSaver @@ -401,6 +410,82 @@ 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 (plus its + parent when it predates v4 and has pending sends), with its pending + writes and the blobs it references, and deletes everything older. 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"` 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 while the + batch runs is refused at its turn, 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`. + """ + 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: + self._refuse_delta_threads(cur, 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)) + 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)}" + ) + @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..a27826c08 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py @@ -27,10 +27,19 @@ 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_REFERENCED_BLOBS_SQL, + PRUNE_STRATEGIES, + PRUNE_WRITES_SQL, BasePostgresSaver, _build_delta_stage1_sql, _build_delta_stage2_sql, _DeltaStage2Row, + _garbage_blobs, ) from langgraph.checkpoint.postgres.shallow import AsyncShallowPostgresSaver @@ -360,6 +369,71 @@ 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`. + """ + 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 self._refuse_delta_threads(cur, 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)) + 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)}" + ) + @asynccontextmanager async def _cursor( self, *, pipeline: bool = False @@ -688,6 +762,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..cdeb4425f 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py @@ -160,6 +160,114 @@ 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`), and keep_latest works one namespace +# at a time. + +# 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. +PRUNE_DELTA_THREADS_SQL = """ +SELECT thread_id FROM checkpoints +WHERE thread_id = ANY(%s) AND metadata ? 'counters_since_delta_snapshot' +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 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. +PRUNE_HEADS_SQL = f""" +SELECT + h.checkpoint_ns, + 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 +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 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_CHECKPOINTS_SQL = """ +WITH deleted AS ( + DELETE FROM checkpoints + WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id < %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_WRITES_SQL = """ +DELETE FROM checkpoint_writes +WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id < %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) +""" + + +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. + """ + 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 + + 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..6e0cd9061 --- /dev/null +++ b/libs/checkpoint-postgres/tests/test_prune.py @@ -0,0 +1,378 @@ +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, +) -> 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=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 = "") -> 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) + 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]] +) -> None: + pending = [action] + 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: + await self.cur.execute(query, params, **kwargs) + if query is sql and pending: + await pending.pop()() + + @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]}}, + ) + else: + first = await _put( + saver, thread_id, None, 0, {"m": _DeltaSnapshot([1])}, {"m": "1"} + ) + await _put(saver, thread_id, first, 1, {}, {"m": "2"}) + + +@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] + assert await _rows(saver, other_id) == other_before + + +@pytest.mark.parametrize("evidence", ["counters", "snapshot"]) +@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"] + + +@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"], + } + + +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 diff --git a/libs/langgraph/tests/test_prune_postgres.py b/libs/langgraph/tests/test_prune_postgres.py new file mode 100644 index 000000000..294711121 --- /dev/null +++ b/libs/langgraph/tests/test_prune_postgres.py @@ -0,0 +1,367 @@ +import operator +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, 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"), +] + +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 _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))] + + +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 _run_once_after(saver: PostgresSaver, sql: str, action: Callable[[], Any]) -> None: + pending = [action] + 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: + self.cur.execute(query, params, **kwargs) + if query is sql and pending: + pending.pop()() + + @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")) == 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")) == 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")) == 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) + namespaces = { + c.config["configurable"]["checkpoint_ns"] + for c in checkpointer.list(_config("pruned")) + } + + checkpointer.prune(["pruned"]) + + assert len(_checkpoint_ids(checkpointer, "pruned")) == len(namespaces) + assert graph.get_state(_config("pruned"), subgraphs=True) == 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 + ) + + +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 + ) + + +@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 + + +@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")) == 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