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"]