diff --git a/backend/AGENTS.md b/backend/AGENTS.md index 2c166b4d6..641f7a0ba 100644 --- a/backend/AGENTS.md +++ b/backend/AGENTS.md @@ -505,7 +505,7 @@ JSONL event stores when `GATEWAY_WORKERS > 1`. - Gateway checkpoint mutations outside run execution must use `services.reserve_checkpoint_write()`, which composes the process-local thread lock with the durable `checkpoint_write` reservation. Manual compaction, `POST /threads/{id}/state`, and both goal mutation routes (`PUT` / `DELETE /threads/{id}/goal`, including creation of a missing goal checkpoint) use this boundary, so an existing run blocks the write and the reservation blocks new reject/interrupt/rollback runs across workers. - `POST /wait` (both thread-scoped and `/api/runs/wait`) drains the stream bridge via `wait_for_run_completion()` instead of bare `await record.task`, so it honours the run's `on_disconnect` setting and cancels the background run on real client disconnect rather than returning a stale checkpoint (issue #3265). - Redis `StreamBridge` keys use a rolling retained-buffer TTL (`stream_bridge.stream_ttl_seconds`, refreshed on `publish()` / `publish_end()`) as a leak safety net, not as a run timeout. Startup and lease-driven periodic orphan recovery share one Gateway stream-terminalization path: after `RunManager` durably marks a run `error` with `stop_reason=orphan_recovered`, Gateway publishes `END_SENTINEL` and schedules stream cleanup. The periodic store scan, per-row status writes, and Gateway callback run as one supervised single-flight task, so a slow pass is skipped at the next interval instead of piling up or pausing the sole lease-renewal loop. Store retries have bounded attempts/backoff; an individual operation still relies on the database driver/pool timeout. `RunManager.shutdown()` gives active user runs priority within its shared deadline, then drains or cancels orphan recovery. Gateway tracks delayed recovered-stream cleanups and converts unfinished delays to immediate deletes before closing the bridge; the Redis TTL remains the outage safety net. Only startup recovery, before the runtime yields to requests, projects the latest affected thread to `error`; periodic recovery deliberately avoids that non-atomic projection because `ThreadMetaStore` has no `latest_run_id` conditional-update contract. Store-only SSE and `/wait` consumers wait for the bridge's real END marker after an ordinary durable terminal status, because status persistence can precede tail events. The explicit `orphan_recovered` signal is the only heartbeat fallback: its publisher is known to be gone, so it supplies the liveness boundary if END publication fails or the retained key expires. Malformed `Last-Event-ID` reconnect values live-tail new Redis events rather than replaying the retained buffer. Keep cross-component recovery orchestration in Gateway through the generic `RunManager.on_orphans_recovered` callback; do not introduce a harness-to-app dependency. Callback failure warnings include every recovered `run_id` so operators can identify rows whose Gateway-side terminalization needs inspection. -- Thread-scoped run creation accepts `checkpoint` / `checkpoint_id`; Gateway validates the checkpoint belongs to the request thread before writing `checkpoint_id` / `checkpoint_ns` into `config.configurable` for LangGraph branching. +- Thread-scoped run creation accepts `checkpoint` / `checkpoint_id`; Gateway validates the checkpoint belongs to the request thread before writing `checkpoint_id` / `checkpoint_ns` into `config.configurable` for LangGraph branching. In `delta` checkpoint mode the worker rewrites that fork into a linear head write before the graph starts (see "A delta-mode run cannot fork" under Checkpoint Channel Modes), because delta state for a fork replays the abandoned sibling's writes. - Thread-scoped Gateway runs evaluate an active `ThreadState.goal` after the visible turn completes. `runtime/goal.py` asks a non-thinking evaluator model to judge only visible conversation evidence and return a typed blocker; the evaluator model is created once per run and reused across hidden continuation checks. The evaluator runs after the graph root's tracing scope has already closed, so `create_goal_evaluator_model`/`evaluate_goal_completion` attach their own model-level tracing callbacks (`attach_tracing=True`) and inject Langfuse trace metadata (`thread_id`/`user_id`/`deerflow_trace_id`) directly onto the `ainvoke` call — the same standalone-caller pattern as `oneshot_llm.run_oneshot_llm` and `MemoryUpdater` (see Tracing System below). Satisfied goals are cleared; every non-satisfied evaluation — continuable or stand-down — is persisted with `last_evaluation` (the blocker, reason, and evidence summary; outcomes that stop the loop additionally record a `stand_down_reason` for observability), but only `goal_not_met_yet` evaluations are streamed as hidden `HumanMessage` continuations, and only when a durable assistant end-of-turn checkpoint exists, the run has not been aborted, the thread did not change during evaluation, and the no-progress breaker has not fired. The continuation cap is 8 — a hard maximum in the `0`–`8` range; callers requesting more are clamped (`set_goal`/TUI) or rejected with 422 (`PUT /goal`). The no-progress breaker keys on the latest visible assistant evidence (not the evaluator's free-text reason, which an LLM rewords every turn), so two consecutive continuations that add no new visible assistant output stop the loop after 2 attempts. Model-response cleanup helpers such as think-block stripping and code-fence stripping live in `deerflow.utils.llm_text` so `runtime/goal.py` and Gateway suggestion parsing share the same JSON-prep behavior. - Run event stream changes must keep producer code, `deerflow/constants.py`, `runtime/events/catalog.py`, `contracts/run_event_stream_contract.json`, `backend/docs/RUN_EVENT_STREAM.md`, and `tests/test_run_event_stream_contract.py` in sync. The dependency-free constants module owns the persisted envelope limits (`event_type` 32 characters, `category` 16) and cross-layer workspace event identity; the catalog owns validated runtime definitions and categories. Dynamic middleware tags are limited to 21 characters after the `middleware:` prefix. The JSON contract owns payload schemas, backend-specific storage semantics, legacy aliases, and compatibility rules; conformance tests require both views and all producer groups to agree. `run.end.content` remains opaque and may retain nested Python values in memory while JSONL/database stores stringify non-JSON nested values, so consumers must not assume backend-identical nested output representations. @@ -908,9 +908,11 @@ Checkpointer storage runs in one of two channel modes, selected by `checkpoint_c **Replay checkpoint lookup prefers lineage and degrades only for an explicitly missing legacy parent link.** Branch and regenerate paths first walk `parent_config`, which prevents a global chronological scan from selecting a sibling created by regeneration. `CheckpointParentMissingError` alone enables the bounded newest-first history fallback in `app/gateway/checkpoint_lineage.py`; cycles, dangling/non-addressable parents, target mismatches, and depth exhaustion raise `CheckpointLineageIntegrityError` and fail closed instead of selecting a sibling. The compatibility scans request 400 raw checkpoints so up to 200 duration-only entries do not consume the effective branch-history budget; the fallback scans oldest-to-newest internally, skips duration-only checkpoints, and accepts only checkpoints with an addressable id as the replay base. A source history with no discoverable pre-user checkpoint preserves the historical single-checkpoint branch behavior instead of rejecting the branch; regeneration remains unavailable for that inherited response. Existing single-checkpoint branches are not mutated by regenerate preparation, and no raw checkpoint tuple is copied across threads because delta state depends on ancestry and pending writes. Regenerate source-run lookup uses the current thread's exact event, then the server-stamped `run_id` on the copied human message, then verified RunManager content matching; it does not read parent-thread events. Storage or checkpoint-mode failures are not treated as a missing base and still fail closed. -**Wholesale message replacement uses a state-only mutation graph + `Overwrite`.** `update_state` values pass through channel reducers (`add_messages` merge in full, append in delta), so replacing `messages` wholesale (run rollback, context compaction) requires `{"messages": Overwrite([...])}`. The write goes through `build_state_mutation_graph(as_node, mode, state_schema)` — when the write carries materialized state, `state_schema` MUST be the thread's effective schema (`graph_state_schema(assistant_graph)`), because the base-ThreadState fallback silently discards written channels contributed by custom `AgentMiddleware.state_schema`. Channels absent from the write are unaffected: forked checkpoints clone the parent's channel blobs, so middleware channels survive rollback/compaction regardless of schema (locked by `test_rollback_preserves_middleware_contributed_channels` and `test_compact_thread_context_preserves_middleware_contributed_channels`) — a compiled graph with one no-op node (entry = finish) whose checkpoint machinery (channels/versions/metadata) is identical to the agent graph's but schedules no pending tasks, so the restored/compacted head stays idle instead of re-triggering the agent. Never hand-write checkpoints via `checkpointer.aput` for this; raw writers elsewhere must preserve checkpoint parentage — severed ancestry breaks delta replay (see `runtime/runs/worker.py` writer parenting and `checkpoint_patches.py`). +**A delta-mode run cannot fork; `runtime/runs/worker.py` linearizes the resume instead.** Resuming from an older checkpoint (regenerate, or any client-supplied `checkpoint`) forks the lineage, and delta state for a fork is not materializable: `BaseCheckpointSaver.get_delta_channel_history` — and the bespoke overrides in `InMemorySaver`/`PostgresSaver` — collect **every** `pending_writes` entry stored on each on-path ancestor, but a shared parent also carries the writes of the sibling child that was abandoned. Those writes replay into the fork, so the run starts from a message list still containing the answer it was supposed to replace (#4458: regenerating in a branched thread showed the superseded assistant message beside the new one after a reload; reproduced on postgres, sqlite, and the in-memory saver). Write-to-child ownership belongs to the upstream delta contract, so DeerFlow does not reimplement the walk: `_linearize_delta_checkpoint_resume` materializes the requested checkpoint's complete state and writes every channel onto the **current head** (which has no siblings) through the state mutation graph, using `Overwrite` for reducer channels and resetting newer head-only channels to their schema default (or `None` when no constructible default exists); it then drops the `checkpoint_id` selector and lets the run proceed linearly, while the abandoned turn stays in history as the rewritten head's ancestry. The worker holds `_checkpoint_thread_lock` across `_capture_rollback_point` and the optional linear rewrite, making the rollback snapshot and rewrite atomic with graph streaming and the preceding run's duration-metadata checkpoint write. Capture preserves the complete real pre-run state; cancel-with-rollback then linearly replaces the current delta head with that captured state rather than forking the now-shared pre-run checkpoint, so the abandoned turn is restored without replaying the resume sibling's writes. The worker also recomputes the current-run message boundary from the rewritten state and fails closed (an unreadable resume checkpoint raises rather than falling back to the corrupt fork). `full` mode keeps forking — its checkpoints carry complete `channel_values` and need no replay — so LangGraph branching semantics are unchanged there. Root namespace only; subgraph namespaces are left alone. -**Run rollback flow** (`runtime/runs/worker.py`): `_capture_rollback_point` materializes pre-run state (messages via accessor, raw `pending_writes` via `aget_tuple`) into an immutable `RollbackPoint` before the run starts — capture failure disables rollback (fail-closed), never restores partial state. Cancel-with-rollback then forks from the pre-run checkpoint via the mutation graph and replays the captured pending writes. +**Wholesale state replacement uses a state-only mutation graph + `Overwrite`.** `update_state` values pass through channel reducers (`add_messages` merge in full, append in delta), so replacing reducer values requires `Overwrite` rather than an ordinary update. Full-mode rollback and context compaction replace `messages`; delta resume and delta rollback replace every materialized channel and reset current-head-only channels to their schema default (or `None`). These writes go through `build_state_mutation_graph(as_node, mode, state_schema)`, and `state_schema` MUST be the thread's effective schema (`graph_state_schema(assistant_graph)`), because the base-ThreadState fallback silently discards written channels contributed by custom `AgentMiddleware.state_schema`. Channels absent from a full-mode fork write inherit the parent's channel blobs, so middleware channels survive rollback/compaction (locked by `test_rollback_preserves_middleware_contributed_channels` and `test_compact_thread_context_preserves_middleware_contributed_channels`). The compiled mutation graph has one no-op node (entry = finish) whose checkpoint machinery (channels/versions/metadata) is identical to the agent graph's but schedules no pending tasks, so the restored/compacted head stays idle instead of re-triggering the agent. Never hand-write checkpoints via `checkpointer.aput` for this; raw writers elsewhere must preserve checkpoint parentage — severed ancestry breaks delta replay (see `runtime/runs/worker.py` writer parenting and `checkpoint_patches.py`). + +**Run rollback flow** (`runtime/runs/worker.py`): `_capture_rollback_point` materializes the complete pre-run state via the accessor and captures raw `pending_writes` via `aget_tuple` into an immutable `RollbackPoint` before the run starts — capture failure disables rollback (fail-closed), never restores partial state. In `full` mode, cancel-with-rollback forks from the pre-run checkpoint via the mutation graph and inherits non-message channels from that parent. In `delta` mode, forking is unsafe once the cancelled path has attached sibling writes to the pre-run checkpoint, so rollback replaces every captured channel on the current head, using `Overwrite` for reducers and schema defaults for current-head-only channels. Both modes reattach only the captured pre-run pending writes to the restored checkpoint. **Where things live**: - `runtime/checkpoint_mode.py` — mode freeze, marker injection, delta detection, compatibility gate, both error types diff --git a/backend/packages/harness/deerflow/runtime/runs/worker.py b/backend/packages/harness/deerflow/runtime/runs/worker.py index cd8cfd17d..7af4c76f8 100644 --- a/backend/packages/harness/deerflow/runtime/runs/worker.py +++ b/backend/packages/harness/deerflow/runtime/runs/worker.py @@ -40,7 +40,13 @@ from deerflow.runtime.checkpoint_mode import ( aensure_checkpoint_mode_compatible, inject_checkpoint_mode, ) -from deerflow.runtime.checkpoint_state import CheckpointStateAccessor, build_state_mutation_graph, graph_state_schema +from deerflow.runtime.checkpoint_state import ( + CheckpointStateAccessor, + build_state_mutation_graph, + graph_reducer_channels, + graph_state_schema, + graph_writable_channels, +) from deerflow.runtime.context_keys import CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY from deerflow.runtime.goal import ( DEFAULT_MAX_GOAL_CONTINUATIONS, @@ -647,14 +653,38 @@ async def run_agent( # failure disables rollback: restoring an empty or partial message # history would silently truncate the thread. if checkpointer is not None: - try: - rollback_point = await _capture_rollback_point(accessor, checkpointer, checkpoint_config) - except Exception: - snapshot_capture_failed = True - logger.warning("Could not capture pre-run checkpoint snapshot for run %s", run_id, exc_info=True) - if rollback_point is not None: - pre_run_checkpoint_id = rollback_point.config.get("configurable", {}).get("checkpoint_id") - pre_existing_message_ids = _collect_pre_existing_message_ids({"messages": list(rollback_point.messages)}) + # A previous successful run may still be persisting duration + # metadata after its active admission slot is released. Share its + # checkpoint lock so the rollback snapshot and any resume rewrite + # are one uninterrupted read/write sequence against the head. + async with _checkpoint_thread_lock(thread_id): + try: + rollback_point = await _capture_rollback_point(accessor, checkpointer, checkpoint_config) + except Exception: + snapshot_capture_failed = True + logger.warning("Could not capture pre-run checkpoint snapshot for run %s", run_id, exc_info=True) + if rollback_point is not None: + pre_run_checkpoint_id = rollback_point.config.get("configurable", {}).get("checkpoint_id") + pre_existing_message_ids = _collect_pre_existing_message_ids({"messages": list(rollback_point.messages)}) + + # Resuming from an older checkpoint is a fork, and a delta fork + # materializes the abandoned sibling's writes back into state + # (#4458). Rewrite it as a linear head write *after* the rollback + # point is captured, so cancel-with-rollback still restores the + # real pre-run head rather than the rolled-back one. + resumed_messages = await _linearize_delta_checkpoint_resume( + accessor=accessor, + checkpointer=checkpointer, + config=config, + thread_id=thread_id, + run_id=run_id, + ) + if resumed_messages is not None: + # The graph now starts from the selected state, so the + # current-run message boundary is that state, not the head we + # captured for rollback. + pre_existing_message_ids = _collect_pre_existing_message_ids({"messages": list(resumed_messages)}) + initial_runnable_config = RunnableConfig(**config) runtime_ctx[CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY] = frozenset(pre_existing_message_ids) _install_runtime_context(config, runtime_ctx) @@ -1354,15 +1384,16 @@ async def _prepare_goal_continuation_input( @dataclass(frozen=True) class RollbackPoint: - """Materialized pre-run state used to fork the pre-run checkpoint lineage. + """Materialized pre-run state used to restore the thread after cancellation. Raw checkpoint blobs cannot reconstruct Delta-channel messages (their - checkpoints omit ``channel_values``), so rollback restores messages by - applying an ``Overwrite`` through a state-mutation graph anchored at the - pre-run checkpoint instead of cloning the raw blob. + checkpoints omit the materialized value), so rollback preserves those + messages plus delta mode's materialized non-message state in addition to + the raw pending writes. """ config: dict[str, Any] + state_values: dict[str, Any] messages: tuple[Any, ...] metadata: dict[str, Any] pending_writes: tuple[tuple[str, str, Any], ...] @@ -1384,8 +1415,9 @@ async def _capture_rollback_point( if not configurable.get("checkpoint_id"): return None checkpoint_tuple = await _call_checkpointer_method(checkpointer, "aget_tuple", "get_tuple", snapshot_config) - values = getattr(snapshot, "values", None) or {} - messages = values.get("messages") if isinstance(values, dict) else None + raw_values = getattr(snapshot, "values", None) or {} + messages = raw_values.get("messages") if isinstance(raw_values, dict) else None + state_values = copy.deepcopy({key: value for key, value in raw_values.items() if key != "messages"}) if accessor.mode == "delta" and isinstance(raw_values, dict) else {} return RollbackPoint( config={ "configurable": { @@ -1394,12 +1426,132 @@ async def _capture_rollback_point( "checkpoint_id": configurable.get("checkpoint_id"), } }, + state_values=state_values, messages=tuple(messages or ()), metadata=dict(getattr(snapshot, "metadata", None) or {}), pending_writes=tuple(getattr(checkpoint_tuple, "pending_writes", ()) or ()), ) +def _complete_state_replacement_values( + *, + mutation_graph: Any, + selected_values: dict[str, Any], + current_values: dict[str, Any], + run_id: str, + operation: str, +) -> dict[str, Any]: + """Build a whole-state replacement through the graph's effective schema.""" + writable_fields = graph_writable_channels(mutation_graph) + reducer_fields = graph_reducer_channels(mutation_graph) + if writable_fields is None or reducer_fields is None: + raise RuntimeError(f"Run {run_id} could not inspect the state schema for {operation}") + + replacement_values: dict[str, Any] = {} + for field_name in writable_fields: + if field_name in selected_values: + replacement = copy.deepcopy(selected_values[field_name]) + elif field_name in current_values: + # LangGraph has no public "unset channel" update. A fresh channel + # exposes its schema default when one exists (for example [] / {}); + # optional and otherwise-unconstructible channels reset to None. + channel = mutation_graph.channels.get(field_name) + replacement = copy.deepcopy(channel.get()) if channel is not None and channel.is_available() else None + else: + continue + replacement_values[field_name] = Overwrite(replacement) if field_name in reducer_fields else replacement + return replacement_values + + +async def _linearize_delta_checkpoint_resume( + *, + accessor: CheckpointStateAccessor, + checkpointer: Any, + config: dict[str, Any], + thread_id: str, + run_id: str, +) -> list[Any] | None: + """Replace a delta-mode checkpoint fork with an equivalent linear write. + + Resuming from an older checkpoint forks the lineage, and in ``delta`` mode + the fork's state cannot be materialized correctly: the delta history walk + collects **every** ``pending_writes`` entry stored on each on-path + ancestor, but a shared parent also carries the writes of the sibling child + that was abandoned. Those writes are replayed into the fork, so the run + starts from a message list that still contains the answer it was supposed + to replace — regenerating in a branched thread surfaced this as the old + assistant message reappearing beside the new one after a reload (#4458). + Reproduced on postgres, sqlite, and the in-memory saver; ``full`` mode is + unaffected because its checkpoints carry complete ``channel_values`` and + need no replay. + + The upstream contract (`BaseCheckpointSaver.get_delta_channel_history` and + the savers overriding it) is where write-to-child ownership belongs, so + this does not reimplement it. Instead the fork is expressed as what it + means: materialize the requested checkpoint's state and write it with + replace semantics on the **current head**, which has no other children, + then run linearly. Every materialized channel is restored; channels that + exist only on the newer head are reset to their schema default (or + ``None`` when the channel has no constructible default). The abandoned + turn stays in checkpoint history as the rewritten head's ancestry. + + Returns the materialized messages when the resume was linearized, or + ``None`` when there was nothing to do (full mode, no checkpoint selector, + a non-root namespace, or a selector that already names the head). Failures + propagate: silently falling back to the fork would persist the corrupted + history this exists to prevent. The worker call site holds + ``_checkpoint_thread_lock`` across rollback capture and this rewrite; do + not reacquire that non-reentrant lock inside this helper. + """ + if checkpointer is None or accessor.mode != "delta": + return None + configurable = config.get("configurable") + if not isinstance(configurable, dict): + return None + checkpoint_id = configurable.get("checkpoint_id") + if not isinstance(checkpoint_id, str) or not checkpoint_id: + return None + if configurable.get("checkpoint_ns"): + # Subgraph namespaces have their own lineage; the Gateway only selects + # root checkpoints, so leave anything else untouched. + return None + + head_config: dict[str, Any] = {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}} + head = await accessor.aget(head_config) + if _checkpoint_id(head) == checkpoint_id: + # Selecting the head is already linear — no sibling can exist yet. + return None + + source_config: dict[str, Any] = {"configurable": {"thread_id": thread_id, "checkpoint_ns": "", "checkpoint_id": checkpoint_id}} + snapshot = await accessor.aget(source_config) + values = getattr(snapshot, "values", None) or {} + messages = values.get("messages") if isinstance(values, dict) else None + if not isinstance(messages, list): + raise RuntimeError(f"Run {run_id} could not materialize resume checkpoint {checkpoint_id}") + + # Write through the thread's effective schema so every application and + # middleware channel can be restored. Reducer channels need Overwrite to + # replace their already-aggregated value instead of merging it again. + mutation_graph = build_state_mutation_graph("checkpoint_resume", accessor.mode, graph_state_schema(getattr(accessor, "graph", None))) + selected_values = dict(values) + head_values = getattr(head, "values", None) or {} + head_values = dict(head_values) if isinstance(head_values, dict) else {} + replacement_values = _complete_state_replacement_values( + mutation_graph=mutation_graph, + selected_values=selected_values, + current_values=head_values, + run_id=run_id, + operation="checkpoint resume", + ) + + mutation_accessor = CheckpointStateAccessor.bind(mutation_graph, checkpointer, mode=accessor.mode) + await mutation_accessor.aupdate(head_config, replacement_values, as_node="checkpoint_resume") + configurable.pop("checkpoint_id", None) + configurable.pop("checkpoint_map", None) + logger.info("Run %s linearized a delta-mode resume of checkpoint %s onto thread %s", run_id, checkpoint_id, thread_id) + return list(messages) + + async def _rollback_to_pre_run_checkpoint( *, accessor: CheckpointStateAccessor | None, @@ -1409,13 +1561,14 @@ async def _rollback_to_pre_run_checkpoint( rollback_point: RollbackPoint | None, snapshot_capture_failed: bool, ) -> None: - """Fork the pre-run checkpoint lineage with the pre-run messages restored. + """Restore the complete pre-run state after a cancelled run. - The fork is written through a state-only mutation graph (the synthetic - ``rollback_restore`` node must be registered for ``as_node`` and finishes - immediately so no agent nodes are scheduled). LangGraph owns the restored - checkpoint's source/step/channel versions/parent/timestamp; the parent - pointer back to the pre-run checkpoint is the audit trail. + Full mode forks the captured pre-run checkpoint and overwrites messages; + all other channels inherit from that parent. Delta mode cannot safely fork + once the cancelled path has attached writes to the same parent, so it + replaces every captured channel on the current head instead. Both writes + use a state-only mutation graph whose synthetic ``rollback_restore`` node + finishes immediately and schedules no agent work. """ if checkpointer is None: logger.info("Run %s rollback requested but no checkpointer is configured", run_id) @@ -1441,15 +1594,35 @@ async def _rollback_to_pre_run_checkpoint( logger.warning("Run %s rollback skipped: agent accessor unavailable", run_id) return - # The restored checkpoint inherits every channel from the pre-run fork; - # compile the mutation graph with the thread's effective schema so - # middleware-contributed channels survive (the base ThreadState fallback - # would silently drop them). + # Compile with the thread's effective schema so middleware-contributed + # channels survive (the base ThreadState fallback would silently drop + # them). mutation_graph = build_state_mutation_graph("rollback_restore", accessor.mode, graph_state_schema(getattr(accessor, "graph", None))) mutation_accessor = CheckpointStateAccessor.bind(mutation_graph, checkpointer, mode=accessor.mode) + if accessor.mode == "delta": + # A delta rollback fork has the same write-ownership problem as a + # checkpoint resume: the captured parent now carries writes from the + # cancelled sibling. Restore linearly on the current head instead. + restore_config: dict[str, Any] = {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}} + current = await accessor.aget(restore_config) + raw_current_values = getattr(current, "values", None) or {} + current_values = dict(raw_current_values) if isinstance(raw_current_values, dict) else {} + selected_values = copy.deepcopy(rollback_point.state_values) + selected_values["messages"] = list(rollback_point.messages) + replacement_values = _complete_state_replacement_values( + mutation_graph=mutation_graph, + selected_values=selected_values, + current_values=current_values, + run_id=run_id, + operation="rollback", + ) + else: + restore_config = rollback_point.config + replacement_values = {"messages": Overwrite(list(rollback_point.messages))} + restored_config = await mutation_accessor.aupdate( - rollback_point.config, - {"messages": Overwrite(list(rollback_point.messages))}, + restore_config, + replacement_values, as_node="rollback_restore", ) if not isinstance(restored_config, dict): diff --git a/backend/tests/test_run_worker_delta_resume.py b/backend/tests/test_run_worker_delta_resume.py new file mode 100644 index 000000000..d64152bf4 --- /dev/null +++ b/backend/tests/test_run_worker_delta_resume.py @@ -0,0 +1,495 @@ +"""Delta-mode checkpoint resume linearization (#4458). + +Resuming a run from an older checkpoint forks the lineage. In ``delta`` mode +that fork cannot be materialized correctly — the delta history walk replays +every ``pending_writes`` entry stored on each on-path ancestor, including the +writes of the sibling child that was abandoned — so the run starts from a +message list that still contains the answer it was meant to replace. These +tests pin the worker's linearization: the requested state is written onto the +current head (which has no siblings) and the run proceeds linearly. +""" + +from __future__ import annotations + +import asyncio +from types import SimpleNamespace +from typing import Annotated, Any, TypedDict +from unittest.mock import AsyncMock + +import pytest +from langchain_core.messages import AIMessage, AnyMessage, HumanMessage +from langgraph.channels.delta import DeltaChannel +from langgraph.checkpoint.memory import InMemorySaver +from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver +from langgraph.graph import StateGraph +from langgraph.graph.message import add_messages +from langgraph.types import Overwrite + +from deerflow.agents.thread_state import merge_message_writes +from deerflow.runtime.checkpoint_state import CheckpointStateAccessor, build_state_mutation_graph +from deerflow.runtime.runs.manager import RunManager +from deerflow.runtime.runs.schemas import RunStatus +from deerflow.runtime.runs.worker import RunContext, _checkpoint_thread_lock, _linearize_delta_checkpoint_resume, run_agent + +pytestmark = pytest.mark.anyio + + +@pytest.fixture +def anyio_backend(): + return "asyncio" + + +class _DeltaChannelState(TypedDict): + messages: Annotated[list[AnyMessage], DeltaChannel(merge_message_writes, snapshot_frequency=1000)] + + +class _FullChannelState(TypedDict): + messages: Annotated[list[AnyMessage], add_messages] + + +def _replace_optional_state(existing: dict[str, str] | None, new: dict[str, str] | None) -> dict[str, str] | None: + return existing if new is None else new + + +class _ExtendedDeltaChannelState(_DeltaChannelState): + notes: Annotated[list[AnyMessage], add_messages] + marker: str | None + head_only: Annotated[dict[str, str] | None, _replace_optional_state] + head_last_only: str | None + + +def _build_answer_graph(state_schema: type, checkpointer: Any, answer_id: str): + async def _answer(state: dict[str, Any]) -> dict[str, Any]: + return {"messages": [AIMessage(content=f"answer for {answer_id}", id=answer_id)]} + + builder = StateGraph(state_schema) + builder.add_node("answer", _answer) + builder.set_entry_point("answer") + builder.set_finish_point("answer") + return builder.compile(checkpointer=checkpointer) + + +def _ids(snapshot: Any) -> list[str]: + return [message.id for message in (snapshot.values or {}).get("messages", [])] + + +def _run_config(thread_id: str, checkpoint_id: str | None = None) -> dict[str, Any]: + configurable: dict[str, Any] = {"thread_id": thread_id, "checkpoint_ns": ""} + if checkpoint_id is not None: + configurable["checkpoint_id"] = checkpoint_id + return {"configurable": configurable} + + +async def _seed_two_turns(checkpointer: Any, state_schema: type, thread_id: str): + """h1 -> a1, h2 -> a2, returning (accessor, head, pre-h2 snapshot).""" + config = _run_config(thread_id) + await _build_answer_graph(state_schema, checkpointer, "a1").ainvoke({"messages": [HumanMessage(content="q1", id="h1")]}, config) + graph = _build_answer_graph(state_schema, checkpointer, "a2") + await graph.ainvoke({"messages": [HumanMessage(content="q2", id="h2")]}, config) + + mode = "delta" if state_schema is _DeltaChannelState else "full" + accessor = CheckpointStateAccessor.bind(graph, checkpointer, mode=mode) + head = await accessor.aget(config) + history = await accessor.ahistory(config, limit=20) + pre_turn = next(snapshot for snapshot in history if "h2" not in _ids(snapshot)) + return accessor, head, pre_turn + + +async def test_linearizes_a_delta_resume_onto_the_head(): + checkpointer = InMemorySaver() + accessor, head, pre_turn = await _seed_two_turns(checkpointer, _DeltaChannelState, "thread-1") + base_id = pre_turn.config["configurable"]["checkpoint_id"] + config = _run_config("thread-1", base_id) + + resumed = await _linearize_delta_checkpoint_resume( + accessor=accessor, + checkpointer=checkpointer, + config=config, + thread_id="thread-1", + run_id="run-1", + ) + + assert [message.id for message in resumed] == ["h1", "a1"] + # The selector is consumed: the run continues on the (rewritten) head. + assert "checkpoint_id" not in config["configurable"] + new_head = await accessor.aget(_run_config("thread-1")) + assert _ids(new_head) == ["h1", "a1"] + assert new_head.config["configurable"]["checkpoint_id"] != head.config["configurable"]["checkpoint_id"] + + +async def test_linearization_restores_all_selected_state_and_clears_newer_channels(): + checkpointer = InMemorySaver() + config = _run_config("thread-state") + + first_graph = _build_answer_graph(_ExtendedDeltaChannelState, checkpointer, "a1") + await first_graph.ainvoke( + { + "messages": [HumanMessage(content="q1", id="h1")], + "notes": [HumanMessage(content="selected", id="note-selected")], + "marker": "selected", + }, + config, + ) + accessor = CheckpointStateAccessor.bind(first_graph, checkpointer, mode="delta") + selected = await accessor.aget(config) + + second_graph = _build_answer_graph(_ExtendedDeltaChannelState, checkpointer, "a2") + await second_graph.ainvoke( + { + "messages": [HumanMessage(content="q2", id="h2")], + "notes": [HumanMessage(content="newer", id="note-newer")], + "marker": "newer", + "head_only": {"value": "must-not-leak"}, + "head_last_only": "must-not-leak", + }, + config, + ) + accessor = CheckpointStateAccessor.bind(second_graph, checkpointer, mode="delta") + selected_id = selected.config["configurable"]["checkpoint_id"] + + await _linearize_delta_checkpoint_resume( + accessor=accessor, + checkpointer=checkpointer, + config=_run_config("thread-state", selected_id), + thread_id="thread-state", + run_id="run-state", + ) + + rewritten = await accessor.aget(config) + assert _ids(rewritten) == ["h1", "a1"] + assert [message.id for message in rewritten.values["notes"]] == ["note-selected"] + assert rewritten.values["marker"] == "selected" + assert rewritten.values.get("head_only") is None + assert rewritten.values.get("head_last_only") is None + + +async def test_regenerating_in_a_branched_thread_does_not_resurrect_the_old_answer(tmp_path): + """The #4458 shape end-to-end on a persistent saver. + + A branch writes two synthetic checkpoints (replay base + visible head); + regenerating the inherited answer resumes from the replay base. Without + linearization the delta walk replays the branch head's own ``Overwrite`` + — which is stored on that shared parent — and the superseded assistant + message comes back alongside the new one. + """ + db_path = tmp_path / "branch.sqlite3" + async with AsyncSqliteSaver.from_conn_string(str(db_path)) as checkpointer: + await checkpointer.setup() + _, source_head, source_pre_turn = await _seed_two_turns(checkpointer, _DeltaChannelState, "parent") + + mutation_graph = build_state_mutation_graph("branch", "delta", _DeltaChannelState) + branch_writer = CheckpointStateAccessor.bind(mutation_graph, checkpointer, mode="delta") + replay_base_config = await branch_writer.aupdate( + _run_config("branch"), + {"messages": Overwrite(list(source_pre_turn.values["messages"]))}, + as_node="branch", + ) + await branch_writer.aupdate( + replay_base_config, + {"messages": Overwrite(list(source_head.values["messages"]))}, + as_node="branch", + ) + + graph = _build_answer_graph(_DeltaChannelState, checkpointer, "a2-new") + accessor = CheckpointStateAccessor.bind(graph, checkpointer, mode="delta") + branch_history = await accessor.ahistory(_run_config("branch"), limit=20) + base = next(snapshot for snapshot in branch_history if "h2" not in _ids(snapshot)) + config = _run_config("branch", base.config["configurable"]["checkpoint_id"]) + + await _linearize_delta_checkpoint_resume( + accessor=accessor, + checkpointer=checkpointer, + config=config, + thread_id="branch", + run_id="run-regen", + ) + # The regenerate run replays the same human message and answers again. + await graph.ainvoke({"messages": [HumanMessage(content="q2", id="h2")]}, config) + + final = await accessor.aget(_run_config("branch")) + assert _ids(final) == ["h1", "a1", "h2", "a2-new"] + + +async def test_run_agent_streams_from_the_linearized_delta_resume(tmp_path): + """Pin selector removal through ``run_agent`` and its stream config. + + Helper-level tests would still pass if the worker later streamed from the + stale ``RunnableConfig`` built before linearization. This exercises the + production boundary that hands the rewritten config to ``agent.astream``. + """ + db_path = tmp_path / "worker-branch.sqlite3" + async with AsyncSqliteSaver.from_conn_string(str(db_path)) as checkpointer: + await checkpointer.setup() + _, source_head, source_pre_turn = await _seed_two_turns(checkpointer, _DeltaChannelState, "worker-parent") + + mutation_graph = build_state_mutation_graph("branch", "delta", _DeltaChannelState) + branch_writer = CheckpointStateAccessor.bind(mutation_graph, checkpointer, mode="delta") + replay_base_config = await branch_writer.aupdate( + _run_config("worker-branch"), + {"messages": Overwrite(list(source_pre_turn.values["messages"]))}, + as_node="branch", + ) + await branch_writer.aupdate( + replay_base_config, + {"messages": Overwrite(list(source_head.values["messages"]))}, + as_node="branch", + ) + + read_graph = _build_answer_graph(_DeltaChannelState, checkpointer, "unused") + read_accessor = CheckpointStateAccessor.bind(read_graph, checkpointer, mode="delta") + branch_history = await read_accessor.ahistory(_run_config("worker-branch"), limit=20) + base = next(snapshot for snapshot in branch_history if "h2" not in _ids(snapshot)) + + run_manager = RunManager() + record = await run_manager.create("worker-branch") + bridge = SimpleNamespace( + publish=AsyncMock(), + publish_end=AsyncMock(), + cleanup=AsyncMock(), + ) + created_graphs: list[Any] = [] + + def agent_factory(*, config): + del config + graph = _build_answer_graph(_DeltaChannelState, checkpointer, "a2-new") + created_graphs.append(graph) + return graph + + await asyncio.wait_for( + run_agent( + bridge, + run_manager, + record, + ctx=RunContext(checkpointer=checkpointer, checkpoint_channel_mode="delta"), + agent_factory=agent_factory, + graph_input={"messages": [HumanMessage(content="q2", id="h2")]}, + config=_run_config("worker-branch", base.config["configurable"]["checkpoint_id"]), + stream_modes=["values"], + ), + timeout=5, + ) + + assert record.status == RunStatus.success + final_accessor = CheckpointStateAccessor.bind(created_graphs[-1], checkpointer, mode="delta") + final = await final_accessor.aget(_run_config("worker-branch")) + assert _ids(final) == ["h1", "a1", "h2", "a2-new"] + + +async def test_run_agent_serializes_resume_preparation_with_checkpoint_writes(monkeypatch): + """Rollback capture and resume linearization share the checkpoint lock. + + A prior successful run can still be persisting duration metadata after its + active run slot is released. Resume preparation must wait for that writer + so its head snapshot and linear rewrite cannot interleave with the + metadata checkpoint's compare-and-write sequence. + """ + checkpointer = InMemorySaver() + graph = _build_answer_graph(_DeltaChannelState, checkpointer, "a1") + run_manager = RunManager() + record = await run_manager.create("worker-lock") + bridge = SimpleNamespace( + publish=AsyncMock(), + publish_end=AsyncMock(), + cleanup=AsyncMock(), + ) + factory_called = asyncio.Event() + startup_calls: list[str] = [] + + def agent_factory(*, config): + del config + factory_called.set() + return graph + + async def capture_rollback_point(*args, **kwargs): + del args, kwargs + startup_calls.append("capture") + return None + + async def linearize_resume(**kwargs): + del kwargs + startup_calls.append("linearize") + return None + + monkeypatch.setattr("deerflow.runtime.runs.worker._capture_rollback_point", capture_rollback_point) + monkeypatch.setattr("deerflow.runtime.runs.worker._linearize_delta_checkpoint_resume", linearize_resume) + + async with _checkpoint_thread_lock("worker-lock"): + run_task = asyncio.create_task( + run_agent( + bridge, + run_manager, + record, + ctx=RunContext(checkpointer=checkpointer, checkpoint_channel_mode="delta"), + agent_factory=agent_factory, + graph_input={"messages": [HumanMessage(content="q1", id="h1")]}, + config=_run_config("worker-lock"), + stream_modes=["values"], + ) + ) + await asyncio.wait_for(factory_called.wait(), timeout=1) + await asyncio.sleep(0) + assert startup_calls == [] + + await run_task + assert startup_calls == ["capture", "linearize"] + + +async def test_cancelled_delta_resume_rolls_back_to_pre_linearization_head(tmp_path): + """Rollback must restore the head captured before resume linearization. + + The worker deliberately captures its rollback point before rewriting a + selected delta checkpoint onto the head. If that ordering is reversed, + cancelling the resumed run would restore the selected historical state + and silently discard the abandoned turn that was present before the run. + """ + db_path = tmp_path / "worker-rollback-branch.sqlite3" + async with AsyncSqliteSaver.from_conn_string(str(db_path)) as checkpointer: + await checkpointer.setup() + _, source_head, source_pre_turn = await _seed_two_turns(checkpointer, _DeltaChannelState, "worker-rollback-parent") + + mutation_graph = build_state_mutation_graph("branch", "delta", _DeltaChannelState) + branch_writer = CheckpointStateAccessor.bind(mutation_graph, checkpointer, mode="delta") + replay_base_config = await branch_writer.aupdate( + _run_config("worker-rollback-branch"), + {"messages": Overwrite(list(source_pre_turn.values["messages"]))}, + as_node="branch", + ) + await branch_writer.aupdate( + replay_base_config, + {"messages": Overwrite(list(source_head.values["messages"]))}, + as_node="branch", + ) + + read_graph = _build_answer_graph(_DeltaChannelState, checkpointer, "unused") + read_accessor = CheckpointStateAccessor.bind(read_graph, checkpointer, mode="delta") + pre_run_head = await read_accessor.aget(_run_config("worker-rollback-branch")) + assert _ids(pre_run_head) == ["h1", "a1", "h2", "a2"] + branch_history = await read_accessor.ahistory(_run_config("worker-rollback-branch"), limit=20) + base = next(snapshot for snapshot in branch_history if "h2" not in _ids(snapshot)) + + run_manager = RunManager() + record = await run_manager.create("worker-rollback-branch") + bridge = SimpleNamespace( + publish=AsyncMock(), + publish_end=AsyncMock(), + cleanup=AsyncMock(), + ) + created_graphs: list[Any] = [] + + def agent_factory(*, config): + del config + + async def _answer_then_abort(state: dict[str, Any]) -> dict[str, Any]: + del state + record.abort_action = "rollback" + record.abort_event.set() + return {"messages": [AIMessage(content="answer for a2-new", id="a2-new")]} + + builder = StateGraph(_DeltaChannelState) + builder.add_node("answer", _answer_then_abort) + builder.set_entry_point("answer") + builder.set_finish_point("answer") + graph = builder.compile(checkpointer=checkpointer) + created_graphs.append(graph) + return graph + + await run_agent( + bridge, + run_manager, + record, + ctx=RunContext(checkpointer=checkpointer, checkpoint_channel_mode="delta"), + agent_factory=agent_factory, + graph_input={"messages": [HumanMessage(content="q2", id="h2")]}, + config=_run_config("worker-rollback-branch", base.config["configurable"]["checkpoint_id"]), + stream_modes=["values"], + ) + + assert record.status == RunStatus.error + final_accessor = CheckpointStateAccessor.bind(created_graphs[-1], checkpointer, mode="delta") + final = await final_accessor.aget(_run_config("worker-rollback-branch")) + assert _ids(final) == ["h1", "a1", "h2", "a2"] + + +async def test_full_mode_keeps_the_fork(): + """Full checkpoints carry complete channel values, so the fork materializes + correctly and LangGraph's branching semantics stay untouched.""" + checkpointer = InMemorySaver() + accessor, _, pre_turn = await _seed_two_turns(checkpointer, _FullChannelState, "thread-1") + config = _run_config("thread-1", pre_turn.config["configurable"]["checkpoint_id"]) + + resumed = await _linearize_delta_checkpoint_resume( + accessor=accessor, + checkpointer=checkpointer, + config=config, + thread_id="thread-1", + run_id="run-1", + ) + + assert resumed is None + assert config["configurable"]["checkpoint_id"] == pre_turn.config["configurable"]["checkpoint_id"] + + +async def test_ordinary_run_without_a_checkpoint_selector_is_untouched(): + checkpointer = InMemorySaver() + accessor, _, _ = await _seed_two_turns(checkpointer, _DeltaChannelState, "thread-1") + config = _run_config("thread-1") + + assert ( + await _linearize_delta_checkpoint_resume( + accessor=accessor, + checkpointer=checkpointer, + config=config, + thread_id="thread-1", + run_id="run-1", + ) + is None + ) + + +async def test_selecting_the_head_is_already_linear(): + """No sibling can exist under the head yet, so there is nothing to rewrite + and the thread keeps its checkpoint count.""" + checkpointer = InMemorySaver() + accessor, head, _ = await _seed_two_turns(checkpointer, _DeltaChannelState, "thread-1") + head_id = head.config["configurable"]["checkpoint_id"] + before = len(await accessor.ahistory(_run_config("thread-1"), limit=50)) + + assert ( + await _linearize_delta_checkpoint_resume( + accessor=accessor, + checkpointer=checkpointer, + config=_run_config("thread-1", head_id), + thread_id="thread-1", + run_id="run-1", + ) + is None + ) + assert len(await accessor.ahistory(_run_config("thread-1"), limit=50)) == before + + +async def test_unmaterializable_resume_state_fails_closed(): + """Falling back to the fork would persist the corrupted history this + exists to prevent, so an unreadable resume checkpoint raises.""" + checkpointer = InMemorySaver() + accessor, _, pre_turn = await _seed_two_turns(checkpointer, _DeltaChannelState, "thread-1") + + class _NoMessages: + def __init__(self, inner): + self._inner = inner + self.mode = inner.mode + self.graph = inner.graph + + async def aget(self, config): + snapshot = await self._inner.aget(config) + if config.get("configurable", {}).get("checkpoint_id"): + return type(snapshot)(**{**snapshot._asdict(), "values": {}}) + return snapshot + + with pytest.raises(RuntimeError, match="could not materialize resume checkpoint"): + await _linearize_delta_checkpoint_resume( + accessor=_NoMessages(accessor), + checkpointer=checkpointer, + config=_run_config("thread-1", pre_turn.config["configurable"]["checkpoint_id"]), + thread_id="thread-1", + run_id="run-1", + ) diff --git a/backend/tests/test_run_worker_rollback.py b/backend/tests/test_run_worker_rollback.py index dfd1aa493..81813b3dc 100644 --- a/backend/tests/test_run_worker_rollback.py +++ b/backend/tests/test_run_worker_rollback.py @@ -88,6 +88,7 @@ async def test_pending_cancel_stops_waiting_for_prior_finalization(): def _make_rollback_point(*, checkpoint_id="ckpt-1", messages=("before",), pending_writes=()): + materialized_messages = tuple(messages) return RollbackPoint( config={ "configurable": { @@ -96,7 +97,8 @@ def _make_rollback_point(*, checkpoint_id="ckpt-1", messages=("before",), pendin "checkpoint_id": checkpoint_id, } }, - messages=tuple(messages), + state_values={}, + messages=materialized_messages, metadata={"source": "input"}, pending_writes=tuple(pending_writes), ) @@ -1318,11 +1320,14 @@ async def test_rollback_propagates_aput_writes_failure(monkeypatch): @pytest.mark.anyio -async def test_rollback_forks_pre_run_lineage_for_delta_checkpoints(): - """Delta channels omit messages from checkpoint blobs, so rollback must - fork the pre-run lineage through the graph: the restored checkpoint's - messages are reconstructed by replaying ancestor writes, and writes from - the cancelled run (not ancestors of the fork) must not leak back in.""" +async def test_rollback_linearizes_delta_restore_onto_cancelled_head(): + """Delta rollback replaces state on the head instead of creating a fork. + + The cancelled path has already attached writes to the pre-run checkpoint, + so forking that checkpoint would replay sibling writes. The restored + checkpoint must instead descend from the cancelled head and replace the + captured messages there. + """ checkpointer = InMemorySaver() graph = _build_message_append_graph(_DeltaChannelState, checkpointer) accessor = CheckpointStateAccessor.bind(graph, checkpointer, mode="delta") @@ -1355,7 +1360,7 @@ async def test_rollback_forks_pre_run_lineage_for_delta_checkpoints(): assert [message.content for message in latest_snapshot.values["messages"]] == ["turn-0", "turn-1"] restored_tuple = await checkpointer.aget_tuple({"configurable": {"thread_id": "thread-1", "checkpoint_ns": "", "checkpoint_id": restored_checkpoint_id}}) - assert restored_tuple.parent_config["configurable"]["checkpoint_id"] == pre_run_checkpoint_id + assert restored_tuple.parent_config["configurable"]["checkpoint_id"] == cancelled_checkpoint_id assert restored_tuple.metadata.get("source") == "update" # A non-snapshot Delta checkpoint must not persist the full message list. raw_messages = restored_tuple.checkpoint.get("channel_values", {}).get("messages") @@ -1364,7 +1369,7 @@ async def test_rollback_forks_pre_run_lineage_for_delta_checkpoints(): @pytest.mark.anyio async def test_rollback_restores_pre_run_pending_writes_for_delta_checkpoints(): - """Pre-run pending writes are re-attached to the restored fork; writes + """Pre-run pending writes are re-attached to the restored checkpoint; writes attached to the cancelled run are not.""" checkpointer = InMemorySaver() graph = _build_message_append_graph(_DeltaChannelState, checkpointer) @@ -1408,8 +1413,8 @@ async def test_rollback_restores_pre_run_pending_writes_for_delta_checkpoints(): @pytest.mark.anyio -async def test_rollback_forks_pre_run_lineage_for_delta_checkpoints_sqlite_reopen(tmp_path): - """Same lineage contract against a disk-backed saver, verified after a +async def test_rollback_linearizes_delta_restore_sqlite_reopen(tmp_path): + """Same linear restore contract against a disk-backed saver, verified after a close/reopen so only persisted bytes can satisfy the assertions.""" db_path = tmp_path / "rollback.sqlite3" @@ -1426,6 +1431,8 @@ async def test_rollback_forks_pre_run_lineage_for_delta_checkpoints_sqlite_reope pre_run_checkpoint_id = rollback_point.config["configurable"]["checkpoint_id"] await graph.ainvoke({}, thread_config) # cancelled run + cancelled_snapshot = await accessor.aget(thread_config) + cancelled_checkpoint_id = cancelled_snapshot.config["configurable"]["checkpoint_id"] await _rollback_to_pre_run_checkpoint( accessor=accessor, @@ -1438,8 +1445,10 @@ async def test_rollback_forks_pre_run_lineage_for_delta_checkpoints_sqlite_reope restored_snapshot = await accessor.aget(thread_config) restored_checkpoint_id = restored_snapshot.config["configurable"]["checkpoint_id"] - assert restored_checkpoint_id != pre_run_checkpoint_id + assert restored_checkpoint_id not in (pre_run_checkpoint_id, cancelled_checkpoint_id) assert [message.content for message in restored_snapshot.values["messages"]] == ["turn-0", "turn-1"] + restored_tuple = await checkpointer.aget_tuple({"configurable": {"thread_id": "thread-1", "checkpoint_ns": "", "checkpoint_id": restored_checkpoint_id}}) + assert restored_tuple.parent_config["configurable"]["checkpoint_id"] == cancelled_checkpoint_id async with AsyncSqliteSaver.from_conn_string(str(db_path)) as checkpointer: graph = _build_message_append_graph(_DeltaChannelState, checkpointer) @@ -1451,7 +1460,7 @@ async def test_rollback_forks_pre_run_lineage_for_delta_checkpoints_sqlite_reope assert [message.content for message in latest_snapshot.values["messages"]] == ["turn-0", "turn-1"] restored_tuple = await checkpointer.aget_tuple({"configurable": {"thread_id": "thread-1", "checkpoint_ns": "", "checkpoint_id": restored_checkpoint_id}}) - assert restored_tuple.parent_config["configurable"]["checkpoint_id"] == pre_run_checkpoint_id + assert restored_tuple.parent_config["configurable"]["checkpoint_id"] == cancelled_checkpoint_id raw_messages = restored_tuple.checkpoint.get("channel_values", {}).get("messages") assert not isinstance(raw_messages, list)