From a2808e82923a71666a5a36a7590bdc047e221b48 Mon Sep 17 00:00:00 2001 From: zeng-bohan Date: Tue, 8 Sep 2026 19:21:57 +0800 Subject: [PATCH] test(checkpoint): retention deletion contract + growth baseline (#4189 item 3) (#5255) * test(checkpoint): retention deletion contract + growth baseline Six contract scenarios x memory/sqlite/postgres pin what retention deletions must never break (branch ancestors, explicit resume targets, pending writes, duration-only chain links), prove the two safe shapes (leaf sibling branches, trailing duration leaves), record the full-vs-delta growth baseline in the normalized bench shape, and add an item 4 probe showing the default ToolOutputBudgetMiddleware already externalizes oversized tool results. Refs #4189 * test(checkpoint): make the retention contract load-bearing per review Review findings from willem-bd and Ricky-7-Yan: - scenario D pins its own row: before/after stats delta plus a serde round-trip of the stored write, instead of an always-true > 0 check - _delete_checkpoint now performs the joint delete the doc mandates (checkpoint row + writes rows + blobs unreachable from surviving checkpoints), so E1/E2 exercise the shape they prescribe - E1 builds the real runtime duration shape via persist_run_durations (parent dict clone, fresh id/ts, real metadata), which surfaces the shared-version case: the leaf's blobs are the surviving parent's rows - contract doc: blob reachability must be computed from surviving checkpoints in a whole-thread pass; shared-version/duration-only hazard called out explicitly; memory data model includes saver.blobs - _stats counts memory blob rows and returns the full normalized shape (logical byte totals included) - probe: drops the unused middleware/outputs_dir graph parameters and discloses the manual-harness scope limit in the module docstring - E1/E2 assert default head resolution (protected set item 5); unused graph_for helper and DURATION_ONLY_METADATA stand-in removed Signed-off-by: zengbohan1 <310902929+zengbohan1@users.noreply.github.com> * fix(checkpoint): scope probe cleanup to owned dirs, key report by backend Second-round review findings on #5255: - [P1] bench_tool_result_probe.py removed the whole user-supplied --outputs-dir (and the shared .probe-tmp) in its finally block, so pre-existing files were deleted on success and failure alike. The run now writes into (and removes) a fresh owned probe-run-* child beneath the requested directory, and SQLite databases live in a unique mkdtemp'd temp directory that is removed with the run. Regression tests pin that unrelated pre-existing files survive both a successful and a simulated failing run. - [P2] the optional retention report keyed every backend's measurements under one shared name, so a multi-backend invocation kept only the last backend's numbers. _report() now takes the parameterized backend explicitly (saver_env.kind); regression pins that memory and sqlite entries coexist in one report file. Signed-off-by: zengbohan1 <310902929+zengbohan1@users.noreply.github.com> --------- Signed-off-by: zengbohan1 <310902929+zengbohan1@users.noreply.github.com> Co-authored-by: zengbohan1 <310902929+zengbohan1@users.noreply.github.com> --- backend/docs/checkpoint-retention-contract.md | 110 ++++ .../checkpoint/bench_tool_result_probe.py | 189 ++++++ .../test_bench_tool_result_probe_cleanup.py | 58 ++ .../test_checkpoint_retention_contract.py | 595 ++++++++++++++++++ 4 files changed, 952 insertions(+) create mode 100644 backend/docs/checkpoint-retention-contract.md create mode 100644 backend/scripts/benchmark/checkpoint/bench_tool_result_probe.py create mode 100644 backend/tests/test_bench_tool_result_probe_cleanup.py create mode 100644 backend/tests/test_checkpoint_retention_contract.py diff --git a/backend/docs/checkpoint-retention-contract.md b/backend/docs/checkpoint-retention-contract.md new file mode 100644 index 000000000..50e34673a --- /dev/null +++ b/backend/docs/checkpoint-retention-contract.md @@ -0,0 +1,110 @@ +# Checkpoint Retention Contract (DRAFT) + +Status: **draft** — the deletion contract for #4189 item 3. No retention or +deletion implementation should land before this contract (or a successor +revision of it) is accepted, and every deletion proposal must be validated +against `backend/tests/test_checkpoint_retention_contract.py`. + +## Why a contract is needed + +LangGraph checkpoints form a per-thread **parent chain**. Gateway features +depend on that chain being intact: + +- **Branch / regenerate** resolves the replay base by walking + `parent_config` links from a checkpoint that contains the target message + (`app/gateway/checkpoint_lineage.py::find_checkpoint_before_message`). +- **Explicit resume** replays from a `checkpoint_id` a client still holds. + +Deleting checkpoint rows by recency or table size can therefore break those +features **silently** — a missing ancestor surfaces as +`CheckpointLineageIntegrityError` at branch time, or as a lost resume target, +never as an obvious storage bug. The contract below separates deletable rows +from protected rows and pins the verification method. + +## Data model + +| Backend | State rows | Writes rows | +| --------- | ----------------------- | ------------------ | +| SQLite | `checkpoints` | `writes` | +| Postgres | `checkpoints`, `checkpoint_blobs` | `checkpoint_writes` | +| Memory | `saver.storage`, `saver.blobs` | `saver.writes` | + +(Note: SQLite has no separate blob table; channel values live inside the +serialized checkpoint payload. Postgres splits blobs out.) + +Measurement shape: per-thread rows + bytes per table, normalized by +`bench_channels._normalized_storage_stats`. + +## Protected set (MUST NOT delete without the stated compensation) + +1. **Explicit resume targets** — any `checkpoint_id` a client may still + resume to. Deleting it removes the replay surface + (`test_deleting_explicit_resume_target_breaks_resume`). A retention policy + may expire these, but only with an explicit TTL semantic agreed here. +2. **Branch ancestors** — every checkpoint on the parent chain from a + branchable head back to (and including) the checkpoint *before* the oldest + branchable message. Deleting any node on that walk breaks branch/regenerate + with `CheckpointLineageIntegrityError` + (`test_deleting_branch_ancestor_breaks_lineage_loudly`). +3. **Pending writes** — rows in the writes table are uncommitted/in-flight + state, not garbage (`test_pending_writes_are_retained_state_not_garbage`). +4. **Duration-only chain links** — `persist_run_durations` appends + metadata-only checkpoints. A duration-only checkpoint that a later run has + forked from is a *chain link*: the walk steps through it, so deleting it + requires **grafting** the fork onto the grandparent (rewriting the fork's + `parent_config`) in the same change. A bare leaf (below) is safe; a link is + not. The link shape can only be produced by the real runtime, so the graft + path is specified here and intentionally not covered by a storage-level + test. +5. **Latest resumable state per thread** — the newest checkpoint must remain + addressable so a thread can always continue. + +## Provably safe forms (validated by tests) + +1. **Leaf sibling branches** — a checkpoint forked off an older turn that has + no children (`test_leaf_sibling_branch_deletion_is_safe`). Pruning it does + not affect the main line's walk, explicit resume, or head. +2. **Trailing duration-only leaves** — a duration-only checkpoint no later run + has forked from (`test_leaf_duration_checkpoint_deletion_is_safe`). + +New deletion proposals must add their shape as a test here: construct the +chain, delete, then verify (a) latest resume, (b) explicit `checkpoint_id` +resume, (c) branch from an older visible turn, and (d) orphan row counts. + +## Deletion mechanics + +- Deletion must cover the backend's tables jointly and account for orphans, and + blob reachability must be computed from the **surviving checkpoints in a + whole-thread pass**: after deleting a checkpoint row, a `checkpoint_blobs` / + `checkpoint_writes` row is an orphan only if *no surviving checkpoint* + references it. The shared-version case is not hypothetical — the real + duration-only checkpoint is a copy of the head checkpoint dict + (`persist_run_history_metadata` replaces only id/ts), so it inherits the + parent's `channel_versions` verbatim, and on Postgres the blob rows + reachable from the deleted duration row are the same rows backing its + parent. An implementation that deletes blobs keyed by the removed + checkpoint's own `channel_versions` would corrupt the thread's newest + surviving state — exactly the failure class this contract exists to + prevent. (For the same reason a real duration-only leaf is not + payload-free: it materializes the parent's values under + `{"writes": {"runtime_run_duration": {...}}, "source": "update", "step": + ...}` metadata, which is what makes reclaiming it worthwhile.) +- Failure semantics: if a proposed deletion cannot be proven safe against the + protected set, it must not ship. Partial deletion that leaves a dangling + `parent_config` converts a cleanup into a thread-level outage (branch and + regenerate fail loudly for every later turn). +- Measurement first: proposals must include before/after numbers from + `scripts/benchmark/checkpoint/bench_channels.py` (per-thread rows/bytes, + SQLite and Postgres) plus the contract test suite passing. + +## Item 4 note (large tool results) + +`ToolOutputBudgetMiddleware` externalizes oversized tool outputs before they +reach state (preview + file reference under `.tool-results/`), so the +"50 KB result re-snapshotted every step" scenario from the original report +depends on which tools/paths bypass it. The probe +(`scripts/benchmark/checkpoint/bench_tool_result_probe.py`) measures the +on-disk checkpoint delta for the wrapped vs unwrapped paths on the lead +graph; subagent chains instantiate the same middleware by default. Any PR +claiming a residual gap must name the concrete bypassing path and show its +probe numbers. diff --git a/backend/scripts/benchmark/checkpoint/bench_tool_result_probe.py b/backend/scripts/benchmark/checkpoint/bench_tool_result_probe.py new file mode 100644 index 000000000..03cbd7e00 --- /dev/null +++ b/backend/scripts/benchmark/checkpoint/bench_tool_result_probe.py @@ -0,0 +1,189 @@ +#!/usr/bin/env python +"""Probe: do oversized tool results reach checkpoint state as full text? (#4189 item 4) + +`ToolOutputBudgetMiddleware` is registered by default and externalizes tool +results above `externalize_min_chars` (preview + file reference under +`.tool-results/`). This probe quantifies the storage effect of that +transformation: the same oversized tool result is driven through the +middleware's `awrap_tool_call` (or run raw), the resulting ToolMessage is +written into a checkpointed graph state, and the per-thread checkpoint +storage (rows + bytes, same normalized shape as bench_channels) is reported +for SQLite. + +Scope limit (state this when citing the numbers): the middleware is invoked +manually and the resulting message is injected with ``aupdate_state`` — the +production agent-factory path (middleware stack wiring, ThreadDataMiddleware +runtime state, tool-node task writes) is NOT exercised. The probe therefore +bounds the middleware's own transformation and the checkpoint cost of its +output; by itself it cannot establish that "the default configuration +covers item 4". A residual gap claim must name the concrete factory path +and come with its own measurements. + +Item 4 of #4189 can be closed as covered if the wrapped path's checkpoint +bytes stay flat as the result size grows and no factory-path measurement +shows full text landing in state. + +Usage: + cd backend + python scripts/benchmark/checkpoint/bench_tool_result_probe.py \ + [--result-bytes 50000] [--outputs-dir .tool-results-probe] + +--outputs-dir may contain unrelated files: the probe creates and removes +only its own ``probe-run-*`` child inside it (externalized samples land +there), while SQLite databases go to a unique per-run temp directory. + +Output: JSON on stdout. +""" + +from __future__ import annotations + +import argparse +import asyncio +import json +import tempfile +from pathlib import Path +from types import SimpleNamespace +from typing import Annotated, Any, TypedDict +from uuid import uuid4 + +from langchain_core.messages import AnyMessage, ToolMessage +from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver +from langgraph.graph import StateGraph +from langgraph.graph.message import add_messages + +from deerflow.agents.middlewares.tool_output_budget_middleware import ToolOutputBudgetMiddleware + +PROBE_TOOL = "probe_tool" + + +class ProbeState(TypedDict): + messages: Annotated[list[AnyMessage], add_messages] + + +async def _stats_sqlite(saver: AsyncSqliteSaver, thread_id: str) -> dict[str, int]: + async def one(sql: str) -> tuple[int, int]: + async with saver.conn.execute(sql, (thread_id,)) as cursor: + row = await cursor.fetchone() + return int(row[0]), int(row[1] or 0) + + cp = await one("SELECT COUNT(*), COALESCE(SUM(LENGTH(checkpoint) + LENGTH(metadata)), 0) FROM checkpoints WHERE thread_id = ?") + wr = await one("SELECT COUNT(*), COALESCE(SUM(LENGTH(value)), 0) FROM writes WHERE thread_id = ?") + return { + "checkpoint_rows": cp[0], + "checkpoint_bytes": cp[1], + "write_rows": wr[0], + "write_bytes": wr[1], + } + + +def _make_graph(saver: Any) -> Any: + def call_probe_tool(state: dict[str, Any]) -> dict[str, Any]: + # the oversized result never flows through this node: the probe + # injects the (possibly externalized) ToolMessage via aupdate_state + # below, so the graph only provides a checkpointed state container + return {} + + builder = StateGraph(ProbeState) + builder.add_node("tool", call_probe_tool) + builder.set_entry_point("tool") + builder.set_finish_point("tool") + return builder.compile(checkpointer=saver) + + +async def _run_path( + label: str, + result_bytes: int, + *, + wrapped: bool, + outputs_dir: Path | None, + tmp_dir: Path, +) -> dict[str, Any]: + tmp_dir.mkdir(parents=True, exist_ok=True) + async with AsyncSqliteSaver.from_conn_string(str(tmp_dir / f"probe-{label}.sqlite")) as saver: + await saver.setup() + middleware = ToolOutputBudgetMiddleware() if wrapped else None + graph = _make_graph(saver) + + thread_id = f"probe-{label}" + config = {"configurable": {"thread_id": thread_id}} + + # invoke the tool through the middleware's wrap (or raw), then persist + # the resulting ToolMessage into graph state and take a checkpoint + oversized = "A" * result_bytes + request = SimpleNamespace( + tool_call={"name": PROBE_TOOL, "id": "probe-call-1"}, + runtime=SimpleNamespace(state={"thread_data": {"outputs_path": str(outputs_dir)}} if outputs_dir else {"thread_data": None}), + ) + + async def handler(_request: Any) -> ToolMessage: + return ToolMessage(content=oversized, tool_call_id="probe-call-1") + + if wrapped and middleware is not None: + message = await middleware.awrap_tool_call(request, handler) + else: + message = await handler(request) + + await graph.aupdate_state(config, {"messages": [message]}) + stats = await _stats_sqlite(saver, thread_id) + content_chars = len(message.content) if isinstance(message.content, str) else -1 + + return { + "path": label, + "result_bytes": result_bytes, + "tool_message_content_chars": content_chars, + "externalized_file_bytes": (sum(f.stat().st_size for f in outputs_dir.rglob("*") if f.is_file()) if outputs_dir and outputs_dir.exists() else 0), + "checkpoint": stats, + } + + +async def _main(result_bytes: int, outputs_dir: Path, tmp_dir: Path) -> dict[str, Any]: + raw = await _run_path("raw-unwrapped", result_bytes, wrapped=False, outputs_dir=None, tmp_dir=tmp_dir) + externalized = await _run_path("budget-externalized", result_bytes, wrapped=True, outputs_dir=outputs_dir, tmp_dir=tmp_dir) + truncated = await _run_path("budget-truncated", result_bytes, wrapped=True, outputs_dir=None, tmp_dir=tmp_dir) + return { + "result_bytes": result_bytes, + "paths": [raw, externalized, truncated], + "verdict": { + "wrapped_content_stays_small": externalized["tool_message_content_chars"] < result_bytes, + "raw_content_is_full": raw["tool_message_content_chars"] == result_bytes, + }, + } + + +def run_probe(result_bytes: int, outputs_dir: Path, tmp_dir: Path) -> dict[str, Any]: + """Drive one probe run, cleaning up only directories this run owns. + + ``outputs_dir`` may be user-supplied and may pre-exist with unrelated + files: the run writes into (and removes) a fresh owned ``probe-run-*`` + child of it, never the directory itself or anything beside it. ``tmp_dir`` + hosts the per-run SQLite databases and is emptied by the cleanup. + """ + import shutil + + owned_outputs = outputs_dir / f"probe-run-{uuid4().hex[:8]}" + owned_outputs.mkdir(parents=True, exist_ok=False) + tmp_dir.mkdir(parents=True, exist_ok=True) + try: + return asyncio.run(_main(result_bytes, owned_outputs, tmp_dir)) + finally: + shutil.rmtree(owned_outputs, ignore_errors=True) + shutil.rmtree(tmp_dir, ignore_errors=True) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--result-bytes", type=int, default=50_000) + parser.add_argument("--outputs-dir", type=Path, default=Path(".tool-results-probe")) + args = parser.parse_args() + + outputs_dir: Path = args.outputs_dir + outputs_dir.mkdir(parents=True, exist_ok=True) + tmp_dir = Path(tempfile.mkdtemp(prefix="deerflow-probe-")) + report = run_probe(args.result_bytes, outputs_dir, tmp_dir) + + json.dump(report, __import__("sys").stdout, indent=2) + print() + + +if __name__ == "__main__": + main() diff --git a/backend/tests/test_bench_tool_result_probe_cleanup.py b/backend/tests/test_bench_tool_result_probe_cleanup.py new file mode 100644 index 000000000..806236ebe --- /dev/null +++ b/backend/tests/test_bench_tool_result_probe_cleanup.py @@ -0,0 +1,58 @@ +"""The probe may only clean directories it owns (review finding on #5255). + +``--outputs-dir`` is user-supplied and may pre-exist with unrelated files: +on success or failure, the probe must remove only its own per-run +``probe-run-*`` child (plus its unique temp directory), never the requested +directory itself or anything beside it. +""" + +from __future__ import annotations + +import importlib.util +import sys +from pathlib import Path + +import pytest + + +def _load_probe_module() -> object: + path = Path(__file__).resolve().parents[1] / "scripts" / "benchmark" / "checkpoint" / "bench_tool_result_probe.py" + spec = importlib.util.spec_from_file_location("bench_tool_result_probe_under_test", path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +def test_successful_run_leaves_unrelated_files_alone(tmp_path: Path) -> None: + module = _load_probe_module() + outputs_dir = tmp_path / "outputs" + outputs_dir.mkdir() + (outputs_dir / "user-file.txt").write_text("keep me", encoding="utf-8") + tmp_dir = tmp_path / "scratch" + + report = module.run_probe(20_000, outputs_dir, tmp_dir) + + assert report["verdict"]["raw_content_is_full"] is True + assert report["verdict"]["wrapped_content_stays_small"] is True + assert (outputs_dir / "user-file.txt").read_text(encoding="utf-8") == "keep me", "unrelated content must survive" + assert not list(outputs_dir.glob("probe-run-*")), "the owned child must be removed" + assert not tmp_dir.exists(), "the owned temp directory must be removed with the run" + + +def test_failing_run_still_leaves_unrelated_files_alone(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + module = _load_probe_module() + outputs_dir = tmp_path / "outputs" + outputs_dir.mkdir() + (outputs_dir / "user-file.txt").write_text("keep me", encoding="utf-8") + + def boom(*args: object, **kwargs: object) -> None: + raise RuntimeError("simulated probe failure") + + monkeypatch.setattr(module, "_main", boom) + with pytest.raises(RuntimeError): + module.run_probe(20_000, outputs_dir, tmp_path / "scratch") + + assert (outputs_dir / "user-file.txt").read_text(encoding="utf-8") == "keep me", "cleanup on failure must still be scoped to the owned child" + assert not list(outputs_dir.glob("probe-run-*")) diff --git a/backend/tests/test_checkpoint_retention_contract.py b/backend/tests/test_checkpoint_retention_contract.py new file mode 100644 index 000000000..8dcf1e5cf --- /dev/null +++ b/backend/tests/test_checkpoint_retention_contract.py @@ -0,0 +1,595 @@ +"""Executable contract for checkpoint retention: what may be deleted, what may not. + +Companion to ``docs/checkpoint-retention-contract.md`` and the #4189 item 3 +design discussion. LangGraph checkpoints form a per-thread parent chain, so a +deletion that looks harmless by recency can silently break branch/regenerate +(``find_checkpoint_before_message`` raises ``CheckpointLineageError`` when a +parent link is no longer addressable) or explicit ``checkpoint_id`` resume. +Each test pins one side of that boundary: + +- growth baseline: per-step rows/bytes across the LangGraph tables, full vs + delta, in the same normalized shape as ``bench_channels``; +- branch ancestor: deleting the checkpoint a branch point depends on must + fail *loudly* (integrity error), never silently; +- explicit resume: deleting a referenced ``checkpoint_id`` removes the + ability to resume to it; +- pending writes: uncommitted writes are retained state, not garbage; +- duration-only checkpoints: the runtime appends metadata-only checkpoints + (``persist_run_durations``); one *inside* a lineage is a chain link the + walk relies on, so blind deletion breaks the walk loudly; +- leaf sibling branch: a checkpoint forked off an older turn (the production + branch path) can be deleted without affecting the main line — the one + proven-safe deletion shape so far. + +All contracts run against InMemorySaver, AsyncSqliteSaver, and — when +``TEST_POSTGRES_URI`` is set — AsyncPostgresSaver, mirroring +``test_delta_channel_checkpointers.py``. +""" + +from __future__ import annotations + +import json +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.channels import DeltaChannel +from langgraph.checkpoint.memory import InMemorySaver +from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver +from langgraph.graph import StateGraph +from langgraph.graph.message import add_messages + +from app.gateway.checkpoint_lineage import ( + CheckpointLineageError, + find_checkpoint_before_message, +) +from deerflow.agents.thread_state import merge_message_writes +from deerflow.runtime.runs.worker import persist_run_durations + + +class FullState(TypedDict): + messages: Annotated[list[AnyMessage], add_messages] + + +class DeltaState(TypedDict): + messages: Annotated[ + list[AnyMessage], + DeltaChannel(merge_message_writes, snapshot_frequency=2), + ] + + +def _thread_id() -> str: + return f"retention-contract-{uuid4().hex}" + + +def _config(thread_id: str) -> dict[str, Any]: + return {"configurable": {"thread_id": thread_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: + """One saver instance over one backend (same shape as the delta contract fixture).""" + + 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-contract.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: + """Minimal checkpoint accessor for ``find_checkpoint_before_message``.""" + + 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) + + +# --------------------------------------------------------------------------- +# Normalized storage stats (same shape as bench_channels._normalized_storage_stats) +# --------------------------------------------------------------------------- + +_SQLITE_TABLES = ( + ("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 = ?"), +) + +_POSTGRES_TABLES = ( + ("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"), +) + + +def _normalized( + *, + checkpoint_rows: int, + checkpoint_bytes: int, + blob_rows: int, + blob_bytes: int, + write_rows: int, + write_bytes: int, +) -> dict[str, int]: + """Same backend-neutral shape as ``bench_channels._normalized_storage_stats``.""" + 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, + } + + +async def _stats(env: _SaverEnv, thread_id: str) -> dict[str, int]: + """Per-thread rows/bytes in the backend-neutral measurement shape. + + The memory branch must count ``saver.blobs``: InMemorySaver keeps the + serialized channel values there, so on the delta workload those rows are + the main payload and a storage-only baseline would undercount the very + growth this contract is supposed to measure. + """ + saver = env.saver + if env.kind == "memory": + 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 _normalized( + 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 env.kind == "sqlite": + stats: dict[str, int] = {} + for row_key, bytes_key, sql in _SQLITE_TABLES: + 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 + stats = {} + for row_key, bytes_key, sql in _POSTGRES_TABLES: + async with saver._cursor() as cursor: + await cursor.execute(sql, (thread_id,)) + row = await cursor.fetchone() + stats[row_key] = int(row["rows"]) + stats[bytes_key] = int(row["bytes"] or 0) + stats["logical_checkpoint_bytes"] = stats["checkpoint_bytes"] + stats["blob_bytes"] + stats["logical_write_bytes"] = stats["write_bytes"] + return stats + + +async def _surviving_channel_versions(saver: Any, thread_id: str, deleted_id: str) -> set[Any]: + """Whole-thread pass over the checkpoints that are NOT being deleted. + + Contract deletion mechanics: a row is an orphan only if no *surviving* + checkpoint references it. A real duration-only checkpoint copies its + parent's ``channel_versions`` verbatim, so the blob rows reachable from + the deleted node can be the very rows backing the surviving parent. + """ + versions: set[Any] = set() + async for tuple_ in saver.alist(_config(thread_id), limit=None): + if tuple_.checkpoint.get("id") == deleted_id: + continue + channel_versions = (tuple_.checkpoint or {}).get("channel_versions") + if isinstance(channel_versions, dict): + versions.update(channel_versions.values()) + return versions + + +async def _delete_checkpoint(env: _SaverEnv, thread_id: str, checkpoint_id: str) -> None: + """Jointly remove one checkpoint row, its writes rows, and the blob rows + exclusively owned by it — the deletion shape the contract doc mandates, + so the provably-safe scenarios exercise the same rule they prescribe.""" + saver = env.saver + survivor_versions = await _surviving_channel_versions(saver, thread_id, checkpoint_id) + if env.kind == "memory": + for namespace in saver.storage.get(thread_id, {}).values(): + namespace.pop(checkpoint_id, None) + for key in [key for key in saver.writes if key[0] == thread_id and key[2] == checkpoint_id]: + saver.writes.pop(key, None) + 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 env.kind == "sqlite": + await saver.conn.execute( + "DELETE FROM checkpoints WHERE thread_id = ? AND checkpoint_id = ?", + (thread_id, checkpoint_id), + ) + await saver.conn.execute( + "DELETE FROM writes WHERE thread_id = ? AND checkpoint_id = ?", + (thread_id, checkpoint_id), + ) + await saver.conn.commit() + 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() + orphan_versions = [row["version"] for row in rows if row["version"] not in survivor_versions] + async with saver._cursor() as cursor: + await cursor.execute( + "DELETE FROM checkpoints WHERE thread_id = %s AND checkpoint_id = %s", + (thread_id, checkpoint_id), + ) + await cursor.execute( + "DELETE FROM checkpoint_writes WHERE thread_id = %s AND checkpoint_id = %s", + (thread_id, checkpoint_id), + ) + if orphan_versions: + await cursor.execute( + "DELETE FROM checkpoint_blobs WHERE thread_id = %s AND version = ANY(%s)", + (thread_id, orphan_versions), + ) + + +async def _task_write_blobs(env: _SaverEnv, thread_id: str, task_id: str) -> list[Any]: + """Deserialize every writes row a task owns, so a scenario can prove its + row exists AND round-trips (a bare row-count can pass on rows that were + already there).""" + saver = env.saver + if env.kind == "memory": + found: list[Any] = [] + for (stored_thread, _ns, _cp_id), writes in saver.writes.items(): + if stored_thread != thread_id: + continue + for stored_task_id, _channel, typed, _path in writes.values(): + if stored_task_id == task_id: + found.append(saver.serde.loads_typed(typed)) + return found + if env.kind == "sqlite": + async with saver.conn.execute( + "SELECT type, value FROM writes WHERE thread_id = ? AND task_id = ?", + (thread_id, task_id), + ) as cursor: + rows = await cursor.fetchall() + return [saver.serde.loads_typed((row[0], row[1])) for row in rows] + async with saver._cursor() as cursor: + await cursor.execute( + "SELECT type, blob FROM checkpoint_writes WHERE thread_id = %s AND task_id = %s", + (thread_id, task_id), + ) + rows = await cursor.fetchall() + return [saver.serde.loads_typed((row["type"], row["blob"])) for row in rows] + + +def _report(name: str, data: dict[str, Any], backend: str) -> None: + """Append one scenario result to the optional JSON report file, keyed by + the parameterized backend so one multi-backend pytest invocation keeps one + entry per backend instead of overwriting a single shared key.""" + path = os.environ.get("DEERFLOW_RETENTION_REPORT") + if not path: + return + report: dict[str, Any] = {} + if os.path.exists(path): + with open(path, encoding="utf-8") as handle: + report = json.load(handle) + report.setdefault(name, {})[backend] = data + with open(path, "w", encoding="utf-8") as handle: + json.dump(report, handle, indent=2, sort_keys=True) + + +# --------------------------------------------------------------------------- +# Shared writers +# --------------------------------------------------------------------------- + + +async def _write_turns( + env: _SaverEnv, + schema: Any, + steps: int, + *, + payload_bytes: int = 256, +) -> tuple[str, list[str], list[str]]: + """Write *steps* one-message turns; return (thread_id, checkpoint_ids, message_ids).""" + graph = _build_graph(schema, env.saver) + thread_id = _thread_id() + checkpoint_ids: list[str] = [] + message_ids: list[str] = [] + for index in range(steps): + message = HumanMessage(content=f"turn {index}: " + "x" * payload_bytes, id=f"turn-{index}") + message_ids.append(message.id) + 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, message_ids + + +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, + ) + + +# --------------------------------------------------------------------------- +# Contracts +# --------------------------------------------------------------------------- + + +@pytest.mark.anyio +async def test_growth_baseline_full_vs_delta(saver_env: _SaverEnv) -> None: + """Scenario A: per-step growth is recorded for both schemas; delta must not regress. + + Full mode re-snapshots cumulative messages every step; delta mode appends + writes and only snapshots every ``snapshot_frequency`` steps. + """ + measurements: dict[str, list[dict[str, int]]] = {} + for schema_name, schema in (("full", FullState), ("delta", DeltaState)): + graph = _build_graph(schema, saver_env.saver) + thread_id = _thread_id() + series: list[dict[str, int]] = [] + for index in range(4): + message = HumanMessage(content=f"turn {index}: " + "y" * 512, id=f"turn-{index}") + await graph.ainvoke({"messages": [message]}, _config(thread_id)) + series.append(await _stats(saver_env, thread_id)) + measurements[schema_name] = series + + rows = [sample["checkpoint_rows"] for sample in series] + assert rows == sorted(rows), f"{schema_name} checkpoint rows must be non-decreasing: {rows}" + + # storage-shape contract: delta mode carries per-step payloads in the + # writes table (snapshotted only every snapshot_frequency), while full + # mode re-snapshots everything into the checkpoints payload. Absolute + # byte comparisons are cadence- and backend-dependent — the report above + # is what feeds the retention design, these assertions pin the shape. + assert measurements["delta"][-1]["write_rows"] > 0, "delta mode must land per-step payloads in writes" + _report("growth_baseline", measurements, saver_env.kind) + + +@pytest.mark.anyio +async def test_deleting_branch_ancestor_breaks_lineage_loudly(saver_env: _SaverEnv) -> None: + """Scenario B: a checkpoint an older turn's branch depends on cannot be silently removed. + + Regenerate/branch resolves the replay base by walking the parent chain from + the head. Deleting the chain node the branch point needs must surface as + ``CheckpointLineageError`` — never as a wrong-but-plausible replay base. + """ + thread_id, checkpoint_ids, message_ids = await _write_turns(saver_env, FullState, steps=4) + head_config = _config(thread_id) + + base = await _walk(saver_env, head_config, message_ids[1]) + assert base is not None + branch_point_id = base.config["configurable"]["checkpoint_id"] + + await _delete_checkpoint(saver_env, thread_id, branch_point_id) + + with pytest.raises(CheckpointLineageError): + await _walk(saver_env, head_config, message_ids[1]) + _report("branch_ancestor_deletion", {"deleted": branch_point_id}, saver_env.kind) + + +@pytest.mark.anyio +async def test_deleting_explicit_resume_target_breaks_resume(saver_env: _SaverEnv) -> None: + """Scenario C: a ``checkpoint_id`` someone may resume to is part of the protected set.""" + thread_id, checkpoint_ids, _message_ids = await _write_turns(saver_env, FullState, steps=4) + target_id = checkpoint_ids[1] + + before = await saver_env.saver.aget_tuple({"configurable": {"thread_id": thread_id, "checkpoint_id": target_id}}) + assert before is not None + + await _delete_checkpoint(saver_env, thread_id, target_id) + + after = await saver_env.saver.aget_tuple({"configurable": {"thread_id": thread_id, "checkpoint_id": target_id}}) + assert after is None, "resume to a deleted checkpoint_id must fail, not silently fall back" + + +@pytest.mark.anyio +async def test_pending_writes_are_retained_state_not_garbage(saver_env: _SaverEnv) -> None: + """Scenario D: uncommitted writes are visible state; their rows are protected.""" + thread_id, checkpoint_ids, _message_ids = await _write_turns(saver_env, FullState, steps=2) + + write = ("messages", b"pending-write") + latest_config = {"configurable": {"thread_id": thread_id, "checkpoint_ns": "", "checkpoint_id": checkpoint_ids[-1]}} + before = await _stats(saver_env, thread_id) + await saver_env.saver.aput_writes(latest_config, [write], task_id="pending-task") + after = await _stats(saver_env, thread_id) + assert after["write_rows"] == before["write_rows"] + 1, "put_writes must land exactly its own row" + blobs = await _task_write_blobs(saver_env, thread_id, "pending-task") + assert blobs == [b"pending-write"], "the stored write must round-trip byte-identically" + _report("pending_writes", {"stats": after}, saver_env.kind) + + +@pytest.mark.anyio +async def test_leaf_duration_checkpoint_deletion_is_safe(saver_env: _SaverEnv) -> None: + """Scenario E1 via the real runtime writer: a trailing duration-only leaf + can be deleted. + + ``persist_run_durations`` appends the shape production actually writes: + a copy of the head checkpoint dict (``channel_values``/``channel_versions`` + verbatim) with a fresh id/ts and metadata ``{"writes": + {"runtime_run_duration": {...}}, "source": "update", "step": ...}``. The + leaf therefore materializes the parent's payload, and on version-deduped + backends its blobs are the *same rows* backing the surviving parent — the + joint delete must leave them alone. A duration-only checkpoint that a + later run has forked from is instead a chain link; deleting that shape + requires grafting the fork onto the grandparent (contract doc) and is not + exercised here. + """ + thread_id, checkpoint_ids, message_ids = await _write_turns(saver_env, FullState, steps=3) + stats_before = await _stats(saver_env, thread_id) + + written = await persist_run_durations(checkpointer=saver_env.saver, thread_id=thread_id, durations={"run-1": 7}) + assert written, "the real duration writer must append its metadata-only checkpoint" + head = await saver_env.saver.aget_tuple(_config(thread_id)) + duration_id = head.checkpoint["id"] + assert duration_id not in checkpoint_ids + stats_after_append = await _stats(saver_env, thread_id) + # the real clone materializes the parent payload; version-deduped storage + # must not grow blob rows when it lands (sqlite has no blob table: 0 == 0) + assert stats_after_append["blob_rows"] == stats_before["blob_rows"] + + # a trailing metadata-only leaf can be dropped (a cleanup that prunes + # trailing duration checkpoints) without affecting the run's lineage + await _delete_checkpoint(saver_env, thread_id, duration_id) + + # the run's final checkpoint still resolves its lineage and stays + # explicitly addressable + base = await _walk(saver_env, _config_thread(thread_id, checkpoint_ids[-1]), "turn-2") + assert base is not None + resumed = await saver_env.saver.aget_tuple(_config_thread(thread_id, checkpoint_ids[-1])) + assert resumed is not None + # protected set item 5: the next turn resolves the head without an id + default_head = await saver_env.saver.aget_tuple(_config(thread_id)) + assert default_head.checkpoint["id"] == checkpoint_ids[-1] + # shared-version safety: every blob backing the surviving parent survives + stats_after_delete = await _stats(saver_env, thread_id) + assert stats_after_delete["blob_rows"] == stats_after_append["blob_rows"] + assert stats_after_delete["checkpoint_rows"] == stats_before["checkpoint_rows"] + _report("leaf_duration_deletion", {"deleted": duration_id, "head": checkpoint_ids[-1]}, saver_env.kind) + + +def _config_thread(thread_id: str, checkpoint_id: str) -> dict[str, Any]: + return {"configurable": {"thread_id": thread_id, "checkpoint_id": checkpoint_id}} + + +@pytest.mark.anyio +async def test_leaf_sibling_branch_deletion_is_safe(saver_env: _SaverEnv) -> None: + """Scenario E2: a forked-off leaf branch can be deleted without touching the main line. + + This is the one deletion shape proven safe so far: create a real branch by + resuming from an older checkpoint and writing a new turn (the production + branch path), then delete the resulting leaf checkpoint. The original + head, its lineage walk, and explicit resume all keep working. + """ + thread_id, checkpoint_ids, message_ids = await _write_turns(saver_env, FullState, steps=4) + original_head_id = checkpoint_ids[-1] + + # fork a real branch from turn 1 via the production path (resume + write) + graph = _build_graph(FullState, saver_env.saver) + fork_config = _config_thread(thread_id, checkpoint_ids[1]) + fork_message = HumanMessage(content="fork turn: " + "z" * 256, id="fork-turn") + await graph.ainvoke({"messages": [fork_message]}, fork_config) + # the fork leaf is the newest checkpoint on the thread; querying with the + # fork config would return the *source* checkpoint instead + fork_state = await graph.aget_state(_config(thread_id)) + fork_checkpoint_id = fork_state.config["configurable"]["checkpoint_id"] + assert fork_checkpoint_id not in (original_head_id, checkpoint_ids[1]) + + await _delete_checkpoint(saver_env, thread_id, fork_checkpoint_id) + + # the main line is untouched: the lineage walk still resolves (the forked + # checkpoint was a leaf), and protected set item 5 holds — default head + # resolution stays addressable, landing on the deleted leaf's surviving + # parent rather than the deleted id + base = await _walk(saver_env, _config_thread(thread_id, original_head_id), message_ids[0]) + assert base is not None + default_head = await saver_env.saver.aget_tuple(_config(thread_id)) + assert default_head is not None + assert default_head.checkpoint["id"] != fork_checkpoint_id + resumed = await saver_env.saver.aget_tuple(_config_thread(thread_id, original_head_id)) + assert resumed is not None + _report("leaf_sibling_deletion", {"fork": fork_checkpoint_id, "head": original_head_id}, saver_env.kind) + + +def test_report_keeps_one_entry_per_backend(tmp_path: Any, monkeypatch: Any) -> None: + """The report must retain one entry per parameterized backend: memory and + SQLite results coexist in the same file instead of overwriting a shared + key (which silently discarded the memory baseline).""" + monkeypatch.setenv("DEERFLOW_RETENTION_REPORT", str(tmp_path / "report.json")) + _report("growth_baseline", {"checkpoint_rows": 7}, "memory") + _report("growth_baseline", {"checkpoint_rows": 9}, "sqlite") + with open(tmp_path / "report.json", encoding="utf-8") as handle: + data = json.load(handle) + assert data["growth_baseline"] == {"memory": {"checkpoint_rows": 7}, "sqlite": {"checkpoint_rows": 9}}