From ce3e64242b80c1ff4cdafdd49390d36c4e6793cc Mon Sep 17 00:00:00 2001 From: zeng-bohan Date: Fri, 18 Sep 2026 16:46:19 +0800 Subject: [PATCH] feat(gateway): checkpoint retention service on the #4189 deletion contract (#5308) * feat(gateway): thread checkpoint retention service on the #4189 deletion contract Implements exactly the two contract-proven deletion shapes (trailing duration-only leaves, opt-in leaf sibling branches) with head-chain protection, explicit id protection, a strict pending-writes guard, and joint writes-row cleanup. Head resolution uses LangGraph's time-ordered checkpoint ids; storage deletion mirrors the contract's per-backend data model. Ships without a production trigger by design. Validated against the contract suite (12 passed) plus 14 service scenarios across memory and SQLite; Postgres paths are gated on TEST_POSTGRES_URI. Signed-off-by: zengbohan1 <310902929+zengbohan1@users.noreply.github.com> * fix(gateway): survivor-reachability blob GC and memory blob stats in retention service Aligns the deletion service with the review-hardened contract: blob rows are garbage-collected in a whole-thread pass against surviving checkpoints' channel_versions (a real duration-only leaf shares its parent's versions, so per-checkpoint version deletion would corrupt the surviving state), the memory branch of the stats helper counts saver.blobs and returns the full normalized shape, and per-node channel versions are collected during the graph pass that already exists. Signed-off-by: zengbohan1 <310902929+zengbohan1@users.noreply.github.com> * fix(gateway): address review findings on checkpoint retention service Resolves the review at a479cfe (willem-bd): - Untested savers now fail fast: an explicit isinstance allowlist (InMemorySaver / AsyncSqliteSaver / AsyncPostgresSaver) raises NotImplementedError before any row is read or deleted, so a shallow or third-party saver can never issue partial DELETEs. - The chain walk ends (break) instead of raising KeyError when the head's ancestor row is missing, matching the deletable loop's tolerance for missing parents. - enforce_thread_retention takes an optional per-thread lock and documents the concurrency requirement: classification and deletion are two separate passes, so callers must serialize per-thread mutation (runtime _checkpoint_thread_lock) or guarantee quiescence. - Dropped the dead mid-run guard: CheckpointTuple has no `next` field in langgraph-checkpoint 4.1.1, and pending_writes is populated for committed writes too (verified on the list path), so neither is a usable mid-run signal; the caller-held thread lock is the actual protection. - Removed the write-only _node_step/_Node.step and fixed the head-selection docstring (newest by checkpoint id, not (step, checkpoint_id)). - Documented the E1 leaf / history fast-path interaction in the contract doc and module docstring: the wiring PR must sequence retention away from history reads or adopt a policy that spares cache-carrying leaves. - Added regression tests: unsupported saver, missing ancestor row, thread lock parameter. Validation: test_checkpoint_retention_service 18 passed / 8 postgres-gated skipped; contract + lineage suites 18 passed / 6 skipped; ruff check and format clean. * fix(retention): count non-empty writes dicts on memory saver - _checkpoint_ids_with_writes now requires a non-empty writes dict on InMemorySaver: the empty phantom entry for checkpoints whose task wrote nothing no longer counts as "owns writes rows", so the default E1 pruning reaches the memory backend again (it was a silent no-op there). - test_runtime_duration_leaf_pruned_by_default runs the shipping default (strict_pending_write_guard=True) and proves E1 is reachable out of the box on every backend; the stale override and its wrong SQLite premise are dropped. - document that _checkpoint_thread_lock is non-reentrant: a caller already holding it must not pass it in, or retention self-deadlocks. * test(checkpoint-retention): fix stray duplicated def token in test_duration_link_protected_after_next_run The previous push left `async def def test_...` at line 244, which made the module unimportable and failed collection of the whole suite (and ruff format --check). Local copy was already correct; this commit re-pushes the clean file. 18 passed / 8 postgres-skipped verified from a head worktree. * fix(gateway): make retention correct on Postgres and fail closed on a bad cap * validate max_delete_per_run before any store read: a negative cap used to widen the batch (Python slicing) instead of being rejected; * report identical before/after stats for an empty thread instead of returning before stats_after is collected; * protect each namespace's resume head and ancestor chain, so a persistent subgraph's latest checkpoint is no longer treated as a sibling leaf; * read Postgres columns through a row-factory-agnostic helper (the PG savers open cursors with dict_row, where positional access raises KeyError: 0); * classify the duration-only leaf without relying on metadata["writes"], which the Postgres saver strips via get_serializable_checkpoint_metadata. Verified locally on memory, SQLite and a real Postgres 16 instance (62 passed, 0 skipped): the E1 shape now fires on Postgres, which no backend test covered before CI ran the Postgres lig. Signed-off-by: zeng-bohan * test(gateway): pin the Postgres-shape duration classifier; report per-namespace heads - Deterministic regression for _mark_duration_leaves_without_the_marker: hand-put the Postgres round-trip shape (writes marker popped, source= update + accumulated run_durations + channel_versions identical to the parent) and assert the shipping default prunes it; a control that bumps one channel version (the client update_state shape) with otherwise identical metadata stays protected. Both legs run on memory and SQLite, so the class cannot silently re-widen (a resumable head losing head protection) or re-narrow (E1 never firing on Postgres) without a locally-executing test failing. - RetentionReport.protected_head_id -> protected_head_ids: heads are now selected per namespace, so the report carries every namespace's head (root key = what an unsaved aget_tuple resolves) instead of only the global max - reshape it before the wiring PR starts consuming reports for audit/aggregation. --------- Signed-off-by: zengbohan1 <310902929+zengbohan1@users.noreply.github.com> Signed-off-by: zeng-bohan Co-authored-by: zengbohan1 <310902929+zengbohan1@users.noreply.github.com> Co-authored-by: Willem Jiang --- backend/app/gateway/checkpoint_retention.py | 519 ++++++++++++++ backend/docs/checkpoint-retention-contract.md | 35 + .../test_checkpoint_retention_service.py | 669 ++++++++++++++++++ 3 files changed, 1223 insertions(+) create mode 100644 backend/app/gateway/checkpoint_retention.py create mode 100644 backend/tests/test_checkpoint_retention_service.py diff --git a/backend/app/gateway/checkpoint_retention.py b/backend/app/gateway/checkpoint_retention.py new file mode 100644 index 000000000..2d859bee7 --- /dev/null +++ b/backend/app/gateway/checkpoint_retention.py @@ -0,0 +1,519 @@ +"""Thread-level checkpoint retention enforcing the #4189 item 3 contract. + +Implements exactly the two deletion shapes proven safe by +``docs/checkpoint-retention-contract.md`` and its executable suite +(``tests/test_checkpoint_retention_contract.py``) — nothing else: + +- **Trailing duration-only leaves** — ``persist_run_durations`` appends + metadata-only checkpoints after a run finishes; while no later run has + forked from one, it is nobody's ancestor and can be dropped (contract + scenario E1). +- **Leaf sibling branches** *(opt-in)* — a checkpoint forked off an older + turn that has no children (contract scenario E2). Opt-in because a + superseded line's checkpoints may still be explicit resume targets a + client holds (protected set item 1); ``RetentionPolicy.protect_checkpoint_ids`` + is the escape hatch until a TTL semantic is agreed for that item. + +Everything on the resume head's ancestor chain — including duration-only +chain links, which would need grafting before deletion — every explicitly +protected id, and any checkpoint that still owns ``writes`` rows (protected +set item 3) is never deleted. Deletion runs at the storage layer, mirroring +the contract's per-backend data model, and removes the writes rows orphaned +by a deleted checkpoint in the same step (contract "deletion mechanics"). +Postgres ``checkpoint_blobs`` rows are keyed by ``version`` = the id of the +checkpoint that wrote the blob, so the same join keys clean them. + +The service ships **without a production trigger**: where retention is +invoked from (post-run hook vs scheduler vs explicit admin action) is a +maintainer decision that lands with the contract itself. Measurement-first: +reports carry before/after per-thread stats in the same normalized shape as +``scripts/benchmark/checkpoint/bench_channels.py``. + +History fast-path interaction: the trailing duration-only leaf is also the +carrier of the run-history metadata cache (``run_durations`` / +``run_message_ids``) that ``get_thread_history`` reads from the latest +checkpoint, and the parent it clones does not carry that map. Deleting the +leaf makes the next history read fall back to store scans and re-persist a +fresh leaf, so the wiring PR must sequence retention away from history reads +or adopt a policy that spares cache-carrying leaves — see the contract doc, +"History fast-path interaction (wiring requirement)". +""" + +from __future__ import annotations + +from collections import defaultdict +from contextlib import AbstractAsyncContextManager, nullcontext +from dataclasses import dataclass, field +from typing import Any + +from langgraph.checkpoint.base import BaseCheckpointSaver +from langgraph.checkpoint.memory import InMemorySaver +from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver + +from app.gateway.checkpoint_lineage import ( + checkpoint_configurable, + is_duration_only_checkpoint, +) + +__all__ = ["RetentionPolicy", "RetentionReport", "enforce_thread_retention"] + + +@dataclass(frozen=True) +class RetentionPolicy: + """What this run may delete. Defaults prune only what the contract proves safe unconditionally.""" + + prune_trailing_duration_leaves: bool = True + prune_leaf_sibling_branches: bool = False + protect_checkpoint_ids: frozenset[str] = frozenset() + strict_pending_write_guard: bool = True + max_delete_per_run: int | None = None + + +@dataclass +class RetentionReport: + """Outcome of one retention pass over one thread.""" + + thread_id: str + # Resume heads that survived this pass, per namespace (checkpoint_ns -> id). + # The root namespace (the "" key) is what an unsaved ``aget_tuple`` resolves + # as the thread's latest state; a persistent subgraph contributes its own + # child namespace whose head is protected too and would be invisible in a + # singular field. Consumers that only care about the thread's latest state + # read the root key. + protected_head_ids: dict[str, str] = field(default_factory=dict) + deleted_checkpoint_ids: list[str] = field(default_factory=list) + stats_before: dict[str, int] = field(default_factory=dict) + stats_after: dict[str, int] = field(default_factory=dict) + + +@dataclass +class _Node: + ns: str + cp_id: str + parent_ns: str | None + parent_id: str | None + duration_only: bool + metadata_source: str | None = None + carries_run_durations: bool = False + versions: frozenset = frozenset() + + +def _thread_config(thread_id: str) -> dict[str, Any]: + return {"configurable": {"thread_id": thread_id}} + + +def _row_field(row: Any, name: str, index: int) -> Any: + """Read one column from a DB-API row whatever row factory produced it. + + Postgres cursors are opened with ``row_factory=dict_row`` — both by this + codebase (``deerflow/runtime/checkpointer/async_provider.py``) and inside + ``langgraph-checkpoint-postgres`` itself (``aio.py`` opens every cursor as + ``conn.cursor(binary=True, row_factory=dict_row)``) — so their rows are + name-addressable and positional access raises ``KeyError: 0``. aiosqlite + rows are plain tuples. Reading through this helper keeps the per-backend + helpers independent of which saver opened the cursor: the Postgres leg of + ``tests/test_checkpoint_retention_service.py`` failed on exactly that + assumption, on six scenarios, the first time CI ran it. + """ + try: + return row[name] + except (TypeError, KeyError, IndexError): + return row[index] + + +def _mark_duration_leaves_without_the_marker(nodes: dict[tuple[str, str], _Node]) -> None: + """Classify metadata-only duration leaves on backends that drop the marker. + + ``persist_run_history_metadata`` is the only writer of these checkpoints and + stamps ``metadata["writes"]["runtime_run_duration"]``, which is what + :func:`is_duration_only_checkpoint` reads. The Postgres saver funnels + metadata through langgraph's ``get_serializable_checkpoint_metadata``, which + pops ``writes`` before the row is written (``checkpoint/base/__init__.py``; + only the Postgres savers call it), so on Postgres that marker never comes + back — every duration leaf would look resumable and the default E1 shape + would silently never fire on the production backend. + + The writer's other stamps survive a Postgres round trip and together are + exact: ``source == "update"``, a non-empty accumulated ``run_durations`` + map, and the leaf being a verbatim copy of its parent (``channel_versions`` + copied unchanged — the writer only replaces ``id``/``ts``). A client + ``update_state`` writes a newly versioned channel, so it fails the last + condition and is never swept into this class. + """ + for node in nodes.values(): + if node.duration_only or node.parent_id is None: + continue + parent = nodes.get((node.parent_ns, node.parent_id)) + if parent is None or node.metadata_source != "update" or not node.carries_run_durations: + continue + if node.versions and node.versions == parent.versions: + node.duration_only = True + + +def _ensure_supported_saver(saver: BaseCheckpointSaver) -> None: + """Fail fast before any row is read or written on an untested saver. + + The per-backend helpers below model exactly three storage layouts + (memory, SQLite, Postgres) and would otherwise fall through to "assume + Postgres" SQL for any other ``BaseCheckpointSaver``. A shallow Postgres + saver (no ``checkpoint_blobs``/``checkpoint_writes`` tables) or a + third-party saver would then issue DELETEs and die partway — with + autocommit on the Postgres connection, after the ``checkpoints`` row is + already gone. For a destructive tool, an explicit allowlist that raises + ``NotImplementedError`` beats a partial deletion and a confusing + traceback. + """ + if isinstance(saver, (InMemorySaver, AsyncSqliteSaver)): + return + async_postgres_saver: type | None = None + try: + from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver + + async_postgres_saver = AsyncPostgresSaver + except Exception: + pass + if async_postgres_saver is not None and isinstance(saver, async_postgres_saver): + return + raise NotImplementedError(f"checkpoint retention supports InMemorySaver, AsyncSqliteSaver and AsyncPostgresSaver, got {type(saver).__module__}.{type(saver).__name__}") + + +async def _thread_storage_stats(saver: Any, thread_id: str) -> dict[str, int]: + """Per-thread rows/bytes, same normalized shape as ``bench_channels``.""" + if isinstance(saver, InMemorySaver): + checkpoint_rows = checkpoint_bytes = blob_rows = blob_bytes = write_rows = write_bytes = 0 + for namespace in saver.storage.get(thread_id, {}).values(): + for checkpoint, metadata, _parent in namespace.values(): + checkpoint_rows += 1 + checkpoint_bytes += len(checkpoint[1]) + len(metadata[1]) + for (stored_thread, _ns, _channel, _version), (_type_tag, blob) in saver.blobs.items(): + if stored_thread != thread_id: + continue + blob_rows += 1 + blob_bytes += len(blob) + for (stored_thread, _ns, _cp_id), writes in saver.writes.items(): + if stored_thread != thread_id: + continue + for _task_id, _channel, (_type_tag, blob), _path in writes.values(): + write_rows += 1 + write_bytes += len(blob) + return { + "logical_checkpoint_bytes": checkpoint_bytes + blob_bytes, + "logical_write_bytes": write_bytes, + "checkpoint_rows": checkpoint_rows, + "checkpoint_bytes": checkpoint_bytes, + "blob_rows": blob_rows, + "blob_bytes": blob_bytes, + "write_rows": write_rows, + "write_bytes": write_bytes, + } + if isinstance(saver, AsyncSqliteSaver): + sqls = ( + ("checkpoint_rows", "checkpoint_bytes", "SELECT COUNT(*), COALESCE(SUM(LENGTH(checkpoint) + LENGTH(metadata)), 0) FROM checkpoints WHERE thread_id = ?"), + ("write_rows", "write_bytes", "SELECT COUNT(*), COALESCE(SUM(LENGTH(value)), 0) FROM writes WHERE thread_id = ?"), + ) + stats: dict[str, int] = {} + for row_key, bytes_key, sql in sqls: + async with saver.conn.execute(sql, (thread_id,)) as cursor: + row = await cursor.fetchone() + stats[row_key] = int(row[0]) + stats[bytes_key] = int(row[1] or 0) + stats["blob_rows"] = 0 + stats["blob_bytes"] = 0 + stats["logical_checkpoint_bytes"] = stats["checkpoint_bytes"] + stats["logical_write_bytes"] = stats["write_bytes"] + return stats + # Reachable only for the allowlisted Postgres saver (see + # ``_ensure_supported_saver``): its checkpoint rows are split into + # ``checkpoints`` + ``checkpoint_blobs``. + sqls = ( + ("checkpoint_rows", "checkpoint_bytes", "SELECT COUNT(*) AS rows, COALESCE(SUM(pg_column_size(checkpoint) + pg_column_size(metadata)), 0) AS bytes FROM checkpoints WHERE thread_id = %s"), + ("blob_rows", "blob_bytes", "SELECT COUNT(*) AS rows, COALESCE(SUM(octet_length(blob)), 0) AS bytes FROM checkpoint_blobs WHERE thread_id = %s"), + ("write_rows", "write_bytes", "SELECT COUNT(*) AS rows, COALESCE(SUM(octet_length(blob)), 0) AS bytes FROM checkpoint_writes WHERE thread_id = %s"), + ) + stats = {} + for row_key, bytes_key, sql in sqls: + async with saver._cursor() as cursor: + await cursor.execute(sql, (thread_id,)) + row = await cursor.fetchone() + stats[row_key] = int(_row_field(row, "rows", 0)) + stats[bytes_key] = int(_row_field(row, "bytes", 1) or 0) + stats["logical_checkpoint_bytes"] = stats["checkpoint_bytes"] + stats["blob_bytes"] + stats["logical_write_bytes"] = stats["write_bytes"] + return stats + + +async def _checkpoint_ids_with_writes(saver: Any, thread_id: str) -> set[tuple[str, str]]: + """(checkpoint_ns, checkpoint_id) pairs that still own writes rows. + + Protected set item 3: pending/uncommitted writes are retained state, not + garbage, so v1 refuses to delete any checkpoint that still owns writes + rows. On the memory backend ``InMemorySaver.writes`` also holds an *empty* + dict for checkpoints whose task produced no writes — counting key presence + would spare those phantom entries and silently disable pruning on the + memory saver, so a pair qualifies only when its writes dict is non-empty, + matching how :func:`_thread_storage_stats` counts rows rather than keys. + Production checkpoints normally accumulate their own committed writes rows + too, so the conservative default still makes pruning a no-op on hot + threads; a policy that distinguishes in-flight from orphaned writes belongs + to the contract's next revision, not to a fast path here. + """ + if isinstance(saver, InMemorySaver): + return {(ns, cp_id) for (stored_thread, ns, cp_id), writes in saver.writes.items() if stored_thread == thread_id and writes} + if isinstance(saver, AsyncSqliteSaver): + async with saver.conn.execute( + "SELECT DISTINCT checkpoint_ns, checkpoint_id FROM writes WHERE thread_id = ?", + (thread_id,), + ) as cursor: + rows = await cursor.fetchall() + # aiosqlite rows are plain tuples; the Postgres rows below are not. + return {(row[0] or "", row[1]) for row in rows} + async with saver._cursor() as cursor: + await cursor.execute( + "SELECT DISTINCT checkpoint_ns, checkpoint_id FROM checkpoint_writes WHERE thread_id = %s", + (thread_id,), + ) + rows = await cursor.fetchall() + return {(_row_field(row, "checkpoint_ns", 0) or "", _row_field(row, "checkpoint_id", 1)) for row in rows} + + +async def _delete_checkpoint_rows(saver: Any, thread_id: str, key: tuple[str, str]) -> None: + """Remove one checkpoint row and the writes rows it owns, jointly. + + Blob rows are handled by the survivor-reachability pass + (:func:`_delete_unreachable_blobs`), never per-checkpoint: versions are + shared between a checkpoint and its clones (see the contract doc). + """ + ns, cp_id = key + if isinstance(saver, InMemorySaver): + saver.storage.get(thread_id, {}).get(ns, {}).pop(cp_id, None) + saver.writes.pop((thread_id, ns, cp_id), None) + return + if isinstance(saver, AsyncSqliteSaver): + await saver.conn.execute( + "DELETE FROM checkpoints WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ?", + (thread_id, ns, cp_id), + ) + await saver.conn.execute( + "DELETE FROM writes WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ?", + (thread_id, ns, cp_id), + ) + await saver.conn.commit() + return + # Reachable only for the allowlisted Postgres saver (see + # ``_ensure_supported_saver``). + async with saver._cursor() as cursor: + await cursor.execute( + "DELETE FROM checkpoints WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id = %s", + (thread_id, ns, cp_id), + ) + await cursor.execute( + "DELETE FROM checkpoint_writes WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id = %s", + (thread_id, ns, cp_id), + ) + + +async def _delete_unreachable_blobs(saver: Any, thread_id: str, survivor_versions: set) -> None: + """Whole-thread blob GC: drop exactly the versions no surviving checkpoint + references (memory ``saver.blobs`` / Postgres ``checkpoint_blobs``).""" + if isinstance(saver, InMemorySaver): + for key in [key for key in saver.blobs if key[0] == thread_id and key[3] not in survivor_versions]: + del saver.blobs[key] + return + if isinstance(saver, AsyncSqliteSaver): + return + async with saver._cursor() as cursor: + await cursor.execute("SELECT DISTINCT version FROM checkpoint_blobs WHERE thread_id = %s", (thread_id,)) + rows = await cursor.fetchall() + orphans = [version for version in (_row_field(row, "version", 0) for row in rows) if version not in survivor_versions] + if orphans: + async with saver._cursor() as cursor: + await cursor.execute( + "DELETE FROM checkpoint_blobs WHERE thread_id = %s AND version = ANY(%s)", + (thread_id, orphans), + ) + + +async def enforce_thread_retention( + saver: BaseCheckpointSaver, + thread_id: str, + policy: RetentionPolicy | None = None, + *, + thread_lock: AbstractAsyncContextManager[None] | None = None, + collect_stats: bool = True, +) -> RetentionReport: + """Apply *policy* to one thread's checkpoints and return what happened. + + Classification walks the parent chain the same way + ``app/gateway/checkpoint_lineage.py`` does; each namespace's resume head is + that namespace's newest non-duration-only checkpoint by checkpoint id + (LangGraph ids are time-ordered), and every head's whole ancestor chain is + protected. Anything off those chains is only deletable when it is a leaf, + not explicitly protected, free of writes rows under the strict guard, and + matches one of the two contract-proven shapes. A node whose parent is + already missing is left alone: partial damage must not be silently + compounded. Only savers the deletion mechanics have been validated on are + accepted (see :func:`_ensure_supported_saver`). + + ``max_delete_per_run`` is validated before any row is read: a negative cap + raises ``ValueError`` rather than reaching the slice that applies it, where + Python's negative indexing would widen the batch instead of disabling it. + An empty thread returns a no-op report whose ``stats_after`` mirrors + ``stats_before``, so the measurement shape does not depend on whether the + thread had anything to classify. + + Concurrency requirement: classification and deletion are two separate + passes over the store, so a run that forks from a node classified as a + leaf in between leaves a dangling ``parent_config`` — the contract's own + "converts a cleanup into a thread-level outage" failure class. Callers + must therefore serialize per-thread mutation against the runtime writer + by passing the thread's checkpoint lock + (``deerflow.runtime.runs.worker._checkpoint_thread_lock(thread_id)``) as + *thread_lock*; without one, retention must only run while the thread is + guaranteed quiescent. + + The lock returned by ``_checkpoint_thread_lock`` is a plain, *non-reentrant* + ``asyncio.Lock`` (``AsyncKeyedLockTable.hold``). A caller that already + holds it and then passes it here self-deadlocks — which matters because + the runtime's own ``persist_run_history_metadata`` enters that same lock + before writing, so a post-run-hook call site must invoke retention *after* + releasing it, not from inside the held section. + """ + _ensure_supported_saver(saver) + effective = policy or RetentionPolicy() + if effective.max_delete_per_run is not None and effective.max_delete_per_run < 0: + # Fail closed *before* any store read. ``deletable[:cap]`` with a + # negative cap selects every candidate except the last few, so an + # invalid value meant to disable the run would instead maximize it — + # a batch bound must never widen a destructive pass. + raise ValueError(f"max_delete_per_run must be >= 0, got {effective.max_delete_per_run}") + report = RetentionReport(thread_id=thread_id) + lock: AbstractAsyncContextManager[None] = thread_lock if thread_lock is not None else nullcontext() + async with lock: + if collect_stats: + report.stats_before = await _thread_storage_stats(saver, thread_id) + + nodes: dict[tuple[str, str], _Node] = {} + async for tuple_ in saver.alist(_thread_config(thread_id), limit=None): + configurable = checkpoint_configurable(tuple_) + cp_id = configurable.get("checkpoint_id") + if not cp_id: + continue + ns = configurable.get("checkpoint_ns") or "" + parent_ns: str | None = None + parent_id: str | None = None + parent_config = getattr(tuple_, "parent_config", None) + if isinstance(parent_config, dict): + parent = parent_config.get("configurable") or {} + parent_ns = parent.get("checkpoint_ns") or "" + parent_id = parent.get("checkpoint_id") + checkpoint = getattr(tuple_, "checkpoint", None) or {} + channel_versions = checkpoint.get("channel_versions") + metadata = getattr(tuple_, "metadata", None) or {} + metadata = metadata if isinstance(metadata, dict) else {} + run_durations = metadata.get("run_durations") + nodes[(ns, cp_id)] = _Node( + ns=ns, + cp_id=cp_id, + parent_ns=parent_ns, + parent_id=parent_id, + duration_only=is_duration_only_checkpoint(tuple_), + metadata_source=metadata.get("source"), + carries_run_durations=isinstance(run_durations, dict) and bool(run_durations), + versions=frozenset(channel_versions.values()) if isinstance(channel_versions, dict) else frozenset(), + ) + if not nodes: + # Nothing to classify: an empty (or unknown) thread is a no-op, but + # the before/after pair is part of the report contract, so it is + # mirrored rather than truncated — a caller aggregating measurements + # must not have to special-case "thread had no checkpoints". + if collect_stats: + report.stats_after = dict(report.stats_before) + return report + _mark_duration_leaves_without_the_marker(nodes) + + children: dict[tuple[str, str], list[tuple[str, str]]] = defaultdict(list) + for key, node in nodes.items(): + if node.parent_id is not None: + children[(node.parent_ns, node.parent_id)].append(key) + + resumable = [key for key, node in nodes.items() if not node.duration_only] + # Head = newest by checkpoint id. LangGraph ids are time-ordered (uuid7): + # metadata step restarts from the fork point after a branch-resume, so it + # is not a thread-global sequence, while max-id matches what an unsaved + # ``aget_tuple`` resolves as the thread's latest state. + # + # Heads are selected *per namespace*. ``alist`` is called with the thread + # id alone, so ``nodes`` holds every namespace in the thread — a + # persistent subgraph (compiled with ``checkpointer=True``) keeps its own + # checkpoints under e.g. ``tools:``. Protecting only one global + # head leaves the child namespace's latest checkpoint looking like an + # off-chain sibling leaf, which the opt-in E2 shape would then delete, + # rolling that graph's saved state back one step. + heads: dict[str, tuple[str, str]] = {} + for key in resumable: + current = heads.get(key[0]) + if current is None or key[1] > current[1]: + heads[key[0]] = key + report.protected_head_ids = {ns: key[1] for ns, key in heads.items()} + + chain: set[tuple[str, str]] = set() + for head in heads.values(): + cursor: tuple[str, str] | None = head + while cursor is not None: + if cursor not in nodes: + # A head whose ancestor row is missing (partial damage from + # an earlier policy revision or manual cleanup) ends this + # walk instead of crashing the whole pass; the deletable loop + # below already leaves nodes with missing parents alone. + break + chain.add(cursor) + node = nodes[cursor] + cursor = (node.parent_ns, node.parent_id) if node.parent_id is not None else None + + guarded: set[tuple[str, str]] = set() + if effective.strict_pending_write_guard: + guarded = await _checkpoint_ids_with_writes(saver, thread_id) + + deletable: list[tuple[str, str]] = [] + for key, node in nodes.items(): + if key in chain: + continue + if children.get(key): + continue + if key[1] in effective.protect_checkpoint_ids: + continue + if key in guarded: + continue + if node.parent_id is not None and (node.parent_ns, node.parent_id) not in nodes: + continue + if node.duration_only: + if not effective.prune_trailing_duration_leaves: + continue + elif not effective.prune_leaf_sibling_branches: + continue + deletable.append(key) + + if effective.max_delete_per_run is not None: + deletable = deletable[: effective.max_delete_per_run] + + # Blob GC, contract deletion mechanics: a blob row is an orphan only if no + # SURVIVING checkpoint references its version. A real duration-only leaf + # copies its parent's channel_versions verbatim, so its blobs are the + # parent's rows — deleting "blobs keyed by the removed checkpoint's own + # versions" would corrupt the surviving state. + deleted_keys = set(deletable) + survivor_versions: set[Any] = set() + for key, node in nodes.items(): + if key not in deleted_keys: + survivor_versions.update(node.versions) + + for key in deletable: + await _delete_checkpoint_rows(saver, thread_id, key) + report.deleted_checkpoint_ids.append(key[1]) + + if deletable: + await _delete_unreachable_blobs(saver, thread_id, survivor_versions) + + if collect_stats: + report.stats_after = await _thread_storage_stats(saver, thread_id) + return report diff --git a/backend/docs/checkpoint-retention-contract.md b/backend/docs/checkpoint-retention-contract.md index 50e34673a..33d76845b 100644 --- a/backend/docs/checkpoint-retention-contract.md +++ b/backend/docs/checkpoint-retention-contract.md @@ -97,6 +97,41 @@ resume, (c) branch from an older visible turn, and (d) orphan row counts. `scripts/benchmark/checkpoint/bench_channels.py` (per-thread rows/bytes, SQLite and Postgres) plus the contract test suite passing. +## History fast-path interaction (wiring requirement) + +The trailing duration-only leaf is also the carrier of the run-history +metadata cache: `persist_run_history_metadata` accumulates `run_durations` +and `run_message_ids` in the leaf's metadata, and +`app/gateway/routers/threads.py::get_thread_history` reads that map from the +latest checkpoint (`_checkpoint_run_durations` / +`_checkpoint_run_message_ids`, gated on `is_latest_checkpoint`) to answer +every known turn's duration and message-to-run attribution without scanning +the event store. The parent checkpoint the leaf clones does **not** carry +that map. + +Deleting the leaf (scenario E1) therefore removes the fast-path cache: the +next history read sees no durations, falls back to event-store + run-manager +scans, and `_persist_run_history_metadata_background` re-writes a fresh +duration-only leaf — which the next retention pass deletes again. Net effect +without sequencing: the reclaimed row comes straight back, plus recurring +store scans and an extra write per read. + +The wiring PR that introduces the production trigger must therefore either: + +1. **Sequence retention away from history reads** — e.g. run retention on a + schedule whose next pass re-reclaims the re-created leaf, or run it when + the thread is not being read; or +2. **Adopt a policy that spares cache-carrying leaves** — e.g. a + `RetentionPolicy` flag that keeps any trailing duration-only leaf whose + metadata still carries `run_durations` / `run_message_ids` (same spirit + as the strict pending-writes guard), at the cost of not reclaiming that + leaf's rows. + +Without either, E1 pruning and history reads churn against each other. This +decision belongs to the wiring PR, not to the storage-level service: the +service cannot tell a cache-carrying leaf from a payload-free one on the +alist path without re-implementing the writer's merge semantics. + ## Item 4 note (large tool results) `ToolOutputBudgetMiddleware` externalizes oversized tool outputs before they diff --git a/backend/tests/test_checkpoint_retention_service.py b/backend/tests/test_checkpoint_retention_service.py new file mode 100644 index 000000000..753845171 --- /dev/null +++ b/backend/tests/test_checkpoint_retention_service.py @@ -0,0 +1,669 @@ +"""Behavioral tests for ``app/gateway/checkpoint_retention.py``. + +Each test drives the retention service over real saver backends (memory, +SQLite, and Postgres when ``TEST_POSTGRES_URI`` is set) using the same chain +constructions as the contract suite, then verifies the contract's four +post-deletion properties: latest resume, explicit ``checkpoint_id`` resume, +branch/regenerate lineage walk, and orphan row accounting. The duration-only +checkpoints are produced by the real runtime writer (``persist_run_durations``), +not hand-rolled ``aput`` calls, so the classification runs against the exact +metadata shape production emits. +""" + +from __future__ import annotations + +import asyncio +import copy +import os +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager +from typing import Annotated, Any, TypedDict +from uuid import uuid4 + +import pytest +from langchain_core.messages import AnyMessage, HumanMessage +from langgraph.checkpoint.memory import InMemorySaver +from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver +from langgraph.graph import END, START, StateGraph +from langgraph.graph.message import add_messages + +from app.gateway import checkpoint_retention +from app.gateway.checkpoint_lineage import find_checkpoint_before_message +from app.gateway.checkpoint_retention import ( + RetentionPolicy, + _row_field, + enforce_thread_retention, +) +from deerflow.runtime.runs.worker import _new_checkpoint_marker, persist_run_durations + + +class FullState(TypedDict): + messages: Annotated[list[AnyMessage], add_messages] + + +def _thread_id() -> str: + return f"retention-service-{uuid4().hex}" + + +def _config(thread_id: str) -> dict[str, Any]: + return {"configurable": {"thread_id": thread_id}} + + +def _config_thread(thread_id: str, checkpoint_id: str) -> dict[str, Any]: + return {"configurable": {"thread_id": thread_id, "checkpoint_id": checkpoint_id}} + + +def _noop(state: dict[str, Any]) -> dict[str, Any]: + return {} + + +def _build_graph(schema: Any, checkpointer: Any) -> Any: + builder = StateGraph(schema) + builder.add_node("noop", _noop) + builder.set_entry_point("noop") + builder.set_finish_point("noop") + return builder.compile(checkpointer=checkpointer) + + +class _SaverEnv: + def __init__(self, kind: str, open_saver: Any) -> None: + self.kind = kind + self._open_saver = open_saver + self._cm: Any | None = None + self.saver: Any | None = None + + async def __aenter__(self) -> _SaverEnv: + self._cm = self._open_saver() + self.saver = await self._cm.__aenter__() + setup = getattr(self.saver, "setup", None) + if setup is not None: + await setup() + return self + + async def __aexit__(self, *exc: Any) -> None: + if self._cm is not None: + await self._cm.__aexit__(*exc) + self._cm = None + self.saver = None + + +@asynccontextmanager +async def _open_sqlite(db_path: Any) -> AsyncIterator[Any]: + async with AsyncSqliteSaver.from_conn_string(str(db_path)) as saver: + await saver.setup() + yield saver + + +@asynccontextmanager +async def _open_postgres(uri: str) -> AsyncIterator[Any]: + aio = pytest.importorskip("langgraph.checkpoint.postgres.aio", reason="postgres extra not installed") + async with aio.AsyncPostgresSaver.from_conn_string(uri) as saver: + await saver.setup() + yield saver + + +@pytest.fixture(params=["memory", "sqlite", "postgres"]) +async def saver_env(request: pytest.FixtureRequest, tmp_path: Any) -> AsyncIterator[_SaverEnv]: + kind = request.param + if kind == "memory": + saver = InMemorySaver() + + @asynccontextmanager + async def open_memory() -> AsyncIterator[Any]: + yield saver + + open_saver = open_memory + elif kind == "sqlite": + db_path = tmp_path / "retention-service.sqlite" + + def open_sqlite() -> Any: + return _open_sqlite(db_path) + + open_saver = open_sqlite + else: + uri = os.environ.get("TEST_POSTGRES_URI") + if not uri: + pytest.skip("TEST_POSTGRES_URI is not set") + + def open_postgres() -> Any: + return _open_postgres(uri) + + open_saver = open_postgres + + async with _SaverEnv(kind, open_saver) as env: + yield env + + +class _SaverAccessor: + def __init__(self, saver: Any) -> None: + self._saver = saver + + async def aget(self, config: dict[str, Any]) -> Any: + return await self._saver.aget_tuple(config) + + +async def _walk(env: _SaverEnv, head_config: dict[str, Any], message_id: str) -> Any: + return await find_checkpoint_before_message( + _SaverAccessor(env.saver), + await env.saver.aget_tuple(head_config), + message_id, + max_depth=50, + ) + + +async def _write_turns( + env: _SaverEnv, + steps: int, + *, + payload_bytes: int = 256, +) -> tuple[str, list[str], list[str]]: + graph = _build_graph(FullState, env.saver) + thread_id = _thread_id() + checkpoint_ids: list[str] = [] + for index in range(steps): + message = HumanMessage(content=f"turn {index}: " + "x" * payload_bytes, id=f"turn-{index}") + await graph.ainvoke({"messages": [message]}, _config(thread_id)) + snapshot = await graph.aget_state(_config(thread_id)) + checkpoint_ids.append(snapshot.config["configurable"]["checkpoint_id"]) + return thread_id, checkpoint_ids, [f"turn-{index}" for index in range(steps)] + + +async def _append_duration_checkpoint(env: _SaverEnv, thread_id: str, run_id: str = "run-1") -> str: + """Append a duration-only checkpoint through the real runtime writer.""" + written = await persist_run_durations(checkpointer=env.saver, thread_id=thread_id, durations={run_id: 7}) + assert written, "persist_run_durations must land a metadata-only checkpoint" + head = await env.saver.aget_tuple(_config(thread_id)) + assert head is not None + return head.checkpoint["id"] + + +async def _listed_checkpoint_ids(env: _SaverEnv, thread_id: str) -> set[str]: + return {tuple_.checkpoint["id"] async for tuple_ in env.saver.alist(_config(thread_id), limit=None)} + + +async def _write_count(env: _SaverEnv, thread_id: str, checkpoint_id: str) -> int: + if env.kind == "memory": + return len(env.saver.writes.get((thread_id, "", checkpoint_id), {})) + if env.kind == "sqlite": + async with env.saver.conn.execute( + "SELECT COUNT(*) FROM writes WHERE thread_id = ? AND checkpoint_id = ?", + (thread_id, checkpoint_id), + ) as cursor: + row = await cursor.fetchone() + return int(_row_field(row, "COUNT(*)", 0)) + async with env.saver._cursor() as cursor: + await cursor.execute( + "SELECT COUNT(*) AS write_count FROM checkpoint_writes WHERE thread_id = %s AND checkpoint_id = %s", + (thread_id, checkpoint_id), + ) + row = await cursor.fetchone() + return int(_row_field(row, "write_count", 0)) + + +# --------------------------------------------------------------------------- +# Scenarios +# --------------------------------------------------------------------------- + + +@pytest.mark.anyio +async def test_linear_thread_prunes_nothing(saver_env: _SaverEnv) -> None: + """A plain linear thread is one protected ancestor chain: nothing may go.""" + thread_id, checkpoint_ids, _message_ids = await _write_turns(saver_env, steps=4) + + report = await enforce_thread_retention(saver_env.saver, thread_id) + + assert report.deleted_checkpoint_ids == [] + assert report.protected_head_ids == {"": checkpoint_ids[-1]} + # one invoke lands several checkpoints (input/task/result); the ids we + # collected are the resumable results, and all of them must survive + listed = await _listed_checkpoint_ids(saver_env, thread_id) + assert set(checkpoint_ids).issubset(listed) + + +@pytest.mark.anyio +async def test_runtime_duration_leaf_pruned_by_default(saver_env: _SaverEnv) -> None: + """Contract E1 via the real writer: a trailing duration-only leaf is pruned, + the finished run's final checkpoint stays resumable and walkable, and the + row accounting reflects exactly one reclaimed checkpoint.""" + thread_id, checkpoint_ids, message_ids = await _write_turns(saver_env, steps=3) + duration_id = await _append_duration_checkpoint(saver_env, thread_id) + + # Shipping default (strict_pending_write_guard=True): the duration-only + # leaf owns no writes rows — on memory its ``writes`` entry is the phantom + # empty dict that ``_checkpoint_ids_with_writes`` does not count — so the + # default policy itself prunes it. This is the proof that the headline + # "enabled by default" E1 shape is reachable on every backend, not only + # after relaxing the guard. + report = await enforce_thread_retention(saver_env.saver, thread_id) + + assert report.deleted_checkpoint_ids == [duration_id] + assert report.protected_head_ids == {"": checkpoint_ids[-1]} + assert duration_id not in await _listed_checkpoint_ids(saver_env, thread_id) + resumed = await saver_env.saver.aget_tuple(_config_thread(thread_id, checkpoint_ids[-1])) + assert resumed is not None + base = await _walk(saver_env, _config_thread(thread_id, checkpoint_ids[-1]), message_ids[-1]) + assert base is not None + assert report.stats_after["checkpoint_rows"] == report.stats_before["checkpoint_rows"] - 1 + + +@pytest.mark.anyio +async def test_postgres_round_trip_shape_pruned_by_default(saver_env: _SaverEnv) -> None: + """Deterministic pin for the shape-based fallback classifier, without Postgres. + + On memory/SQLite the writer's ``writes`` marker survives a round trip, so + ``is_duration_only_checkpoint`` already classifies every writer-produced + leaf and the fallback in ``_mark_duration_leaves_without_the_marker`` is + only reached on the TEST_POSTGRES_URI-gated leg. This test hand-puts the + Postgres round-trip shape — a verbatim clone of its parent (fresh id/ts, + ``channel_versions`` copied unchanged) whose metadata has the ``writes`` + marker popped but the writer's surviving stamps intact (``source == + "update"``, a non-empty accumulated ``run_durations`` map) — and asserts + the shipping default still prunes it. The control bumps one channel + version (the shape a client ``update_state`` produces) with otherwise + identical metadata and must stay protected, so the class can neither + silently re-widen (a resumable head would lose head protection) nor + re-narrow (E1 would never fire on the production backend) without this + test failing. + """ + thread_id, checkpoint_ids, _message_ids = await _write_turns(saver_env, steps=3) + head = await saver_env.saver.aget_tuple(_config(thread_id)) + assert head is not None + head_id = head.checkpoint["id"] + + def _pg_round_trip_meta(parent_meta: dict[str, Any]) -> dict[str, Any]: + meta = dict(parent_meta or {}) + meta.pop("writes", None) # what get_serializable_checkpoint_metadata does on Postgres + meta["source"] = "update" + meta["run_durations"] = {"run-1": 7} + meta["step"] = (meta["step"] + 1) if isinstance(meta.get("step"), int) else 1 + return meta + + def _leaf_config(parent_id: str) -> dict[str, Any]: + # aput needs the full configurable (namespace included) for the parent link. + return {"configurable": {"thread_id": thread_id, "checkpoint_ns": "", "checkpoint_id": parent_id}} + + # Control: one NEW channel version on an otherwise verbatim clone — a + # client ``update_state`` can produce this, the runtime writer cannot. + control = copy.deepcopy(dict(head.checkpoint)) + control.update(_new_checkpoint_marker()) + control["channel_versions"] = dict(control["channel_versions"]) + # A version value no real channel would carry; the classifier compares the + # frozenset of versions against the parent's, so any strictly-different + # set is what matters. + control["channel_versions"]["_control"] = "__control_version__" + control_id = control["id"] + control_meta = _pg_round_trip_meta(head.metadata) + await saver_env.saver.aput(_leaf_config(head_id), control, control_meta, {}) + + # The Postgres round-trip shape on top: verbatim clone of the control + # (identical ``channel_versions``), marker popped from the metadata. + pg_leaf = copy.deepcopy(control) + pg_leaf.update(_new_checkpoint_marker()) + pg_leaf_id = pg_leaf["id"] + await saver_env.saver.aput(_leaf_config(control_id), pg_leaf, _pg_round_trip_meta(control_meta), {}) + + report = await enforce_thread_retention(saver_env.saver, thread_id) + + assert report.deleted_checkpoint_ids == [pg_leaf_id] + assert report.protected_head_ids == {"": control_id} + assert pg_leaf_id not in await _listed_checkpoint_ids(saver_env, thread_id) + assert control_id in await _listed_checkpoint_ids(saver_env, thread_id) + + +@pytest.mark.anyio +async def test_duration_link_protected_after_next_run(saver_env: _SaverEnv) -> None: + """Contract protected set item 4: once a later run has been written on top + of a duration-only checkpoint, that checkpoint is a chain link on the new + head's ancestor chain and must not be touched (deleting it would need + grafting, which v1 does not attempt).""" + thread_id, checkpoint_ids, _message_ids = await _write_turns(saver_env, steps=3) + duration_id = await _append_duration_checkpoint(saver_env, thread_id) + + graph = _build_graph(FullState, saver_env.saver) + follow_up = HumanMessage(content="turn after duration: " + "x" * 64, id="turn-after-duration") + await graph.ainvoke({"messages": [follow_up]}, _config(thread_id)) + snapshot = await graph.aget_state(_config(thread_id)) + assert snapshot.config["configurable"]["checkpoint_id"] not in (*checkpoint_ids, duration_id) + + report = await enforce_thread_retention(saver_env.saver, thread_id) + + assert report.deleted_checkpoint_ids == [] + assert duration_id in await _listed_checkpoint_ids(saver_env, thread_id) + assert report.protected_head_ids == {"": snapshot.config["configurable"]["checkpoint_id"]} + + +@pytest.mark.anyio +async def test_regenerated_old_head_pruned_opt_in(saver_env: _SaverEnv) -> None: + """Contract E2 shape in its production direction: after a regenerate, the + fork is the live resume head and the superseded old head becomes a leaf + sibling. With the opt-in flag the old head is pruned while the fork line + keeps working end to end.""" + thread_id, checkpoint_ids, _message_ids = await _write_turns(saver_env, steps=4) + old_head_id = checkpoint_ids[-1] + + graph = _build_graph(FullState, saver_env.saver) + fork_message = HumanMessage(content="regenerated turn: " + "z" * 256, id="fork-turn") + await graph.ainvoke({"messages": [fork_message]}, _config_thread(thread_id, checkpoint_ids[1])) + snapshot = await graph.aget_state(_config(thread_id)) + fork_head_id = snapshot.config["configurable"]["checkpoint_id"] + assert fork_head_id not in (*checkpoint_ids,) + + policy = RetentionPolicy(prune_leaf_sibling_branches=True, strict_pending_write_guard=False) + report = await enforce_thread_retention(saver_env.saver, thread_id, policy) + + assert report.deleted_checkpoint_ids == [old_head_id] + assert report.protected_head_ids == {"": fork_head_id} + + fork_tuple = await saver_env.saver.aget_tuple(_config_thread(thread_id, fork_head_id)) + assert fork_tuple is not None + base = await _walk(saver_env, _config_thread(thread_id, fork_head_id), "fork-turn") + assert base is not None + remaining = await _listed_checkpoint_ids(saver_env, thread_id) + assert old_head_id not in remaining + assert {checkpoint_ids[0], checkpoint_ids[1], fork_head_id}.issubset(remaining) + + +@pytest.mark.anyio +async def test_explicit_protect_ids_spare_the_superseded_head(saver_env: _SaverEnv) -> None: + """Protected set item 1: a client-held checkpoint id wins over pruning. + The same thread prunes only once the id leaves the protect list.""" + thread_id, checkpoint_ids, _message_ids = await _write_turns(saver_env, steps=4) + old_head_id = checkpoint_ids[-1] + + graph = _build_graph(FullState, saver_env.saver) + fork_message = HumanMessage(content="regenerated turn: " + "z" * 256, id="fork-turn") + await graph.ainvoke({"messages": [fork_message]}, _config_thread(thread_id, checkpoint_ids[1])) + + protected_policy = RetentionPolicy(prune_leaf_sibling_branches=True, protect_checkpoint_ids=frozenset({old_head_id})) + report = await enforce_thread_retention(saver_env.saver, thread_id, protected_policy) + assert report.deleted_checkpoint_ids == [] + kept = await saver_env.saver.aget_tuple(_config_thread(thread_id, old_head_id)) + assert kept is not None + + report = await enforce_thread_retention( + saver_env.saver, + thread_id, + RetentionPolicy(prune_leaf_sibling_branches=True, strict_pending_write_guard=False), + ) + assert report.deleted_checkpoint_ids == [old_head_id] + + +@pytest.mark.anyio +async def test_strict_pending_write_guard_spares_leaf_and_cleans_orphans_after(saver_env: _SaverEnv) -> None: + """Protected set item 3 + deletion mechanics: a leaf that still owns writes + rows is spared under the strict guard; deleting it with the guard relaxed + removes its writes rows with it, and writes held by retained checkpoints + survive.""" + thread_id, checkpoint_ids, _message_ids = await _write_turns(saver_env, steps=4) + old_head_id = checkpoint_ids[-1] + + graph = _build_graph(FullState, saver_env.saver) + fork_message = HumanMessage(content="regenerated turn: " + "z" * 256, id="fork-turn") + await graph.ainvoke({"messages": [fork_message]}, _config_thread(thread_id, checkpoint_ids[1])) + snapshot = await graph.aget_state(_config(thread_id)) + fork_head_id = snapshot.config["configurable"]["checkpoint_id"] + + write = ("messages", ("human", b"pending-write")) + old_head_config = {"configurable": {"thread_id": thread_id, "checkpoint_ns": "", "checkpoint_id": old_head_id}} + fork_head_config = {"configurable": {"thread_id": thread_id, "checkpoint_ns": "", "checkpoint_id": fork_head_id}} + await saver_env.saver.aput_writes(old_head_config, [write], task_id="orphan-task") + await saver_env.saver.aput_writes(fork_head_config, [write], task_id="pending-task") + assert await _write_count(saver_env, thread_id, old_head_id) > 0 + + guarded_report = await enforce_thread_retention( + saver_env.saver, + thread_id, + RetentionPolicy(prune_leaf_sibling_branches=True, strict_pending_write_guard=True), + ) + assert guarded_report.deleted_checkpoint_ids == [] + assert await _write_count(saver_env, thread_id, old_head_id) > 0 + + relaxed_report = await enforce_thread_retention( + saver_env.saver, + thread_id, + RetentionPolicy(prune_leaf_sibling_branches=True, strict_pending_write_guard=False), + ) + assert relaxed_report.deleted_checkpoint_ids == [old_head_id] + assert await _write_count(saver_env, thread_id, old_head_id) == 0, "orphaned writes rows must go with their checkpoint" + assert await _write_count(saver_env, thread_id, fork_head_id) > 0, "writes of retained checkpoints must survive" + + fork_tuple = await saver_env.saver.aget_tuple(_config_thread(thread_id, fork_head_id)) + assert fork_tuple is not None + base = await _walk(saver_env, _config_thread(thread_id, fork_head_id), "fork-turn") + assert base is not None + + +@pytest.mark.anyio +async def test_max_delete_per_run_caps_the_batch(saver_env: _SaverEnv) -> None: + """Two prunable leaves (a trailing duration-only leaf and a superseded old + head) with a cap of one: exactly one row goes, the other survives.""" + thread_id, checkpoint_ids, _message_ids = await _write_turns(saver_env, steps=3) + old_head_id = checkpoint_ids[-1] + + graph = _build_graph(FullState, saver_env.saver) + fork_message = HumanMessage(content="regenerated turn: " + "z" * 256, id="fork-turn") + await graph.ainvoke({"messages": [fork_message]}, _config_thread(thread_id, checkpoint_ids[1])) + + duration_id = await _append_duration_checkpoint(saver_env, thread_id) + + policy = RetentionPolicy( + prune_leaf_sibling_branches=True, + strict_pending_write_guard=False, + max_delete_per_run=1, + ) + report = await enforce_thread_retention(saver_env.saver, thread_id, policy) + + assert len(report.deleted_checkpoint_ids) == 1 + assert report.deleted_checkpoint_ids[0] in (old_head_id, duration_id) + remaining = await _listed_checkpoint_ids(saver_env, thread_id) + assert len(remaining & {old_head_id, duration_id}) == 1 + + +@pytest.mark.anyio +async def test_unsupported_saver_raises_before_any_read() -> None: + """An untested saver is rejected up front: no row is read or deleted, so + no partial deletion can happen on a backend the mechanics were not + validated on.""" + + class _FakeSaver: + async def alist(self, *args: Any, **kwargs: Any) -> AsyncIterator[Any]: + raise AssertionError("alist must not run on an unsupported saver") + yield # pragma: no cover + + with pytest.raises(NotImplementedError): + await enforce_thread_retention(_FakeSaver(), _thread_id()) + + +@pytest.mark.anyio +async def test_chain_walk_tolerates_missing_ancestor_row() -> None: + """A head whose ancestor row is missing (partial damage from an earlier + policy revision or manual cleanup) ends the chain walk instead of + crashing the whole pass; the surviving protections still hold.""" + saver = InMemorySaver() + graph = _build_graph(FullState, saver) + thread_id = _thread_id() + checkpoint_ids: list[str] = [] + for index in range(3): + message = HumanMessage(content=f"turn {index}: " + "x" * 256, id=f"turn-{index}") + await graph.ainvoke({"messages": [message]}, _config(thread_id)) + snapshot = await graph.aget_state(_config(thread_id)) + checkpoint_ids.append(snapshot.config["configurable"]["checkpoint_id"]) + + # Simulate a partially pruned thread: the middle resumable row is gone. + namespace = saver.storage[thread_id][""] + assert checkpoint_ids[1] in namespace + namespace.pop(checkpoint_ids[1], None) + + report = await enforce_thread_retention(saver, thread_id) + + assert report.protected_head_ids == {"": checkpoint_ids[-1]} + # The head is protected; the remaining off-chain node (the oldest turn) is + # a leaf sibling, which is not pruned unless opted in. + assert report.deleted_checkpoint_ids == [] + + +@pytest.mark.anyio +async def test_thread_lock_parameter_accepted(saver_env: _SaverEnv) -> None: + """An explicit per-thread lock is honored: with an uncontended lock the + call completes with the same outcome as without one.""" + thread_id, checkpoint_ids, _message_ids = await _write_turns(saver_env, steps=2) + + lock = asyncio.Lock() + report = await enforce_thread_retention(saver_env.saver, thread_id, thread_lock=lock) + + assert report.deleted_checkpoint_ids == [] + assert report.protected_head_ids == {"": checkpoint_ids[-1]} + + +# --------------------------------------------------------------------------- +# Policy validation and report shape +# --------------------------------------------------------------------------- + + +@pytest.mark.anyio +async def test_negative_max_delete_per_run_fails_closed_before_any_read(saver_env: _SaverEnv, monkeypatch: pytest.MonkeyPatch) -> None: + """A negative cap is an invalid configuration, not a small one. It must + raise before the store is touched — the slice that applies the cap would + otherwise turn -1 into "everything except the last candidate", i.e. widen a + destructive pass to almost its maximum — and it must delete nothing.""" + thread_id, checkpoint_ids, _message_ids = await _write_turns(saver_env, steps=3) + old_head_id = checkpoint_ids[-1] + + graph = _build_graph(FullState, saver_env.saver) + fork_message = HumanMessage(content="regenerated turn: " + "z" * 256, id="fork-turn") + await graph.ainvoke({"messages": [fork_message]}, _config_thread(thread_id, checkpoint_ids[1])) + duration_id = await _append_duration_checkpoint(saver_env, thread_id) + before = await _listed_checkpoint_ids(saver_env, thread_id) + # Both leaves are prunable under this policy, so a widened batch would show + # up as deletions rather than as an empty report. + assert {old_head_id, duration_id}.issubset(before) + + async def _no_store_access(*args: Any, **kwargs: Any) -> dict[str, int]: + raise AssertionError("validation must reject the policy before reading the store") + + monkeypatch.setattr(checkpoint_retention, "_thread_storage_stats", _no_store_access) + + with pytest.raises(ValueError, match="max_delete_per_run must be >= 0, got -1"): + await enforce_thread_retention( + saver_env.saver, + thread_id, + RetentionPolicy( + prune_leaf_sibling_branches=True, + strict_pending_write_guard=False, + max_delete_per_run=-1, + ), + ) + + monkeypatch.undo() + assert await _listed_checkpoint_ids(saver_env, thread_id) == before + + +@pytest.mark.anyio +async def test_zero_max_delete_per_run_deletes_nothing(saver_env: _SaverEnv) -> None: + """Zero is the disabled-but-valid boundary: it bounds the batch to nothing + and still reports the same before/after measurement shape.""" + thread_id, checkpoint_ids, _message_ids = await _write_turns(saver_env, steps=3) + duration_id = await _append_duration_checkpoint(saver_env, thread_id) + + report = await enforce_thread_retention( + saver_env.saver, + thread_id, + RetentionPolicy(max_delete_per_run=0), + ) + + assert report.deleted_checkpoint_ids == [] + assert duration_id in await _listed_checkpoint_ids(saver_env, thread_id) + assert report.stats_after == report.stats_before + + +@pytest.mark.anyio +async def test_empty_thread_reports_the_same_before_and_after_stats(saver_env: _SaverEnv) -> None: + """An empty (or unknown) thread is a no-op, but the report contract still + holds: both measurement halves are present and identical, so a caller + aggregating stats does not have to special-case "nothing to classify".""" + thread_id = _thread_id() + + report = await enforce_thread_retention(saver_env.saver, thread_id) + + assert report.deleted_checkpoint_ids == [] + assert report.protected_head_ids == {} + assert report.stats_before == report.stats_after + assert set(report.stats_before) == { + "logical_checkpoint_bytes", + "logical_write_bytes", + "checkpoint_rows", + "checkpoint_bytes", + "blob_rows", + "blob_bytes", + "write_rows", + "write_bytes", + } + + without_stats = await enforce_thread_retention(saver_env.saver, thread_id, collect_stats=False) + assert without_stats.stats_before == {} + assert without_stats.stats_after == {} + + +# --------------------------------------------------------------------------- +# Namespaced (persistent subgraph) histories +# --------------------------------------------------------------------------- + + +class _ChildState(TypedDict): + value: int + + +def _increment(state: _ChildState) -> dict[str, int]: + return {"value": state.get("value", 0) + 1} + + +def _build_nested_graph(checkpointer: Any) -> Any: + """A parent graph whose node is a subgraph that owns its own checkpoints.""" + child_builder = StateGraph(_ChildState) + child_builder.add_node("increment", _increment) + child_builder.add_edge(START, "increment") + child_builder.add_edge("increment", END) + child = child_builder.compile(checkpointer=True) + + parent_builder = StateGraph(_ChildState) + parent_builder.add_node("child", child) + parent_builder.add_edge(START, "child") + parent_builder.add_edge("child", END) + return parent_builder.compile(checkpointer=checkpointer) + + +@pytest.mark.anyio +async def test_persistent_subgraph_resume_head_survives_a_prune(saver_env: _SaverEnv) -> None: + """A persistent subgraph checkpoints under its own namespace, and ``alist`` + returns those rows alongside the parent's. Protecting only one global head + makes the child's latest checkpoint look like an off-chain sibling leaf, so + the opt-in E2 shape deletes it and the child's saved state rolls back to the + preceding checkpoint. Each namespace's resume head must be protected.""" + graph = _build_nested_graph(saver_env.saver) + thread_id = _thread_id() + await graph.ainvoke({"value": 10}, _config(thread_id)) + + namespaces = {(tuple_.config["configurable"].get("checkpoint_ns") or "") async for tuple_ in saver_env.saver.alist(_config(thread_id), limit=None)} + child_namespaces = sorted(ns for ns in namespaces if ns) + assert child_namespaces, "the persistent subgraph must write its own namespace" + child_ns = child_namespaces[0] + + child_config = {"configurable": {"thread_id": thread_id, "checkpoint_ns": child_ns}} + child_head = await saver_env.saver.aget_tuple(child_config) + assert child_head is not None + child_head_id = child_head.checkpoint["id"] + assert child_head.checkpoint["channel_values"]["value"] == 11 + + report = await enforce_thread_retention( + saver_env.saver, + thread_id, + RetentionPolicy(prune_leaf_sibling_branches=True), + ) + + assert child_head_id not in report.deleted_checkpoint_ids + resumed = await saver_env.saver.aget_tuple(child_config) + assert resumed is not None + assert resumed.checkpoint["id"] == child_head_id + assert resumed.checkpoint["channel_values"]["value"] == 11 + listed = {tuple_.checkpoint["id"] async for tuple_ in saver_env.saver.alist(_config(thread_id), limit=None)} + assert child_head_id in listed