From c6ad5ed93c388732cbda64447d058b7e8f35ad80 Mon Sep 17 00:00:00 2001 From: zeng-bohan Date: Fri, 9 Oct 2026 20:30:19 +0800 Subject: [PATCH] fix(gateway): classify duration-only checkpoints without the Postgres-dropped marker (#5988) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(gateway): classify duration-only checkpoints without the Postgres-dropped marker On Postgres the savers funnel metadata through langgraph's get_serializable_checkpoint_metadata, which pops writes before the row lands, so persist_run_history_metadata's runtime_run_duration marker never comes back and every replay/history read path treats metadata-only duration leaves as addressable conversation state. is_duration_only_checkpoint now classifies in two tiers: - writes present (memory/SQLite) is authoritative — the per-writer key is exact, so a title leaf carrying runtime_interrupt_title with inherited run_durations is not a duration checkpoint; - writes absent (Postgres) falls back to the surviving stamps (source == run_durations / run_message_ids keys). Stamps alone are not exact: _ensure_interrupted_title and _rollback_to_pre_run_checkpoint copy the head's metadata and inherit them. Read paths tolerate that because every stamp-bearing leaf either copies its parent's messages verbatim or restores an ancestor's exact message state, so a skipped leaf only advances a message-presence scan onto a state-equivalent checkpoint. Destructive callers opt out via stamp_fallback=False: enforce_thread_retention keeps widening the class only through the parent-shape check in _mark_duration_leaves_without_the_marker. * fix(gateway): confirm the duration stamp fallback by parent shape Review on #5988 surfaced a real false positive: write_thread_goal copies the head's metadata and thereby inherits the run_durations/run_message_ids stamps while changing the goal channel, and on Postgres its own writes.goal marker is popped too — so the stamp fallback classified the goal leaf as duration-only and the branch scan walked past it, copying channel values from before the goal write. The stamp fallback now also requires the duration writer's shape: it clones the parent's channel_versions verbatim, while every state-changing leaf (title, goal) bumps at least one version. Callers resolve the parent from their own context — the lineage walk fetches the grandparent only for stamp candidates, the history-window scans (branch, latest-turn, regenerate filter, chronological fallback) index the window they already hold — and stamps without a resolvable parent are refused: scanning one state-equivalent copy is free, skipping a goal leaf is not. Also corrects the mechanism claims from the first pass: langgraph's update_state paths (rollback, compaction, manual updates) persist metadata built from scratch and cannot inherit the stamps at all; the inheriting writers are title and goal. Co-Authored-By: Claude Sonnet 5 * fix(gateway): share the stamp gate and reuse the walked grandparent Review round 2 on #5988 — all three items, no functional changes: - extract _has_postgres_stamps as the single gate shared by the classifier fallback tier and the walk probe, so widening one widens the other; - carry the grandparent fetched for a stamp candidate into the next iteration (it is exactly the next parent), turning a chain of N Postgres duration leaves from ~2N+1 ancestor reads into N+1; - pin the unresolvable-grandparent promise at walk level: a stamp candidate whose own parent link was pruned stays addressable and remains the replay base. Co-Authored-By: Claude Sonnet 5 * fix(gateway): resolve persisted channel versions for the duration stamp fallback Materialized accessor snapshots expose no checkpoint payload, so the verbatim-copy shape check always failed open and Postgres duration leaves stayed addressable. Carry the version maps alongside the snapshots (raw snapshots expose them directly, graph-backed ones re-read the raw tuple through the accessor) and use them at every replay-base scan site. * perf(gateway): resolve stamp-fallback versions lazily and once per tuple The Postgres stamp fallback needs the raw channel versions for its shape check, and both window scans resolved them eagerly for every entry: a branch request paid up to 2N serial roundtrips for a window of N, and the regenerate scan resolved the same tuples a third time through resolve_history_versions. Resolve them only for stamp candidates - the marker tier and the stamp gate decide without versions - and memoize by config identity, so a marker-decided window pays nothing and each distinct tuple pays at most one raw read. The lineage walk had the same shape: a stamp candidate the shape check refuses re-read its grandparent as the next iteration's parent, because the prefetched slot was only filled when the candidate was skipped. Ancestor reads are now memoized by config identity, so no ancestor is fetched twice. Dropped the dead checkpointer parameter of find_checkpoint_before_message_chronologically: no caller passed it, and wiring one in from an async handler would put a blocking get_tuple on the event loop. --------- Co-authored-by: zeng-bohan <310902929+zeng-bohan@users.noreply.github.com> Co-authored-by: Claude Sonnet 5 --- backend/app/gateway/checkpoint_lineage.py | 300 ++++++++++- backend/app/gateway/checkpoint_retention.py | 6 +- backend/app/gateway/routers/thread_runs.py | 19 +- backend/app/gateway/routers/threads.py | 24 +- backend/app/gateway/services.py | 4 +- .../deerflow/runtime/checkpoint_state.py | 10 + backend/tests/test_checkpoint_lineage.py | 473 +++++++++++++++++- 7 files changed, 813 insertions(+), 23 deletions(-) diff --git a/backend/app/gateway/checkpoint_lineage.py b/backend/app/gateway/checkpoint_lineage.py index 275784c7d3..79f7be510a 100644 --- a/backend/app/gateway/checkpoint_lineage.py +++ b/backend/app/gateway/checkpoint_lineage.py @@ -40,9 +40,245 @@ def checkpoint_metadata(checkpoint_tuple: Any) -> dict[str, Any]: return dict(metadata) if isinstance(metadata, dict) else {} -def is_duration_only_checkpoint(checkpoint_tuple: Any) -> bool: - writes = checkpoint_metadata(checkpoint_tuple).get("writes") - return isinstance(writes, dict) and "runtime_run_duration" in writes +def _has_postgres_stamps(metadata: dict[str, Any]) -> bool: + """The stamp-fallback gate shared by the classifier and the walk probe. + + One gate, two consumers: if the classifier's fallback tier is ever + widened, this predicate widens with it and the lineage walk keeps + fetching the grandparent that the fallback will require. Re-encoding the + gate in the probe would silently break that pairing — the walk would stop + prefetching, every Postgres duration leaf would fall back to + ``parent=None``, and the refusal tier would return the original bug. + """ + + if metadata.get("source") != "update": + return False + return "run_durations" in metadata or "run_message_ids" in metadata + + +def is_duration_only_checkpoint(checkpoint_tuple: Any, *, parent: Any | None = None, stamp_fallback: bool = True, versions: dict[str, Any] | None = None, parent_versions: dict[str, Any] | None = None) -> bool: + """Return whether the tuple is a metadata-only run-duration checkpoint. + + ``persist_run_history_metadata`` (``deerflow.runtime.runs.worker``) is the + only writer of these checkpoints. It stamps + ``metadata["writes"]["runtime_run_duration"]``, and the memory and SQLite + savers round-trip that marker unchanged — so whenever ``writes`` is + present it is authoritative: its per-writer key classifies the leaf + exactly (a title or goal leaf carries ``runtime_interrupt_title`` / + ``goal``, a graph step carries its channel writes, and none of them is a + duration checkpoint). + + The Postgres savers instead funnel metadata through langgraph's + ``get_serializable_checkpoint_metadata``, which pops ``writes`` before + the row lands (``langgraph/checkpoint/postgres/aio.py``; no other saver + calls it) — so on Postgres the marker never comes back and every duration + checkpoint would look like an addressable state to the replay and history + paths. The writer's index keys (``run_durations`` and + ``run_message_ids``) and ``source == "update"`` do survive that round + trip, but they are not proof on their own: ``_ensure_interrupted_title`` + and ``write_thread_goal`` copy the head's metadata and thereby inherit + the stamps while changing a real channel. The fallback therefore also + requires the duration writer's shape — it clones the parent's + ``channel_versions`` verbatim (``id``/``ts`` only), while a state-changing + leaf bumps at least one version — which is what passing ``parent`` + enables. Without a resolvable parent the stamps are refused: the cost of + scanning one metadata-only copy is a state-equivalent base, the cost of + skipping a goal leaf is the goal itself. + + Langgraph's ``update_state`` paths (rollback restore, compaction, manual + updates) cannot inherit the stamps at all: they persist metadata built + from scratch (``{"source": "update", "step", "parents", ...}``), and + ``get_checkpoint_metadata`` merges only scalar values from config + metadata, so dict-valued stamps never leak in either. + + Destructive callers must not rely on the fallback — pass + ``stamp_fallback=False`` and confirm by shape instead, as + ``checkpoint_retention._mark_duration_leaves_without_the_marker`` does. + """ + + metadata = checkpoint_metadata(checkpoint_tuple) + writes = metadata.get("writes") + if isinstance(writes, dict): + return "runtime_run_duration" in writes + if not stamp_fallback or parent is None: + return False + if not _has_postgres_stamps(metadata): + return False + return _copies_parent_verbatim(checkpoint_tuple, parent, versions=versions, parent_versions=parent_versions) + + +def has_duration_stamps_without_marker(checkpoint_tuple: Any) -> bool: + """Whether the tuple reaches the Postgres stamp-fallback tier. + + Reaching the tier makes the tuple a duration-copy *candidate*, not a + verdict: title and goal leaves inherit the same stamps. Callers that can + resolve the parent tuple pass it to :func:`is_duration_only_checkpoint` + for the shape check; this predicate exists so the lineage walk only pays + the extra grandparent read when the marker tier could not decide. + """ + + metadata = checkpoint_metadata(checkpoint_tuple) + if isinstance(metadata.get("writes"), dict): + return False + return _has_postgres_stamps(metadata) + + +def _snapshot_versions(obj: Any) -> dict[str, Any] | None: + """Channel versions persisted on a tuple or snapshot, when exposed. + + Raw checkpoint tuples carry ``checkpoint.channel_versions``; the degraded + ``_RawCheckpointSnapshot`` exposes the same map as ``channel_versions``; + langgraph's materialized ``StateSnapshot`` exposes neither, so the walk + falls back to the persisted raw tuple. + """ + + checkpoint = getattr(obj, "checkpoint", None) + if isinstance(checkpoint, dict): + versions = checkpoint.get("channel_versions") + if isinstance(versions, dict): + return versions + versions = getattr(obj, "channel_versions", None) + if isinstance(versions, dict): + return versions + return None + + +def _history_identity(checkpoint_tuple: Any) -> tuple[str, str, str] | None: + return _config_identity(getattr(checkpoint_tuple, "config", {}) or {}) + + +async def resolve_history_versions( + accessor: Any, + checkpoints: Sequence[Any], + *, + cache: dict[tuple[str, str, str], Any] | None = None, +) -> dict[tuple[str, str, str], Any]: + """Persisted channel versions for a history window, keyed by config identity. + + Pass ``cache`` to share the lookups with an earlier lazy pass over the same + window: the dict is filled in place and returned, so a request that has + already resolved the stamp candidates pays nothing here. + """ + + versions: dict[tuple[str, str, str], Any] = cache if cache is not None else {} + for checkpoint_tuple in checkpoints: + identity = _history_identity(checkpoint_tuple) + if identity is not None and identity not in versions: + versions[identity] = await resolve_channel_versions(accessor, checkpoint_tuple) + return versions + + +async def resolve_channel_versions(accessor: Any, snapshot: Any) -> dict[str, Any] | None: + """Persisted ``channel_versions`` for *snapshot*, fetching raw when absent. + + Materialized accessor snapshots (``StateSnapshot``) do not carry the + checkpoint payload, so the version maps the Postgres stamp fallback needs + are re-read from the raw checkpoint tuple. The degraded raw read path + already exposes the same map, so it never needs the extra read. + """ + + versions = _snapshot_versions(snapshot) + if versions is not None: + return versions + getter = getattr(accessor, "aget_tuple", None) + config = getattr(snapshot, "config", None) + if callable(getter) and isinstance(config, dict): + try: + tup = await getter(config) + except Exception: + return None + return _snapshot_versions(tup) + return None + + +async def resolve_versions_memoized( + accessor: Any, + checkpoint_tuple: Any, + cache: dict[tuple[str, str, str], Any], +) -> dict[str, Any] | None: + """``resolve_channel_versions`` memoized by config identity. + + Newest-first windows make every tuple both a snapshot and its successor's + parent, so resolving candidates lazily without this cache still reads each + distinct tuple twice. Tuples without a config identity are resolved + directly, since there is nothing to key them by. + """ + + identity = _history_identity(checkpoint_tuple) + if identity is None: + return await resolve_channel_versions(accessor, checkpoint_tuple) + if identity not in cache: + cache[identity] = await resolve_channel_versions(accessor, checkpoint_tuple) + return cache[identity] + + +async def resolve_stamp_candidate_versions( + accessor: Any, + checkpoint_tuple: Any, + parent: Any | None, + cache: dict[tuple[str, str, str], Any], +) -> tuple[dict[str, Any] | None, dict[str, Any] | None]: + """Resolve the version maps a stamp candidate needs, and nothing else. + + :func:`is_duration_only_checkpoint` reads ``versions`` only at the + shape-check tier, so the marker tier (``writes`` present) and the stamp + gate must not pay for them: on Postgres every lookup behind a materialized + snapshot is a serial raw read. Scanning a window without this gate costs up + to two such reads per entry, and the memoized cache keeps a stamp candidate + and its parent to one read each even though both are also visited as + snapshots. + """ + + if parent is None or not has_duration_stamps_without_marker(checkpoint_tuple): + return None, None + versions = await resolve_versions_memoized(accessor, checkpoint_tuple, cache) + parent_versions = await resolve_versions_memoized(accessor, parent, cache) + return versions, parent_versions + + +def _copies_parent_verbatim(checkpoint_tuple: Any, parent: Any, *, versions: dict[str, Any] | None = None, parent_versions: dict[str, Any] | None = None) -> bool: + """Whether the leaf cloned its parent's ``channel_versions`` unchanged. + + Same rule as ``checkpoint_retention``'s shape check: the duration writer + replaces only ``id``/``ts``, while any state-changing writer bumps at + least one channel version. + """ + + versions = versions if isinstance(versions, dict) else _snapshot_versions(checkpoint_tuple) + parent_versions = parent_versions if isinstance(parent_versions, dict) else _snapshot_versions(parent) + if not isinstance(versions, dict) or not isinstance(parent_versions, dict) or not versions: + return False + return frozenset(versions.values()) == frozenset(parent_versions.values()) + + +def history_parent_index(checkpoints: Sequence[Any]) -> dict[tuple[str, str], Any]: + """Index a checkpoint window by ``(namespace, id)`` for parent lookups.""" + + index: dict[tuple[str, str], Any] = {} + for checkpoint_tuple in checkpoints: + configurable = checkpoint_configurable(checkpoint_tuple) + checkpoint_id = configurable.get("checkpoint_id") + if isinstance(checkpoint_id, str) and checkpoint_id: + index[(str(configurable.get("checkpoint_ns") or ""), checkpoint_id)] = checkpoint_tuple + return index + + +def parent_from_history_index(checkpoint_tuple: Any, index: dict[tuple[str, str], Any]) -> Any | None: + """Resolve the tuple's parent inside an indexed window, or ``None``. + + ``None`` means the parent is outside the window (or the tuple predates + parent links); the stamp fallback then refuses to classify, which keeps a + goal leaf at the window boundary addressable. + """ + + parent_config = getattr(checkpoint_tuple, "parent_config", None) + if not isinstance(parent_config, dict): + return None + configurable = parent_config.get("configurable") or {} + checkpoint_id = configurable.get("checkpoint_id") + if not isinstance(checkpoint_id, str) or not checkpoint_id: + return None + return index.get((str(configurable.get("checkpoint_ns") or ""), checkpoint_id)) def has_pending_tasks(checkpoint_tuple: Any) -> bool: @@ -126,16 +362,30 @@ async def find_checkpoint_before_message( if current_identity is not None: visited.add(current_identity) - # Each step performs one ancestor read, but normal branch/regenerate - # histories cross the target boundary within 1–3 reads. Keep max_depth as - # a conservative safety cap for valid histories with many intermediate or - # duration-only checkpoints. + # Each distinct ancestor is read once: a Postgres stamp candidate also + # needs its grandparent for the shape check, and a candidate the shape + # check refuses re-reads that grandparent as the next iteration's parent, + # so reads are memoized by config identity instead of carried by hand. + # Normal branch/regenerate histories cross the target boundary within + # 1–3 reads. Keep max_depth as a conservative safety cap for valid + # histories with many intermediate or duration-only checkpoints. + read_cache: dict[tuple[str, str, str], Any] = {} + + async def read(config: dict[str, Any]) -> Any: + identity = _config_identity(config) + if identity is not None and identity in read_cache: + return read_cache[identity] + checkpoint_tuple = await accessor.aget(config) + if identity is not None: + read_cache[identity] = checkpoint_tuple + return checkpoint_tuple + for _ in range(max_depth): parent_config = getattr(current, "parent_config", None) if not isinstance(parent_config, dict): raise CheckpointParentMissingError("Checkpoint lineage ended before the target message") - parent = await accessor.aget(parent_config) + parent = await read(parent_config) parent_identity = _checkpoint_identity(parent) requested_parent_identity = _config_identity(parent_config) if parent_identity is None or not _checkpoint_exists(parent) or (requested_parent_identity is not None and parent_identity != requested_parent_identity): @@ -145,7 +395,25 @@ async def find_checkpoint_before_message( raise CheckpointLineageIntegrityError("Checkpoint lineage contains a cycle") visited.add(parent_identity) - if is_duration_only_checkpoint(parent): + grandparent = None + versions = None + parent_versions = None + if has_duration_stamps_without_marker(parent): + # Postgres popped the marker; the stamps also sit on title and + # goal leaves, so confirm the verbatim-copy shape against the + # grandparent before skipping. An unresolvable grandparent leaves + # the parent addressable, which only ever costs one extra scanned + # state-equivalent copy. The version maps are read only on this + # branch: the marker tier and the stamp gate decide without them. + gp_config = getattr(parent, "parent_config", None) + if isinstance(gp_config, dict): + candidate = await read(gp_config) + if _checkpoint_exists(candidate): + grandparent = candidate + versions = await resolve_channel_versions(accessor, parent) + if grandparent is not None: + parent_versions = await resolve_channel_versions(accessor, grandparent) + if is_duration_only_checkpoint(parent, parent=grandparent, versions=versions, parent_versions=parent_versions): current = parent continue @@ -160,6 +428,8 @@ async def find_checkpoint_before_message( def find_checkpoint_before_message_chronologically( checkpoints: Sequence[Any], message_id: str, + *, + history_versions: dict[tuple[str, str, str], Any] | None = None, ) -> tuple[Any | None, bool]: """Return ``(replay_base, target_found)`` from newest-first history. @@ -169,11 +439,21 @@ def find_checkpoint_before_message_chronologically( checkpoint branches. Duration-only checkpoints are ignored, and only settled checkpoints (see :func:`has_pending_tasks`) with an addressable id can become the replay base. + + ``history_versions`` is the only version source: callers that resolved the + window share it here, and raw tuples still classify through their own + ``checkpoint.channel_versions`` when it is absent. A sync checkpointer + parameter used to re-read them, which put a blocking read on the event loop + for async callers, so it is gone. """ previous_checkpoint = None + history_index = history_parent_index(checkpoints) for checkpoint_tuple in reversed(checkpoints): - if is_duration_only_checkpoint(checkpoint_tuple): + parent = parent_from_history_index(checkpoint_tuple, history_index) + versions = history_versions.get(_history_identity(checkpoint_tuple)) if history_versions is not None else _snapshot_versions(checkpoint_tuple) + parent_versions = history_versions.get(_history_identity(parent)) if (history_versions is not None and parent is not None) else _snapshot_versions(parent) + if is_duration_only_checkpoint(checkpoint_tuple, parent=parent, versions=versions, parent_versions=parent_versions): continue message_ids = {_message_id(message) for message in checkpoint_messages(checkpoint_tuple)} if message_id in message_ids: diff --git a/backend/app/gateway/checkpoint_retention.py b/backend/app/gateway/checkpoint_retention.py index 2d859bee72..ef72ed3f93 100644 --- a/backend/app/gateway/checkpoint_retention.py +++ b/backend/app/gateway/checkpoint_retention.py @@ -416,7 +416,11 @@ async def enforce_thread_retention( cp_id=cp_id, parent_ns=parent_ns, parent_id=parent_id, - duration_only=is_duration_only_checkpoint(tuple_), + # Pruning is destructive, so the stamp fallback is disabled: + # title and goal leaves inherit the writer's stamps, and only + # the shape confirmation in + # _mark_duration_leaves_without_the_marker may widen the class. + duration_only=is_duration_only_checkpoint(tuple_, stamp_fallback=False), 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(), diff --git a/backend/app/gateway/routers/thread_runs.py b/backend/app/gateway/routers/thread_runs.py index 51ec817ed6..665d2d0239 100644 --- a/backend/app/gateway/routers/thread_runs.py +++ b/backend/app/gateway/routers/thread_runs.py @@ -35,7 +35,11 @@ from app.gateway.checkpoint_lineage import ( checkpoint_messages, find_checkpoint_before_message, find_checkpoint_before_message_chronologically, + history_parent_index, is_duration_only_checkpoint, + parent_from_history_index, + resolve_history_versions, + resolve_stamp_candidate_versions, ) from app.gateway.context_usage import build_context_usage from app.gateway.conversation_reader import ( @@ -142,8 +146,8 @@ async def _refresh_store_backed_run(run_mgr: Any, record: Any) -> Any: return record -def _is_duration_only_checkpoint(checkpoint_tuple: Any) -> bool: - return is_duration_only_checkpoint(checkpoint_tuple) +def _is_duration_only_checkpoint(checkpoint_tuple: Any, history_index: dict[tuple[str, str], Any], versions=None, parent_versions=None) -> bool: + return is_duration_only_checkpoint(checkpoint_tuple, parent=parent_from_history_index(checkpoint_tuple, history_index), versions=versions, parent_versions=parent_versions) def compute_run_durations(runs) -> dict[str, int]: @@ -680,12 +684,19 @@ async def _find_base_checkpoint_before_human( raise HTTPException(status_code=409, detail=_UNSAFE_REGENERATE_LINEAGE_DETAIL) from exc try: raw_checkpoints = await accessor.ahistory(base_config, limit=REGENERATE_HISTORY_RAW_SCAN_LIMIT) - checkpoints = [item for item in raw_checkpoints if not _is_duration_only_checkpoint(item)] + history_index = history_parent_index(raw_checkpoints) + version_cache: dict[tuple[str, str, str], Any] = {} + checkpoints = [] + for item in raw_checkpoints: + parent = parent_from_history_index(item, history_index) + versions, parent_versions = await resolve_stamp_candidate_versions(accessor, item, parent, version_cache) + if not _is_duration_only_checkpoint(item, history_index, versions=versions, parent_versions=parent_versions): + checkpoints.append(item) except Exception as exc: logger.exception("Failed to list checkpoints for regenerate thread %s", thread_id) raise HTTPException(status_code=500, detail="Failed to inspect checkpoint history") from exc - previous_checkpoint, target_found = find_checkpoint_before_message_chronologically(raw_checkpoints, human_message_id) + previous_checkpoint, target_found = find_checkpoint_before_message_chronologically(raw_checkpoints, human_message_id, history_versions=await resolve_history_versions(accessor, raw_checkpoints, cache=version_cache)) if target_found: if previous_checkpoint is None: raise HTTPException( diff --git a/backend/app/gateway/routers/threads.py b/backend/app/gateway/routers/threads.py index 38b57eb955..4e64eea1b0 100644 --- a/backend/app/gateway/routers/threads.py +++ b/backend/app/gateway/routers/threads.py @@ -31,7 +31,11 @@ from app.gateway.checkpoint_lineage import ( CheckpointParentMissingError, find_checkpoint_before_message, find_checkpoint_before_message_chronologically, + history_parent_index, is_duration_only_checkpoint, + parent_from_history_index, + resolve_history_versions, + resolve_stamp_candidate_versions, ) from app.gateway.deps import get_checkpointer, get_run_event_store, get_run_manager, get_run_store from app.gateway.internal_auth import get_trusted_internal_owner_user_id @@ -227,8 +231,13 @@ async def _find_branch_checkpoint( target_message_ids: set[str], ) -> Any: try: - for snapshot in await accessor.ahistory(config, limit=_BRANCH_HISTORY_RAW_SCAN_LIMIT): - if is_duration_only_checkpoint(snapshot): + history = await accessor.ahistory(config, limit=_BRANCH_HISTORY_RAW_SCAN_LIMIT) + history_index = history_parent_index(history) + version_cache: dict[tuple[str, str, str], Any] = {} + for snapshot in history: + parent = parent_from_history_index(snapshot, history_index) + versions, parent_versions = await resolve_stamp_candidate_versions(accessor, snapshot, parent, version_cache) + if is_duration_only_checkpoint(snapshot, parent=parent, versions=versions, parent_versions=parent_versions): continue if _matches_branch_target(_checkpoint_messages(snapshot), target_message_ids): return snapshot @@ -248,8 +257,13 @@ async def _branch_targets_latest_turn( ) -> bool: """Return whether the target turn is the final visible turn.""" try: - for snapshot in await accessor.ahistory(config, limit=_BRANCH_HISTORY_RAW_SCAN_LIMIT): - if is_duration_only_checkpoint(snapshot): + history = await accessor.ahistory(config, limit=_BRANCH_HISTORY_RAW_SCAN_LIMIT) + history_index = history_parent_index(history) + version_cache: dict[tuple[str, str, str], Any] = {} + for snapshot in history: + parent = parent_from_history_index(snapshot, history_index) + versions, parent_versions = await resolve_stamp_candidate_versions(accessor, snapshot, parent, version_cache) + if is_duration_only_checkpoint(snapshot, parent=parent, versions=versions, parent_versions=parent_versions): continue messages = _checkpoint_messages(snapshot) if not messages: @@ -306,7 +320,7 @@ async def _find_branch_replay_base( logger.exception("Failed to scan replay checkpoint history for thread %s", sanitize_log_param(thread_id)) raise HTTPException(status_code=500, detail="Failed to inspect checkpoint history") from exc - replay_base, target_found = find_checkpoint_before_message_chronologically(history, target_human_id) + replay_base, target_found = find_checkpoint_before_message_chronologically(history, target_human_id, history_versions=await resolve_history_versions(accessor, history)) if not target_found: logger.warning( "Could not locate branch user message %s in chronological history for thread %s", diff --git a/backend/app/gateway/services.py b/backend/app/gateway/services.py index ef57d83eb4..f88845f5a8 100644 --- a/backend/app/gateway/services.py +++ b/backend/app/gateway/services.py @@ -1296,7 +1296,7 @@ class _RawCheckpointSnapshot: metadata, config ancestry, created_at) comes straight from the tuple. """ - __slots__ = ("checkpoint_exists", "config", "values", "metadata", "parent_config", "created_at", "tasks", "tasks_known", "next") + __slots__ = ("checkpoint_exists", "config", "values", "metadata", "parent_config", "created_at", "tasks", "tasks_known", "next", "channel_versions") def __init__(self, config: dict[str, Any], tup: Any | None) -> None: self.checkpoint_exists = tup is not None @@ -1309,6 +1309,8 @@ class _RawCheckpointSnapshot: self.tasks: tuple = () self.tasks_known = False self.next: tuple = () + versions = checkpoint.get("channel_versions") + self.channel_versions = dict(versions) if isinstance(versions, dict) else None class _RawCheckpointReadAccessor: diff --git a/backend/packages/harness/deerflow/runtime/checkpoint_state.py b/backend/packages/harness/deerflow/runtime/checkpoint_state.py index ee3cbf9e35..8867858f02 100644 --- a/backend/packages/harness/deerflow/runtime/checkpoint_state.py +++ b/backend/packages/harness/deerflow/runtime/checkpoint_state.py @@ -150,6 +150,16 @@ class CheckpointStateAccessor: raise_if_snapshot_incompatible(snapshot, self.mode) return snapshot + def get_tuple(self, config: dict[str, Any]) -> Any | None: + """Raw checkpoint tuple for a config, without materializing state.""" + + return self.checkpointer.get_tuple(self._prepare_config(config)) + + async def aget_tuple(self, config: dict[str, Any]) -> Any | None: + """Raw checkpoint tuple for a config, without materializing state.""" + + return await self.checkpointer.aget_tuple(self._prepare_config(config)) + def get_metadata(self, config: dict[str, Any]) -> dict[str, Any]: """Read checkpoint metadata without materializing channel state.""" checkpoint_tuple = self.checkpointer.get_tuple(self._prepare_config(config)) diff --git a/backend/tests/test_checkpoint_lineage.py b/backend/tests/test_checkpoint_lineage.py index 10d5609245..c35938e5d1 100644 --- a/backend/tests/test_checkpoint_lineage.py +++ b/backend/tests/test_checkpoint_lineage.py @@ -17,12 +17,30 @@ from app.gateway.checkpoint_lineage import ( CheckpointParentMissingError, find_checkpoint_before_message, find_checkpoint_before_message_chronologically, + is_duration_only_checkpoint, + resolve_channel_versions, + resolve_stamp_candidate_versions, ) THREAD_ID = "thread-1" +# The metadata a Postgres row carries for a duration-only checkpoint: the +# writer's ``writes.runtime_run_duration`` stamp was dropped by langgraph's +# Postgres metadata serialisation, so only the index keys and ``source`` +# survive the round trip. +_POSTGRES_DURATION_METADATA = { + "source": "update", + "step": 3, + "run_durations": {"run-1": 12}, + "run_message_ids": {"ai-1": "run-1"}, +} -def _snapshot(checkpoint_id: str, messages: list[object], *, next_tasks: tuple[str, ...] = (), parent_id: str | None = None): +# A Postgres-shaped leaf inherits these stamps whenever a metadata-copying +# writer (title, goal) stacks on a stamped head, but bumps a channel version. +_POSTGRES_INHERITED_STAMPS = {"source": "update", "step": 4, "run_durations": {"run-1": 12}} + + +def _snapshot(checkpoint_id: str, messages: list[object], *, next_tasks: tuple[str, ...] = (), parent_id: str | None = None, metadata: dict | None = None, channel_versions: dict | None = None): parent_config = None if parent_id is not None: parent_config = {"configurable": {"thread_id": THREAD_ID, "checkpoint_ns": "", "checkpoint_id": parent_id}} @@ -36,7 +54,8 @@ def _snapshot(checkpoint_id: str, messages: list[object], *, next_tasks: tuple[s "checkpoint_map": None, } }, - metadata={}, + checkpoint={"channel_versions": dict(channel_versions or {})}, + metadata=metadata or {}, parent_config=parent_config, next=next_tasks, ) @@ -147,3 +166,453 @@ def test_unknown_pending_tasks_do_not_block_selection(): assert found is True assert base.config["configurable"]["checkpoint_id"] == "ckpt-base" + + +def test_memory_backed_duration_checkpoint_is_recognised(): + """The marker memory and SQLite savers round-trip unchanged.""" + snapshot = _snapshot( + "ckpt-duration", + [], + metadata={"source": "update", "step": 3, "writes": {"runtime_run_duration": {"run_ids": ["run-1"], "message_ids": []}}}, + ) + + assert is_duration_only_checkpoint(snapshot) is True + + +def test_postgres_duration_checkpoint_is_recognised_without_the_writes_marker(): + """Postgres serialisation pops ``writes``; stamps plus the verbatim shape identify the leaf.""" + parent = _snapshot("ckpt-real", [], channel_versions={"messages": 5}) + snapshot = _snapshot( + "ckpt-duration", + [], + parent_id="ckpt-real", + metadata=dict(_POSTGRES_DURATION_METADATA), + channel_versions={"messages": 5}, + ) + + assert is_duration_only_checkpoint(snapshot, parent=parent) is True + + +def test_stamp_fallback_refuses_stamps_without_a_parent(): + """Stamps alone are evidence, not proof — with no parent there is no verdict. + + A title or goal leaf at a window boundary carries the same stamps; the + conservative answer keeps it addressable, which only ever costs one + scanned state-equivalent copy. + """ + snapshot = _snapshot("ckpt-duration", [], metadata=dict(_POSTGRES_DURATION_METADATA), channel_versions={"messages": 5}) + + assert is_duration_only_checkpoint(snapshot) is False + assert is_duration_only_checkpoint(snapshot, parent=None) is False + + +def test_goal_leaf_inheriting_stamps_is_not_duration_only(): + """``write_thread_goal`` copies the head's metadata and changes the goal channel. + + On Postgres its own ``writes.goal`` marker is popped too, so the inherited + stamps are the only trace of the duration writer — but the bumped ``goal`` + version breaks the verbatim shape. Classifying it as duration-only made + the branch scan walk past the leaf and copy channel values from before + the goal write. + """ + parent = _snapshot("ckpt-stamped-head", [], channel_versions={"messages": 5, "goal": 2}) + goal_leaf = _snapshot( + "ckpt-goal", + [], + parent_id="ckpt-stamped-head", + metadata=dict(_POSTGRES_INHERITED_STAMPS), + channel_versions={"messages": 5, "goal": 3}, + ) + + assert is_duration_only_checkpoint(goal_leaf, parent=parent) is False + + +def test_postgres_attribution_only_duration_checkpoint_is_recognised(): + """A leaf stamped for message attribution alone has an empty ``run_durations`` map.""" + parent = _snapshot("ckpt-real", [], channel_versions={"messages": 5}) + snapshot = _snapshot( + "ckpt-duration", + [], + parent_id="ckpt-real", + metadata={"source": "update", "step": 3, "run_durations": {}, "run_message_ids": {"ai-1": "run-1"}}, + channel_versions={"messages": 5}, + ) + + assert is_duration_only_checkpoint(snapshot, parent=parent) is True + + +def test_client_update_state_is_not_a_duration_checkpoint(): + """A leaf from langgraph's ``update_state`` carries none of the writer's stamps. + + update_state persists metadata built from scratch on the pinned langgraph + (``{"source": "update", "step", "parents", ...}`` — no ``writes`` key at + all), so this shape reaches the stamp fallback tier and stays addressable + because the index keys are absent. + """ + snapshot = _snapshot("ckpt-update", [], metadata={"source": "update", "step": 3}) + + assert is_duration_only_checkpoint(snapshot) is False + + +def test_graph_step_checkpoints_are_not_duration_checkpoints(): + for source in ("input", "loop", "step"): + snapshot = _snapshot("ckpt-step", [], metadata={"source": source, "step": 1, "writes": {"messages": "..."}}) + assert is_duration_only_checkpoint(snapshot) is False + assert is_duration_only_checkpoint(_snapshot("ckpt-plain", [])) is False + assert is_duration_only_checkpoint(SimpleNamespace(values={}, config={}, metadata=None)) is False + + +def test_chronological_scan_skips_postgres_shaped_duration_checkpoints(): + """An interleaved import must not hand the replay base to a metadata-only leaf. + + The duration leaf belongs to an older incarnation whose timestamps interleave + with the live branch; without the Postgres-shape classification it would be + selected as the replay base for a message that only the newer real state and + the head contain. + """ + human = HumanMessage(id="h1", content="question") + answer = AIMessage(id="ai-1", content="answer") + history = [ + _snapshot("ckpt-x-head", [human, answer]), + _snapshot("ckpt-a-early", [], channel_versions={"messages": 3}), + _snapshot("ckpt-a-duration", [human], parent_id="ckpt-a-early", metadata=dict(_POSTGRES_DURATION_METADATA), channel_versions={"messages": 3}), + _snapshot("ckpt-a-input", [], channel_versions={"messages": 2}), + ] + + base, found = find_checkpoint_before_message_chronologically(history, "h1") + + assert found is True + assert base.config["configurable"]["checkpoint_id"] == "ckpt-a-early" + + +def test_chronological_scan_keeps_goal_leaf_addressable(): + """The imported-history fallback must not walk past a goal write either. + + The goal leaf inherits the stamps but bumped the goal version, so the + verbatim shape fails and it stays a replay-base candidate. + """ + human = HumanMessage(id="h1", content="question") + ai = AIMessage(id="ai-1", content="answer") + human2 = HumanMessage(id="h2", content="follow-up") + ai2 = AIMessage(id="ai-2", content="answer-2") + history = [ + _snapshot("ckpt-turn2-head", [human, ai, human2, ai2]), + _snapshot("ckpt-goal", [human, ai], parent_id="ckpt-turn1-tail", metadata=dict(_POSTGRES_INHERITED_STAMPS), channel_versions={"messages": 5, "goal": 3}), + _snapshot("ckpt-turn1-tail", [human, ai], channel_versions={"messages": 5}), + ] + + base, found = find_checkpoint_before_message_chronologically(history, "h2") + + assert found is True + assert base.config["configurable"]["checkpoint_id"] == "ckpt-goal" + + +def test_lineage_walk_crosses_chained_postgres_duration_checkpoints(): + """The lineage walk skips duration-only parents even when ``writes`` never comes back.""" + swapped_human = HumanMessage(id="h1__user", content="question") + raw_human = HumanMessage(id="h1", content="question") + ai = AIMessage(id="ai-1", content="answer") + history = [ + _snapshot("ckpt-head", [swapped_human, ai], parent_id="ckpt-leaf-2", channel_versions={"messages": 7}), + _snapshot("ckpt-leaf-2", [swapped_human, ai], parent_id="ckpt-leaf-1", metadata=dict(_POSTGRES_DURATION_METADATA), channel_versions={"messages": 5}), + _snapshot("ckpt-leaf-1", [swapped_human, ai], parent_id="ckpt-turn1-tail", metadata=dict(_POSTGRES_DURATION_METADATA), channel_versions={"messages": 5}), + _snapshot("ckpt-turn1-tail", [raw_human], channel_versions={"messages": 5}), + ] + accessor = _Accessor(history) + + base = asyncio.run(find_checkpoint_before_message(accessor, history[0], "h1__user", max_depth=10)) + + assert base.config["configurable"]["checkpoint_id"] == "ckpt-turn1-tail" + + +def test_lineage_walk_keeps_goal_leaf_addressable(): + """A goal write stacked on a stamped head must remain a replay-base candidate. + + The goal leaf inherits the duration stamps on Postgres but bumped the goal + version; treating it as duration-only would walk past it and hand back a + base whose channel values predate the goal write. + """ + human = HumanMessage(id="h1", content="question") + ai = AIMessage(id="ai-1", content="answer") + human2 = HumanMessage(id="h2", content="follow-up") + history = [ + _snapshot("ckpt-turn2-head", [human, ai, human2], parent_id="ckpt-goal", channel_versions={"messages": 6}), + _snapshot("ckpt-goal", [human, ai], parent_id="ckpt-turn1-tail", metadata=dict(_POSTGRES_INHERITED_STAMPS), channel_versions={"messages": 5, "goal": 3}), + _snapshot("ckpt-turn1-tail", [human, ai], channel_versions={"messages": 5}), + ] + accessor = _Accessor(history) + + base = asyncio.run(find_checkpoint_before_message(accessor, history[0], "h2", max_depth=10)) + + assert base.config["configurable"]["checkpoint_id"] == "ckpt-goal" + + +def test_lineage_walk_keeps_stamp_candidate_addressable_when_grandparent_is_gone(): + """A stamp candidate whose own parent link was pruned stays addressable. + + Retention may already have removed the grandparent; the walk must then + treat the candidate as addressable rather than skipping it on inherited + stamps — the returned replay base keeps the candidate's channel values. + """ + human = HumanMessage(id="h1", content="question") + ai = AIMessage(id="ai-1", content="answer") + human2 = HumanMessage(id="h2", content="follow-up") + history = [ + _snapshot("ckpt-turn2-head", [human, ai, human2], parent_id="ckpt-orphan-duration"), + _snapshot("ckpt-orphan-duration", [human, ai], parent_id="ckpt-pruned", metadata=dict(_POSTGRES_INHERITED_STAMPS), channel_versions={"messages": 5}), + ] + accessor = _Accessor(history) + + base = asyncio.run(find_checkpoint_before_message(accessor, history[0], "h2", max_depth=10)) + + assert base.config["configurable"]["checkpoint_id"] == "ckpt-orphan-duration" + + +def test_title_leaf_inheriting_duration_stamps_is_not_duration_only(): + """``_ensure_interrupted_title`` copies the head's metadata, stamps included. + + On memory/SQLite the leaf keeps its own ``writes.runtime_interrupt_title`` + marker, which must win over the inherited ``run_durations`` — otherwise a + title write stacked on a stamped head would be classified as a duration + checkpoint even where the marker is authoritative. + """ + snapshot = _snapshot( + "ckpt-title", + [], + metadata={ + "source": "update", + "step": 4, + "writes": {"runtime_interrupt_title": {"title": "hi"}}, + "run_durations": {"run-1": 12}, + "run_message_ids": {"ai-1": "run-1"}, + }, + ) + + assert is_duration_only_checkpoint(snapshot) is False + + +def test_stamp_fallback_disabled_requires_the_writes_marker(): + """Destructive callers must see a Postgres-shaped leaf as *not* duration-only. + + Without the ``writes`` marker the stamps are only evidence, not proof — + ``checkpoint_retention`` refuses them here and widens the class solely via + the parent-shape check in ``_mark_duration_leaves_without_the_marker``. + """ + parent = _snapshot("ckpt-real", [], channel_versions={"messages": 5}) + postgres_shaped = _snapshot("ckpt-duration", [], metadata=dict(_POSTGRES_DURATION_METADATA), channel_versions={"messages": 5}) + marker_backed = _snapshot( + "ckpt-duration", + [], + metadata={"source": "update", "step": 3, "writes": {"runtime_run_duration": {"run_ids": [], "message_ids": []}}}, + ) + + assert is_duration_only_checkpoint(postgres_shaped, parent=parent, stamp_fallback=False) is False + assert is_duration_only_checkpoint(marker_backed, parent=parent, stamp_fallback=False) is True + + +async def _post_strip_writes(saver, tup) -> None: + """Rewrite *tup*'s persisted metadata without ``writes``, in place.""" + + checkpoint = dict(getattr(tup, "checkpoint", {}) or {}) + metadata = dict(getattr(tup, "metadata", {}) or {}) + metadata.pop("writes", None) + parent_config = getattr(tup, "parent_config", None) + write_config = parent_config if isinstance(parent_config, dict) else {"configurable": {"thread_id": tup.config["configurable"]["thread_id"], "checkpoint_ns": "", "checkpoint_id": None}} + await saver.aput(write_config, checkpoint, metadata, {}) + + +async def _seed_lineage(mode: str): + from langgraph.checkpoint.memory import InMemorySaver + from langgraph.types import Overwrite + + from deerflow.runtime.checkpoint_state import CheckpointStateAccessor, build_state_mutation_graph + + saver = InMemorySaver() + accessor = CheckpointStateAccessor.bind(build_state_mutation_graph("seed", mode), saver, mode=mode) + config = {"configurable": {"thread_id": THREAD_ID, "checkpoint_ns": ""}} + await accessor.aupdate(config, {"messages": Overwrite([HumanMessage(content="q", id="h1"), AIMessage(content="a", id="ai-1")])}, as_node="seed") + return saver, accessor, config + + +async def _stamp_duration(saver, accessor, config): + from deerflow.runtime.runs.worker import persist_run_history_metadata + + await persist_run_history_metadata(checkpointer=saver, thread_id=THREAD_ID, durations={"run-1": 12}, message_run_ids={"ai-1": "run-1"}) + + +def test_classifier_uses_accessor_snapshots_in_full_mode() -> None: + """Real ``StateSnapshot`` objects must classify Postgres duration leaves (full).""" + + async def scenario() -> None: + saver, accessor, config = await _seed_lineage("full") + await _stamp_duration(saver, accessor, config) + head = await saver.aget_tuple(config) + await _post_strip_writes(saver, head) + + history = await accessor.ahistory(config) + leaf = history[0] + parent = history[1] if len(history) > 1 else None + assert is_duration_only_checkpoint(leaf, parent=parent, versions=await resolve_channel_versions(accessor, leaf), parent_versions=await resolve_channel_versions(accessor, parent) if parent is not None else None) is True + + asyncio.run(scenario()) + + +def test_classifier_uses_accessor_snapshots_in_delta_mode() -> None: + """Same integration gap exists on the delta read path; both must be covered.""" + + async def scenario() -> None: + saver, accessor, config = await _seed_lineage("delta") + await _stamp_duration(saver, accessor, config) + head = await saver.aget_tuple(config) + await _post_strip_writes(saver, head) + + history = await accessor.ahistory(config) + leaf = history[0] + parent = history[1] if len(history) > 1 else None + assert is_duration_only_checkpoint(leaf, parent=parent, versions=await resolve_channel_versions(accessor, leaf), parent_versions=await resolve_channel_versions(accessor, parent) if parent is not None else None) is True + + asyncio.run(scenario()) + + +def test_goal_leaf_stays_addressable_with_accessor_snapshots() -> None: + """A goal write stacked on a stamped head bumps the goal version and must stay.""" + + async def scenario() -> None: + saver, accessor, config = await _seed_lineage("full") + await _stamp_duration(saver, accessor, config) + + from deerflow.runtime.goal import write_thread_goal + + await write_thread_goal(saver, THREAD_ID, {"objective": "ship", "continuation_count": 0, "no_progress_count": 0, "updated_at": "2026-10-08T00:00:00Z"}) + goal_tup = await saver.aget_tuple(config) + await _post_strip_writes(saver, goal_tup) + + history = await accessor.ahistory(config) + leaf = history[0] + parent = history[1] if len(history) > 1 else None + assert is_duration_only_checkpoint(leaf, parent=parent, versions=await resolve_channel_versions(accessor, leaf), parent_versions=await resolve_channel_versions(accessor, parent) if parent is not None else None) is False + + asyncio.run(scenario()) + + +def test_degraded_raw_read_accessor_classifies_via_carried_versions() -> None: + """The degraded full-mode accessor ships channel versions on its snapshots.""" + + async def scenario() -> None: + saver, accessor, config = await _seed_lineage("full") + await _stamp_duration(saver, accessor, config) + head = await saver.aget_tuple(config) + await _post_strip_writes(saver, head) + + from app.gateway.services import _RawCheckpointReadAccessor + + raw = _RawCheckpointReadAccessor(saver, "full") + leaf_snap = await raw.aget(config) + parent_config = getattr(leaf_snap, "parent_config", None) + parent_snap = await raw.aget(parent_config) if isinstance(parent_config, dict) else None + assert leaf_snap.channel_versions is not None + assert is_duration_only_checkpoint(leaf_snap, parent=parent_snap, versions=_snapshot_versions_of(leaf_snap), parent_versions=_snapshot_versions_of(parent_snap)) is True + + asyncio.run(scenario()) + + +def _snapshot_versions_of(obj): + checkpoint = getattr(obj, "checkpoint", None) + if isinstance(checkpoint, dict): + versions = checkpoint.get("channel_versions") + if isinstance(versions, dict): + return versions + return getattr(obj, "channel_versions", None) + + +class _CountingAccessor(_Accessor): + """``_Accessor`` that records every read, pinning the walk's read budget.""" + + def __init__(self, snapshots: list[object]) -> None: + super().__init__(snapshots) + self.reads: list[str] = [] + + async def aget(self, config): + self.reads.append(config.get("configurable", {}).get("checkpoint_id")) + return await super().aget(config) + + +class _RawVersionAccessor: + """Accessor whose snapshots hide the payload, like langgraph's materialized ones. + + ``aget_tuple`` is the only way to reach ``channel_versions`` from such a + snapshot, so every call models one serial Postgres roundtrip. + """ + + def __init__(self, tuples: dict[str, object]) -> None: + self.tuples = tuples + self.raw_reads: list[str] = [] + + async def aget_tuple(self, config): + checkpoint_id = config.get("configurable", {}).get("checkpoint_id") + self.raw_reads.append(checkpoint_id) + return self.tuples.get(checkpoint_id) + + +def _hidden_snapshot(checkpoint_id: str, metadata: dict) -> SimpleNamespace: + """A snapshot exposing neither ``checkpoint`` nor ``channel_versions``.""" + + return SimpleNamespace( + config={"configurable": {"thread_id": THREAD_ID, "checkpoint_ns": "", "checkpoint_id": checkpoint_id}}, + metadata=dict(metadata), + ) + + +def test_lineage_walk_reads_a_refused_stamp_candidate_grandparent_once(): + """A refused stamp candidate must not re-read its own parent. + + The goal leaf inherits the duration stamps but bumped a channel version, so + the walk reads its parent for the shape check and then continues from the + candidate — whose parent link points at that very tuple. Both reads must hit + one entry in the walk's read cache. + """ + human = HumanMessage(id="h1", content="question") + ai = AIMessage(id="ai-1", content="answer") + human2 = HumanMessage(id="h2", content="follow-up") + history = [ + _snapshot("ckpt-head", [human, ai, human2], parent_id="ckpt-goal", channel_versions={"messages": 6}), + _snapshot("ckpt-goal", [human, ai, human2], parent_id="ckpt-turn1-tail", metadata=dict(_POSTGRES_INHERITED_STAMPS), channel_versions={"messages": 5, "goal": 3}), + _snapshot("ckpt-turn1-tail", [human, ai], channel_versions={"messages": 5}), + ] + accessor = _CountingAccessor(history) + + base = asyncio.run(find_checkpoint_before_message(accessor, history[0], "h2", max_depth=10)) + + assert base.config["configurable"]["checkpoint_id"] == "ckpt-turn1-tail" + assert accessor.reads.count("ckpt-turn1-tail") == 1 + + +def test_stamp_candidate_versions_are_lazy_and_memoized(): + """Only stamp candidates pay for version maps, and each tuple pays once. + + Behind a materialized snapshot every lookup is a serial raw read, so a + window the marker tier or the stamp gate already decided must issue none, + and resolving a candidate must reuse the pass that resolved its parent. + """ + candidate = _hidden_snapshot("ckpt-candidate", _POSTGRES_DURATION_METADATA) + parent = _hidden_snapshot("ckpt-parent", _POSTGRES_DURATION_METADATA) + marker_decided = _hidden_snapshot("ckpt-marker", {"source": "update", "step": 3, "writes": {"runtime_run_duration": {"run_id": "run-1"}}}) + accessor = _RawVersionAccessor( + { + "ckpt-candidate": SimpleNamespace(checkpoint={"channel_versions": {"messages": 5}}), + "ckpt-parent": SimpleNamespace(checkpoint={"channel_versions": {"messages": 5}}), + } + ) + cache: dict[tuple[str, str, str], object] = {} + + assert asyncio.run(resolve_stamp_candidate_versions(accessor, marker_decided, parent, cache)) == (None, None) + assert accessor.raw_reads == [] + + assert asyncio.run(resolve_stamp_candidate_versions(accessor, candidate, parent, cache)) == ({"messages": 5}, {"messages": 5}) + assert accessor.raw_reads == ["ckpt-candidate", "ckpt-parent"] + + assert asyncio.run(resolve_stamp_candidate_versions(accessor, parent, candidate, cache)) == ({"messages": 5}, {"messages": 5}) + assert accessor.raw_reads == ["ckpt-candidate", "ckpt-parent"] + + assert asyncio.run(resolve_stamp_candidate_versions(accessor, candidate, None, cache)) == (None, None) + assert accessor.raw_reads == ["ckpt-candidate", "ckpt-parent"]