mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-04-25 11:18:22 +00:00
35f141fc48
* feat: implement full checkpoint rollback on user cancellation - Capture pre-run checkpoint snapshot including checkpoint state, metadata, and pending_writes - Add _rollback_to_pre_run_checkpoint() function to restore thread state - Implement _call_checkpointer_method() helper to support both async and sync checkpointer methods - Rollback now properly restores checkpoint, metadata, channel_versions, and pending_writes - Remove obsolete TODO comment (Phase 2) as rollback is now complete This resolves the TODO(Phase 2) comment and enables full thread state restoration when a run is cancelled by the user. * fix: address rollback review feedback * fix: strengthen checkpoint rollback validation and error handling - Validate restored_config structure and checkpoint_id before use - Raise RuntimeError on malformed pending_writes instead of silent skip - Normalize None checkpoint_ns to empty string instead of "None" - Move delete_thread to only execute when pre_run_snapshot is None - Add docstring noting non-atomic rollback as known limitation This addresses review feedback on PR #1867 regarding data integrity in the checkpoint rollback implementation. * test: add comprehensive coverage for checkpoint rollback edge cases - test_rollback_restores_snapshot_without_deleting_thread - test_rollback_deletes_thread_when_no_snapshot_exists - test_rollback_raises_when_restore_config_has_no_checkpoint_id - test_rollback_normalizes_none_checkpoint_ns_to_root_namespace - test_rollback_raises_on_malformed_pending_write_not_a_tuple - test_rollback_raises_on_malformed_pending_write_non_string_channel - test_rollback_propagates_aput_writes_failure Covers all scenarios from PR #1867 review feedback. * test: format rollback worker tests
382 lines
15 KiB
Python
382 lines
15 KiB
Python
"""Background agent execution.
|
|
|
|
Runs an agent graph inside an ``asyncio.Task``, publishing events to
|
|
a :class:`StreamBridge` as they are produced.
|
|
|
|
Uses ``graph.astream(stream_mode=[...])`` which gives correct full-state
|
|
snapshots for ``values`` mode, proper ``{node: writes}`` for ``updates``,
|
|
and ``(chunk, metadata)`` tuples for ``messages`` mode.
|
|
|
|
Note: ``events`` mode is not supported through the gateway — it requires
|
|
``graph.astream_events()`` which cannot simultaneously produce ``values``
|
|
snapshots. The JS open-source LangGraph API server works around this via
|
|
internal checkpoint callbacks that are not exposed in the Python public API.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import copy
|
|
import inspect
|
|
import logging
|
|
from typing import Any, Literal
|
|
|
|
from deerflow.runtime.serialization import serialize
|
|
from deerflow.runtime.stream_bridge import StreamBridge
|
|
|
|
from .manager import RunManager, RunRecord
|
|
from .schemas import RunStatus
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Valid stream_mode values for LangGraph's graph.astream()
|
|
_VALID_LG_MODES = {"values", "updates", "checkpoints", "tasks", "debug", "messages", "custom"}
|
|
|
|
|
|
async def run_agent(
|
|
bridge: StreamBridge,
|
|
run_manager: RunManager,
|
|
record: RunRecord,
|
|
*,
|
|
checkpointer: Any,
|
|
store: Any | None = None,
|
|
agent_factory: Any,
|
|
graph_input: dict,
|
|
config: dict,
|
|
stream_modes: list[str] | None = None,
|
|
stream_subgraphs: bool = False,
|
|
interrupt_before: list[str] | Literal["*"] | None = None,
|
|
interrupt_after: list[str] | Literal["*"] | None = None,
|
|
) -> None:
|
|
"""Execute an agent in the background, publishing events to *bridge*."""
|
|
|
|
run_id = record.run_id
|
|
thread_id = record.thread_id
|
|
requested_modes: set[str] = set(stream_modes or ["values"])
|
|
pre_run_checkpoint_id: str | None = None
|
|
pre_run_snapshot: dict[str, Any] | None = None
|
|
snapshot_capture_failed = False
|
|
|
|
# Track whether "events" was requested but skipped
|
|
if "events" in requested_modes:
|
|
logger.info(
|
|
"Run %s: 'events' stream_mode not supported in gateway (requires astream_events + checkpoint callbacks). Skipping.",
|
|
run_id,
|
|
)
|
|
|
|
try:
|
|
# 1. Mark running
|
|
await run_manager.set_status(run_id, RunStatus.running)
|
|
|
|
# Snapshot the latest pre-run checkpoint so rollback can restore it.
|
|
if checkpointer is not None:
|
|
try:
|
|
config_for_check = {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}}
|
|
ckpt_tuple = await checkpointer.aget_tuple(config_for_check)
|
|
if ckpt_tuple is not None:
|
|
ckpt_config = getattr(ckpt_tuple, "config", {}).get("configurable", {})
|
|
pre_run_checkpoint_id = ckpt_config.get("checkpoint_id")
|
|
pre_run_snapshot = {
|
|
"checkpoint_ns": ckpt_config.get("checkpoint_ns", ""),
|
|
"checkpoint": copy.deepcopy(getattr(ckpt_tuple, "checkpoint", {})),
|
|
"metadata": copy.deepcopy(getattr(ckpt_tuple, "metadata", {})),
|
|
"pending_writes": copy.deepcopy(getattr(ckpt_tuple, "pending_writes", []) or []),
|
|
}
|
|
except Exception:
|
|
snapshot_capture_failed = True
|
|
logger.warning("Could not capture pre-run checkpoint snapshot for run %s", run_id, exc_info=True)
|
|
|
|
# 2. Publish metadata — useStream needs both run_id AND thread_id
|
|
await bridge.publish(
|
|
run_id,
|
|
"metadata",
|
|
{
|
|
"run_id": run_id,
|
|
"thread_id": thread_id,
|
|
},
|
|
)
|
|
|
|
# 3. Build the agent
|
|
from langchain_core.runnables import RunnableConfig
|
|
from langgraph.runtime import Runtime
|
|
|
|
# Inject runtime context so middlewares can access thread_id
|
|
# (langgraph-cli does this automatically; we must do it manually)
|
|
runtime = Runtime(context={"thread_id": thread_id}, store=store)
|
|
# If the caller already set a ``context`` key (LangGraph >= 0.6.0
|
|
# prefers it over ``configurable`` for thread-level data), make
|
|
# sure ``thread_id`` is available there too.
|
|
if "context" in config and isinstance(config["context"], dict):
|
|
config["context"].setdefault("thread_id", thread_id)
|
|
config.setdefault("configurable", {})["__pregel_runtime"] = runtime
|
|
|
|
runnable_config = RunnableConfig(**config)
|
|
agent = agent_factory(config=runnable_config)
|
|
|
|
# 4. Attach checkpointer and store
|
|
if checkpointer is not None:
|
|
agent.checkpointer = checkpointer
|
|
if store is not None:
|
|
agent.store = store
|
|
|
|
# 5. Set interrupt nodes
|
|
if interrupt_before:
|
|
agent.interrupt_before_nodes = interrupt_before
|
|
if interrupt_after:
|
|
agent.interrupt_after_nodes = interrupt_after
|
|
|
|
# 6. Build LangGraph stream_mode list
|
|
# "events" is NOT a valid astream mode — skip it
|
|
# "messages-tuple" maps to LangGraph's "messages" mode
|
|
lg_modes: list[str] = []
|
|
for m in requested_modes:
|
|
if m == "messages-tuple":
|
|
lg_modes.append("messages")
|
|
elif m == "events":
|
|
# Skipped — see log above
|
|
continue
|
|
elif m in _VALID_LG_MODES:
|
|
lg_modes.append(m)
|
|
if not lg_modes:
|
|
lg_modes = ["values"]
|
|
|
|
# Deduplicate while preserving order
|
|
seen: set[str] = set()
|
|
deduped: list[str] = []
|
|
for m in lg_modes:
|
|
if m not in seen:
|
|
seen.add(m)
|
|
deduped.append(m)
|
|
lg_modes = deduped
|
|
|
|
logger.info("Run %s: streaming with modes %s (requested: %s)", run_id, lg_modes, requested_modes)
|
|
|
|
# 7. Stream using graph.astream
|
|
if len(lg_modes) == 1 and not stream_subgraphs:
|
|
# Single mode, no subgraphs: astream yields raw chunks
|
|
single_mode = lg_modes[0]
|
|
async for chunk in agent.astream(graph_input, config=runnable_config, stream_mode=single_mode):
|
|
if record.abort_event.is_set():
|
|
logger.info("Run %s abort requested — stopping", run_id)
|
|
break
|
|
sse_event = _lg_mode_to_sse_event(single_mode)
|
|
await bridge.publish(run_id, sse_event, serialize(chunk, mode=single_mode))
|
|
else:
|
|
# Multiple modes or subgraphs: astream yields tuples
|
|
async for item in agent.astream(
|
|
graph_input,
|
|
config=runnable_config,
|
|
stream_mode=lg_modes,
|
|
subgraphs=stream_subgraphs,
|
|
):
|
|
if record.abort_event.is_set():
|
|
logger.info("Run %s abort requested — stopping", run_id)
|
|
break
|
|
|
|
mode, chunk = _unpack_stream_item(item, lg_modes, stream_subgraphs)
|
|
if mode is None:
|
|
continue
|
|
|
|
sse_event = _lg_mode_to_sse_event(mode)
|
|
await bridge.publish(run_id, sse_event, serialize(chunk, mode=mode))
|
|
|
|
# 8. Final status
|
|
if record.abort_event.is_set():
|
|
action = record.abort_action
|
|
if action == "rollback":
|
|
await run_manager.set_status(run_id, RunStatus.error, error="Rolled back by user")
|
|
try:
|
|
await _rollback_to_pre_run_checkpoint(
|
|
checkpointer=checkpointer,
|
|
thread_id=thread_id,
|
|
run_id=run_id,
|
|
pre_run_checkpoint_id=pre_run_checkpoint_id,
|
|
pre_run_snapshot=pre_run_snapshot,
|
|
snapshot_capture_failed=snapshot_capture_failed,
|
|
)
|
|
logger.info("Run %s rolled back to pre-run checkpoint %s", run_id, pre_run_checkpoint_id)
|
|
except Exception:
|
|
logger.warning("Failed to rollback checkpoint for run %s", run_id, exc_info=True)
|
|
else:
|
|
await run_manager.set_status(run_id, RunStatus.interrupted)
|
|
else:
|
|
await run_manager.set_status(run_id, RunStatus.success)
|
|
|
|
except asyncio.CancelledError:
|
|
action = record.abort_action
|
|
if action == "rollback":
|
|
await run_manager.set_status(run_id, RunStatus.error, error="Rolled back by user")
|
|
try:
|
|
await _rollback_to_pre_run_checkpoint(
|
|
checkpointer=checkpointer,
|
|
thread_id=thread_id,
|
|
run_id=run_id,
|
|
pre_run_checkpoint_id=pre_run_checkpoint_id,
|
|
pre_run_snapshot=pre_run_snapshot,
|
|
snapshot_capture_failed=snapshot_capture_failed,
|
|
)
|
|
logger.info("Run %s was cancelled and rolled back", run_id)
|
|
except Exception:
|
|
logger.warning("Run %s cancellation rollback failed", run_id, exc_info=True)
|
|
else:
|
|
await run_manager.set_status(run_id, RunStatus.interrupted)
|
|
logger.info("Run %s was cancelled", run_id)
|
|
|
|
except Exception as exc:
|
|
error_msg = f"{exc}"
|
|
logger.exception("Run %s failed: %s", run_id, error_msg)
|
|
await run_manager.set_status(run_id, RunStatus.error, error=error_msg)
|
|
await bridge.publish(
|
|
run_id,
|
|
"error",
|
|
{
|
|
"message": error_msg,
|
|
"name": type(exc).__name__,
|
|
},
|
|
)
|
|
|
|
finally:
|
|
await bridge.publish_end(run_id)
|
|
asyncio.create_task(bridge.cleanup(run_id, delay=60))
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
async def _call_checkpointer_method(checkpointer: Any, async_name: str, sync_name: str, *args: Any, **kwargs: Any) -> Any:
|
|
"""Call a checkpointer method, supporting async and sync variants."""
|
|
method = getattr(checkpointer, async_name, None) or getattr(checkpointer, sync_name, None)
|
|
if method is None:
|
|
raise AttributeError(f"Missing checkpointer method: {async_name}/{sync_name}")
|
|
result = method(*args, **kwargs)
|
|
if inspect.isawaitable(result):
|
|
return await result
|
|
return result
|
|
|
|
|
|
async def _rollback_to_pre_run_checkpoint(
|
|
*,
|
|
checkpointer: Any,
|
|
thread_id: str,
|
|
run_id: str,
|
|
pre_run_checkpoint_id: str | None,
|
|
pre_run_snapshot: dict[str, Any] | None,
|
|
snapshot_capture_failed: bool,
|
|
) -> None:
|
|
"""Restore thread state to the checkpoint snapshot captured before run start."""
|
|
if checkpointer is None:
|
|
logger.info("Run %s rollback requested but no checkpointer is configured", run_id)
|
|
return
|
|
|
|
if snapshot_capture_failed:
|
|
logger.warning("Run %s rollback skipped: pre-run checkpoint snapshot capture failed", run_id)
|
|
return
|
|
|
|
if pre_run_snapshot is None:
|
|
await _call_checkpointer_method(checkpointer, "adelete_thread", "delete_thread", thread_id)
|
|
logger.info("Run %s rollback reset thread %s to empty state", run_id, thread_id)
|
|
return
|
|
|
|
checkpoint_to_restore = None
|
|
metadata_to_restore: dict[str, Any] = {}
|
|
checkpoint_ns = ""
|
|
checkpoint = pre_run_snapshot.get("checkpoint")
|
|
if not isinstance(checkpoint, dict):
|
|
logger.warning("Run %s rollback skipped: invalid pre-run checkpoint snapshot", run_id)
|
|
return
|
|
checkpoint_to_restore = checkpoint
|
|
if checkpoint_to_restore.get("id") is None and pre_run_checkpoint_id is not None:
|
|
checkpoint_to_restore = {**checkpoint_to_restore, "id": pre_run_checkpoint_id}
|
|
if checkpoint_to_restore.get("id") is None:
|
|
logger.warning("Run %s rollback skipped: pre-run checkpoint has no checkpoint id", run_id)
|
|
return
|
|
metadata = pre_run_snapshot.get("metadata", {})
|
|
metadata_to_restore = metadata if isinstance(metadata, dict) else {}
|
|
raw_checkpoint_ns = pre_run_snapshot.get("checkpoint_ns")
|
|
checkpoint_ns = raw_checkpoint_ns if isinstance(raw_checkpoint_ns, str) else ""
|
|
|
|
channel_versions = checkpoint_to_restore.get("channel_versions")
|
|
new_versions = dict(channel_versions) if isinstance(channel_versions, dict) else {}
|
|
|
|
restore_config = {"configurable": {"thread_id": thread_id, "checkpoint_ns": checkpoint_ns}}
|
|
restored_config = await _call_checkpointer_method(
|
|
checkpointer,
|
|
"aput",
|
|
"put",
|
|
restore_config,
|
|
checkpoint_to_restore,
|
|
metadata_to_restore if isinstance(metadata_to_restore, dict) else {},
|
|
new_versions,
|
|
)
|
|
if not isinstance(restored_config, dict):
|
|
raise RuntimeError(f"Run {run_id} rollback restore returned invalid config: expected dict")
|
|
restored_configurable = restored_config.get("configurable", {})
|
|
if not isinstance(restored_configurable, dict):
|
|
raise RuntimeError(f"Run {run_id} rollback restore returned invalid config payload")
|
|
restored_checkpoint_id = restored_configurable.get("checkpoint_id")
|
|
if not restored_checkpoint_id:
|
|
raise RuntimeError(f"Run {run_id} rollback restore did not return checkpoint_id")
|
|
|
|
pending_writes = pre_run_snapshot.get("pending_writes", [])
|
|
if not pending_writes:
|
|
return
|
|
|
|
writes_by_task: dict[str, list[tuple[str, Any]]] = {}
|
|
for item in pending_writes:
|
|
if not isinstance(item, (tuple, list)) or len(item) != 3:
|
|
raise RuntimeError(f"Run {run_id} rollback failed: pending_write is not a 3-tuple: {item!r}")
|
|
task_id, channel, value = item
|
|
if not isinstance(channel, str):
|
|
raise RuntimeError(f"Run {run_id} rollback failed: pending_write has non-string channel: task_id={task_id!r}, channel={channel!r}")
|
|
writes_by_task.setdefault(str(task_id), []).append((channel, value))
|
|
|
|
for task_id, writes in writes_by_task.items():
|
|
await _call_checkpointer_method(
|
|
checkpointer,
|
|
"aput_writes",
|
|
"put_writes",
|
|
restored_config,
|
|
writes,
|
|
task_id=task_id,
|
|
)
|
|
|
|
|
|
def _lg_mode_to_sse_event(mode: str) -> str:
|
|
"""Map LangGraph internal stream_mode name to SSE event name.
|
|
|
|
LangGraph's ``astream(stream_mode="messages")`` produces message
|
|
tuples. The SSE protocol calls this ``messages-tuple`` when the
|
|
client explicitly requests it, but the default SSE event name used
|
|
by LangGraph Platform is simply ``"messages"``.
|
|
"""
|
|
# All LG modes map 1:1 to SSE event names — "messages" stays "messages"
|
|
return mode
|
|
|
|
|
|
def _unpack_stream_item(
|
|
item: Any,
|
|
lg_modes: list[str],
|
|
stream_subgraphs: bool,
|
|
) -> tuple[str | None, Any]:
|
|
"""Unpack a multi-mode or subgraph stream item into (mode, chunk).
|
|
|
|
Returns ``(None, None)`` if the item cannot be parsed.
|
|
"""
|
|
if stream_subgraphs:
|
|
if isinstance(item, tuple) and len(item) == 3:
|
|
_ns, mode, chunk = item
|
|
return str(mode), chunk
|
|
if isinstance(item, tuple) and len(item) == 2:
|
|
mode, chunk = item
|
|
return str(mode), chunk
|
|
return None, None
|
|
|
|
if isinstance(item, tuple) and len(item) == 2:
|
|
mode, chunk = item
|
|
return str(mode), chunk
|
|
|
|
# Fallback: single-element output from first mode
|
|
return lg_modes[0] if lg_modes else None, item
|