deer-flow/backend/app/gateway/checkpoint_retention.py

"""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:<task>``. 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