mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-07-27 16:37:55 +00:00
fix(runtime): linearize delta-mode checkpoint resume (#4460)
* fix(runtime): linearize delta-mode checkpoint resume Resuming a run from an older checkpoint forks the lineage, and in delta mode that 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 replay into the fork, so the run starts from a message list that still contains the answer it was meant to replace — regenerating in a branched thread surfaced this as the superseded assistant message reappearing beside the new one after a reload. All three saver implementations are affected, so write-to-child ownership is a gap in the upstream delta contract rather than one saver's slip. Rather than reimplement that walk, express the fork as what it means: materialize the requested checkpoint's state, write it as an Overwrite on the current head (which has no siblings), and run linearly. The abandoned turn stays in history as the rewritten head's ancestry. This runs after the rollback point is captured, so cancel-with-rollback still restores the real pre-run head, and fails closed — an unreadable resume checkpoint raises instead of falling back to the corrupt fork. Full mode keeps forking: its checkpoints carry complete channel_values and need no replay. * fix(runtime): restore complete delta resume state * fix(runtime): linearize delta rollback restoration * Apply suggestions from code review Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> * fix(runtime): serialize delta resume preparation --------- Co-authored-by: Willem Jiang <willem.jiang@gmail.com> Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
This commit is contained in:
Aari
GitHub
parent
bb9f67aaf1
commit
244ce7739f
@ -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
|
||||
|
||||
@ -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):
|
||||
|
||||
495
backend/tests/test_run_worker_delta_resume.py
Normal file
495
backend/tests/test_run_worker_delta_resume.py
Normal file
@ -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",
|
||||
)
|
||||
@ -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)
|
||||
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user