github/deer-flow
1
0
Fork 0
mirror of https://github.com/bytedance/deer-flow.git synced 2026-09-19 19:16:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(gateway): checkpoint retention service on the #4189 deletion contract (#5308)

Browse Source 
* 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 <zengbh1@gmail.com>

* 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 <zengbh1@gmail.com>
Co-authored-by: zengbohan1 <310902929+zengbohan1@users.noreply.github.com>
Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
This commit is contained in:
zeng-bohan 2026-09-18 16:46:19 +08:00 committed by GitHub
parent d540be7e21
commit ce3e64242b
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 1223 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

519
backend/app/gateway/checkpoint_retention.py Normal file
View File

@ -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:<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

35
backend/docs/checkpoint-retention-contract.md
View File

@ -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

669
backend/tests/test_checkpoint_retention_service.py Normal file
View File

@ -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