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

test(checkpoint): retention deletion contract + growth baseline (#4189 item 3) (#5255)

Browse Source 
* 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>
This commit is contained in:
zeng-bohan 2026-09-08 19:21:57 +08:00 committed by GitHub
parent 9fda432ba1
commit a2808e8292
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 952 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

110
backend/docs/checkpoint-retention-contract.md Normal file
View File

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

189
backend/scripts/benchmark/checkpoint/bench_tool_result_probe.py Normal file
View File

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

58
backend/tests/test_bench_tool_result_probe_cleanup.py Normal file
View File

@ -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-*"))

595
backend/tests/test_checkpoint_retention_contract.py Normal file
View File

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