mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-26 21:26:39 +00:00
fix(runtime): defer host cancellation through the whole terminal finalization (#6914)
* fix(runtime): defer host cancellation through the whole terminal finalization Cancelling run_agent() during an early terminal-finalization await (e.g. the delivery-receipt write) escaped the surrounding `except Exception` guards, because asyncio.CancelledError derives from BaseException. The cancellation then skipped the remaining terminal tail: terminal status persistence, the completion hook, task-stop notification, and the stream END marker. Add `_await_despite_host_cancellation`, which shields each terminal step so the started awaitable keeps running under host cancellation, defers the interruption (uncancelling the task), and re-awaits the shield until the step completes. Apply it to the durable finalization stages; the original interruption is still re-raised after publish_end. * fix(runtime): bound terminal deferral and shield abort-path cancellation * fix(runtime): retain owned finalization through cancellation --------- Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
This commit is contained in:
万家良
GitHub
parent
33a3970c6b
commit
4f1c848744
@ -2281,6 +2281,8 @@ async with runtime:
|
||||
|
||||
`SubagentRuntime.stop()` waits for its owned batch service to finish before propagating caller cancellation, including repeated cancellation. This drain has no timeout; repository operations and child cleanup must terminate. If service shutdown also fails or is cancelled, the first caller cancellation is preserved with the service failure as its cause.
|
||||
|
||||
Gateway run finalization also retains ownership through caller cancellation: it finishes terminal persistence, stream completion, and resource cleanup before propagating the interruption. Backend I/O must terminate or enforce its own timeout. Gateway shutdown includes workers still finalizing after their terminal status commits; an expired shutdown drain budget is logged and does not cancel their outstanding writes.
|
||||
|
||||
The factory still does not load YAML or create SQL infrastructure: the caller supplies the config snapshot, repository, and lifecycle. Because it accepts a caller-owned `system_prompt`, direct integrations also own any model-visible wording about those limits; the default middleware enforces the runtime limits regardless. The factory does not mount the Gateway owner-scoped HTTP routes or Web UI, so direct applications must expose their own result API/UI if they need those surfaces. For ordinary delegation only, `SubagentRuntime(...)` needs no asynchronous startup.
|
||||
|
||||
Administrators can add, edit, disable, and delete reusable worker definitions from **Settings → Subagents**. Built-in and `config.yaml` definitions remain visible there as read-only entries. The default Lead Agent can use every enabled runtime sub-agent; each page-created Custom Agent can instead allow all, none, or a selected set. That selection is enforced both in the model-visible directory and by the server-side `task` tool. Managed definitions are deployment-wide in this version and follow `agent_storage.backend`: atomic files for a local deployment or the shared application database for multiple instances.
|
||||
|
||||
@ -151,7 +151,7 @@ The first `RunManager.list_by_thread()` hydration page uses a 100-row floor or
|
||||
the number of required IDs, whichever is larger; missing exact runs use targeted
|
||||
`get()` calls.
|
||||
|
||||
**Terminal run cleanup explicitly breaks graph-scoped references while preserving the existing `RunRecord` grace period.** Every `agent.astream()` iterator is closed in `_stream_once`, including abort/exception/early-break paths. A close failure after an abort is warning-only and cannot replace the user-requested `interrupted` outcome; normal-completion close failures still surface, and an in-flight stream exception remains authoritative over a secondary close failure. Journal construction and cancellable preflight work (including MCP task projection and the prior-finalization wait) live inside the worker's guarded body, so cancellation before agent startup still terminalizes the run and closes its stream. `run_agent()` wraps the complete terminal-finalization sequence in an outer teardown guard, so cancellation or failure from any terminal-stage await cannot skip `RunJournal.close()`, removal of the journal, `__pregel_runtime`, and internal runtime-context values from every runnable config, or release of local graph/payload references. That guard schedules bridge cleanup, run-record cleanup, and cyclic GC even when interruption happens before the terminal stream marker or terminal publication itself fails, so neither a cancelled observer nor a delivery-backend outage can strand process-local run state. A non-`Exception` `BaseException` caught while awaiting the completion hook or task-stop notification (including host-task cancellation) is deferred through the ordinary remaining finalization, with the first interruption preserved and every caught host-task `CancelledError` balanced by calling `Task.uncancel()` until the current task’s cumulative cancellation count is clear. Task-stop fan-out runs in one child task and every host wait uses `shield`, so repeated cancellation of the worker cannot cancel that fan-out or skip later observers; the worker keeps awaiting the same child task. A rogue observer that raises its own `CancelledError` remains contained by the extension dispatcher and distinguishable from host cancellation. This guarantee applies only to cancellation caught during those hook stages: clearing the finalizing barrier and publishing END remain direct awaits, so another cancellation in the subsequent critical tail retains forceful-termination semantics instead of creating an unbounded shield. If that tail completes without another interruption, the first deferred interruption is re-raised after END; a barrier-clear failure prevents END publication, while an END failure is raised after the barrier is clear. `RunJournal.flush()` clears its `_pending_progress_task` after awaiting or cancelling it; ordinary `close()` detaches the event store/progress reporter and clears callback bookkeeping only after that flush succeeds, preserving the buffer for retry on a transient store failure. A fenced worker instead calls `close(flush=False)`, which cancels pending journal work and detaches without initiating another event-store write after lease ownership is lost; its final detach runs even if a second cancellation interrupts pending-task shutdown. `RunManager.cleanup(run_id)` retains the process-local `RunRecord`, completed task, and request payload for its default 300-second local join/status window before releasing them, and evicts only when a `RunStore` backs the manager — without one there is no hydration fallback, so the record is retained rather than silently dropping the run's history. Durable history remains in `RunStore`; `StreamBridge` data keeps its separate 60-second late-subscriber window, and both cleanup coroutines run in a fresh empty `contextvars.Context`. A contextless full cyclic-GC pass, coalesced to at most once every 10 seconds and dispatched through the default executor, bounds the lifetime of unreachable LangGraph callback/loop cycles without synchronously walking the heap in the event-loop timer; passes taking at least 100 ms are logged at INFO because CPython GC may still impose interpreter-level pauses.
|
||||
**Terminal run cleanup explicitly breaks graph-scoped references while preserving the existing `RunRecord` grace period.** Every `agent.astream()` iterator is closed in `_stream_once`, including abort/exception/early-break paths. A close failure after an abort is warning-only and cannot replace the user-requested `interrupted` outcome; normal-completion close failures still surface, and an in-flight stream exception remains authoritative over a secondary close failure. Journal construction and cancellable preflight work (including MCP task projection and the prior-finalization wait) live inside the worker's guarded body, so cancellation before agent startup still terminalizes the run and closes its stream. `run_agent()` wraps the complete terminal-finalization sequence in an outer teardown guard, so cancellation or failure from any terminal-stage await cannot skip `RunJournal.close()`, removal of the journal, `__pregel_runtime`, and internal runtime-context values from every runnable config, or release of local graph/payload references. That guard schedules bridge cleanup, run-record cleanup, and cyclic GC even when interruption happens before the terminal stream marker or terminal publication itself fails, so neither a cancelled observer nor a delivery-backend outage can strand process-local run state. The registered worker owns a named `run-finalization-<run_id>` task through the whole terminal sequence and resource teardown. Every host wait shields and drains that same task across repeated cancellation; there is no per-step cancellation deadline that can detach a write. Backend operations must provide their own I/O bounds. The first host cancellation stays in the worker, independent of step results, and is re-raised after finalization settles; an escaping finalizer failure is retained as its cause. The optional completion hook runs in an owned child task and still receives host cancellation, including one deferred before the hook starts; the finalizer awaits its settlement before task-stop observers and END. Task-stop fan-out retains its existing observer timeout and rogue-cancellation containment. Clearing the finalizing barrier and publishing END preserve their ordering: a barrier-clear failure prevents END, while END failure surfaces after the barrier is clear. Shutdown includes every terminal worker whose task is still active, even after its durable status commits. Its existing bounded drain logs pending workers and retains their ownership; expiry is not permission to cancel or retire a finalizer. `RunJournal.flush()` clears its `_pending_progress_task` after awaiting or cancelling it; ordinary `close()` detaches the event store/progress reporter and clears callback bookkeeping only after that flush succeeds, preserving the buffer for retry on a transient store failure. A fenced worker instead calls `close(flush=False)`, which cancels pending journal work and detaches without initiating another event-store write after lease ownership is lost; its final detach runs even if a second cancellation interrupts pending-task shutdown. `RunManager.cleanup(run_id)` retains the process-local `RunRecord`, completed task, and request payload for its default 300-second local join/status window before releasing them, and evicts only when a `RunStore` backs the manager — without one there is no hydration fallback, so the record is retained rather than silently dropping the run's history. Durable history remains in `RunStore`; `StreamBridge` data keeps its separate 60-second late-subscriber window, and both cleanup coroutines run in a fresh empty `contextvars.Context`. A contextless full cyclic-GC pass, coalesced to at most once every 10 seconds and dispatched through the default executor, bounds the lifetime of unreachable LangGraph callback/loop cycles without synchronously walking the heap in the event-loop timer; passes taking at least 100 ms are logged at INFO because CPython GC may still impose interpreter-level pauses.
|
||||
|
||||
**Checkpoint channel benchmark**: `scripts/benchmark/checkpoint/bench_channels.py`
|
||||
runs paired `full`/`delta` message-only StateGraphs in a fresh child process per
|
||||
|
||||
@ -85,10 +85,13 @@ writes. `RunRecord.terminal_commit_pending` keeps `_renew_leases()` renewing tha
|
||||
still-active row until the commit is attempted; a renewal rejected by the worker's
|
||||
own commit is confirmed by re-reading the row, while a peer claim fences the run.
|
||||
Never select runs for renewal by local status alone. Shutdown's drain treats the
|
||||
same state as active: such runs are awaited (never cancelled — cancellation would
|
||||
skip their terminal tail) so the duration checkpoint and terminal commit land
|
||||
before checkpointer teardown. A drain timeout leaves staged statuses and commit
|
||||
barriers intact without persisting `interrupted`. A finalizer that raises is
|
||||
same state as active, and also awaits committed terminal workers still executing
|
||||
hooks, END or resource cleanup. The registered worker shields and drains one owned
|
||||
finalizer across repeated host cancellation, preserving the first interruption even
|
||||
if a step fails. There is no deadline that cancels or detaches writes; backend I/O
|
||||
must be bounded by its backend. Only the optional completion hook receives host
|
||||
cancellation, and its child task is drained before the terminal tail continues.
|
||||
A drain timeout leaves staged statuses and commit barriers intact without persisting `interrupted`. A finalizer that raises is
|
||||
logged with its run ID, staged status and exception while the other runs continue
|
||||
draining. Coverage: `tests/test_gateway_run_drain_shutdown.py`.
|
||||
|
||||
|
||||
@ -2552,11 +2552,10 @@ class RunManager:
|
||||
during the drain keeps its real terminal status instead of being
|
||||
blanket-overwritten. Runs whose terminal status is staged in memory with
|
||||
a deferred commit are still writing their final duration checkpoint and
|
||||
then commit that status, so they are awaited in the same bounded drain —
|
||||
without cancellation, which would skip their terminal tail (#5542),
|
||||
exactly as ``_renew_leases`` keeps them active through
|
||||
``_awaits_terminal_commit``. The whole drain, including the trailing status
|
||||
persistence, is bounded by ``timeout`` so a run stuck in cleanup (or a
|
||||
then commit that status. Committed terminal workers may still be running
|
||||
hooks, publishing END or closing resources. Both remain visible in this
|
||||
bounded drain and are awaited without cancellation (#5542). The whole
|
||||
drain, including the trailing status persistence, is bounded by ``timeout`` so a run stuck in cleanup (or a
|
||||
slow store under DB pressure) cannot hang worker shutdown — the
|
||||
precondition for the signal-reentrancy deadlock guarded by
|
||||
``app.gateway.app._SHUTDOWN_HOOK_TIMEOUT_SECONDS``. Runs still active
|
||||
@ -2567,12 +2566,11 @@ class RunManager:
|
||||
|
||||
async with self._lock:
|
||||
inflight = [record for record in self._runs.values() if record.status in (RunStatus.pending, RunStatus.running) and record.task is not None and not record.task.done()]
|
||||
# A run whose terminal status is staged in memory (deferred commit)
|
||||
# keeps writing the duration checkpoint before committing that
|
||||
# status. Drain it too — but never cancel it — so the write lands
|
||||
# while the checkpointer is still open. ``_renew_leases`` already
|
||||
# treats these runs as active through ``_awaits_terminal_commit``.
|
||||
finalizing = [record for record in self._runs.values() if self._awaits_terminal_commit(record) and record.status not in (RunStatus.pending, RunStatus.running) and record.task is not None and not record.task.done()]
|
||||
# Terminal workers still own finalization, including after the
|
||||
# durable commit (hooks, END and journal/sandbox cleanup). Await them
|
||||
# without cancellation so resource teardown cannot overtake a worker
|
||||
# that settles within this drain's budget.
|
||||
finalizing = [record for record in self._runs.values() if record.status not in (RunStatus.pending, RunStatus.running) and record.task is not None and not record.task.done()]
|
||||
finalizing_ids = {record.run_id for record in finalizing}
|
||||
for record in inflight:
|
||||
record.abort_action = "interrupt"
|
||||
|
||||
@ -890,6 +890,35 @@ async def _await_task_stop_after_host_cancellation(
|
||||
deferred = _defer_finalization_interrupt(deferred, exc)
|
||||
|
||||
|
||||
async def _drain_run_finalization(task: asyncio.Task[None], *, on_host_cancel: Callable[[], None] | None = None) -> None:
|
||||
"""Retain ownership through repeated cancellation, including failed cleanup.
|
||||
|
||||
The registered worker stays alive until its finalizer settles, keeping it
|
||||
visible to shutdown. There is deliberately no timeout: abandoning an owned
|
||||
write would let resource teardown race it. Backends must bound their own I/O.
|
||||
"""
|
||||
deferred: BaseException | None = None
|
||||
while True:
|
||||
try:
|
||||
await asyncio.shield(task)
|
||||
break
|
||||
except asyncio.CancelledError as exc:
|
||||
host = asyncio.current_task()
|
||||
if host is None or not host.cancelling():
|
||||
if deferred is not None:
|
||||
raise deferred from exc
|
||||
raise
|
||||
deferred = _defer_finalization_interrupt(deferred, exc)
|
||||
if on_host_cancel is not None:
|
||||
on_host_cancel()
|
||||
except BaseException as exc:
|
||||
if deferred is not None:
|
||||
raise deferred from exc
|
||||
raise
|
||||
if deferred is not None:
|
||||
raise deferred
|
||||
|
||||
|
||||
async def run_agent(
|
||||
bridge: StreamBridge,
|
||||
run_manager: RunManager,
|
||||
@ -1648,367 +1677,375 @@ async def run_agent(
|
||||
)
|
||||
|
||||
finally:
|
||||
try:
|
||||
if record.ownership_lost:
|
||||
logger.warning(
|
||||
"Skipping durable finalization for run %s because this worker no longer owns its lease",
|
||||
run_id,
|
||||
)
|
||||
completion_hook_task: asyncio.Task[None] | None = None
|
||||
host_interrupted = False
|
||||
|
||||
if not record.ownership_lost and _is_edit_replay_run(record) and record.status != RunStatus.success:
|
||||
if not record.finalizing:
|
||||
await run_manager.set_finalizing(run_id, True)
|
||||
try:
|
||||
if not checkpoint_rollback_completed:
|
||||
checkpoint_rollback_completed = await _rollback_to_pre_run_checkpoint(
|
||||
accessor=accessor,
|
||||
checkpointer=checkpointer,
|
||||
thread_id=thread_id,
|
||||
run_id=run_id,
|
||||
rollback_point=rollback_point,
|
||||
snapshot_capture_failed=snapshot_capture_failed,
|
||||
)
|
||||
if checkpoint_rollback_completed:
|
||||
await _publish_restored_checkpoint_values(
|
||||
bridge=bridge,
|
||||
run_id=run_id,
|
||||
accessor=accessor,
|
||||
thread_id=thread_id,
|
||||
)
|
||||
logger.info("Run %s edit replay restored pre-run checkpoint %s", run_id, pre_run_checkpoint_id)
|
||||
except Exception:
|
||||
logger.warning("Run %s edit replay rollback failed", run_id, exc_info=True)
|
||||
def interrupt_completion_hook() -> None:
|
||||
nonlocal host_interrupted
|
||||
host_interrupted = True
|
||||
if completion_hook_task is not None:
|
||||
completion_hook_task.cancel()
|
||||
|
||||
# Persist any subagent step events still buffered (#3779) — including on
|
||||
# abort/exception paths, where the stream loop broke before its own flush.
|
||||
if not record.ownership_lost and subagent_events is not None:
|
||||
await subagent_events.flush()
|
||||
async def finalize() -> None:
|
||||
nonlocal completion_hook_task
|
||||
nonlocal accessor, agent, agent_result, cancel_action, checkpoint_rollback_completed
|
||||
nonlocal completion_data, deferred_finalization_interrupt, delivery_content, goal_completion, goal_evaluator_model
|
||||
nonlocal graph_input, pre_run_workspace_snapshot, produced_output_paths, rollback_point, runtime
|
||||
nonlocal runtime_ctx, subagent_events, task_info, task_store
|
||||
|
||||
if not record.ownership_lost and event_store is not None and pre_run_workspace_snapshot is not None:
|
||||
try:
|
||||
await record_workspace_changes(
|
||||
event_store,
|
||||
thread_id,
|
||||
try:
|
||||
if record.ownership_lost:
|
||||
logger.warning(
|
||||
"Skipping durable finalization for run %s because this worker no longer owns its lease",
|
||||
run_id,
|
||||
pre_run_workspace_snapshot,
|
||||
user_id=workspace_changes_user_id,
|
||||
extra_excluded_dir_names=workspace_excluded_dir_names,
|
||||
)
|
||||
except Exception:
|
||||
logger.warning("Failed to record workspace changes for run %s", run_id, exc_info=True)
|
||||
|
||||
# Flush buffered journal events before the terminal receipt. The
|
||||
# receipt uses a run-scoped idempotent write shared with recovery, then
|
||||
# the staged terminal status is persisted. This ordering closes the
|
||||
# crash window where a terminal run could otherwise outlive its receipt.
|
||||
# A fenced worker leaves receipt recovery to the peer that claimed it.
|
||||
if not record.ownership_lost and journal is not None:
|
||||
try:
|
||||
await journal.flush()
|
||||
except Exception:
|
||||
logger.warning("Failed to flush journal for run %s", run_id, exc_info=True)
|
||||
if not record.ownership_lost and _is_edit_replay_run(record) and record.status != RunStatus.success:
|
||||
if not record.finalizing:
|
||||
await run_manager.set_finalizing(run_id, True)
|
||||
try:
|
||||
if not checkpoint_rollback_completed:
|
||||
checkpoint_rollback_completed = await _rollback_to_pre_run_checkpoint(
|
||||
accessor=accessor,
|
||||
checkpointer=checkpointer,
|
||||
thread_id=thread_id,
|
||||
run_id=run_id,
|
||||
rollback_point=rollback_point,
|
||||
snapshot_capture_failed=snapshot_capture_failed,
|
||||
)
|
||||
if checkpoint_rollback_completed:
|
||||
await _publish_restored_checkpoint_values(
|
||||
bridge=bridge,
|
||||
run_id=run_id,
|
||||
accessor=accessor,
|
||||
thread_id=thread_id,
|
||||
)
|
||||
logger.info("Run %s edit replay restored pre-run checkpoint %s", run_id, pre_run_checkpoint_id)
|
||||
except Exception:
|
||||
logger.warning("Run %s edit replay rollback failed", run_id, exc_info=True)
|
||||
|
||||
if delivery_content is None:
|
||||
if produced_output_paths is None:
|
||||
produced_output_paths = await _produced_output_paths(
|
||||
# Persist any subagent step events still buffered (#3779) — including on
|
||||
# abort/exception paths, where the stream loop broke before its own flush.
|
||||
if not record.ownership_lost and subagent_events is not None:
|
||||
await subagent_events.flush()
|
||||
|
||||
if not record.ownership_lost and event_store is not None and pre_run_workspace_snapshot is not None:
|
||||
try:
|
||||
await record_workspace_changes(
|
||||
event_store,
|
||||
thread_id,
|
||||
run_id,
|
||||
pre_run_workspace_snapshot,
|
||||
thread_id=thread_id,
|
||||
user_id=workspace_changes_user_id,
|
||||
extra_excluded_dir_names=workspace_excluded_dir_names,
|
||||
)
|
||||
delivery_content = _delivery_content_with_outputs(journal.get_delivery_content(), produced_output_paths)
|
||||
receipt_persisted = await _persist_delivery_receipt(
|
||||
event_store,
|
||||
thread_id=thread_id,
|
||||
run_id=run_id,
|
||||
content=delivery_content,
|
||||
)
|
||||
if produced_output_paths and record.status == RunStatus.success and not receipt_persisted:
|
||||
await run_manager.set_status(
|
||||
run_id,
|
||||
RunStatus.error,
|
||||
error=_DELIVERY_RECEIPT_FAILED_ERROR,
|
||||
persist=False,
|
||||
)
|
||||
|
||||
if not record.ownership_lost and journal is not None and persist_completion:
|
||||
try:
|
||||
# Advance the final completion fields and timestamp without
|
||||
# terminalizing the durable row. That active row continues to
|
||||
# fence peer checkpoint writers through the duration write.
|
||||
completion_data = journal.get_completion_data()
|
||||
await run_manager.update_finalizing_progress(run_id, **completion_data)
|
||||
except Exception:
|
||||
logger.warning("Failed to persist finalizing run progress for %s (non-fatal)", run_id, exc_info=True)
|
||||
|
||||
# Keep the durable run row active through its final duration checkpoint
|
||||
# write. A peer Gateway admits history migration from the durable row,
|
||||
# not this worker's staged terminal status; terminalizing first would
|
||||
# let that migration read an unfinished lifetime and race this write.
|
||||
if started and not record.ownership_lost and checkpointer is not None and record.status == RunStatus.success:
|
||||
try:
|
||||
created = datetime.fromisoformat(record.created_at.replace("Z", "+00:00"))
|
||||
updated = datetime.fromisoformat(record.updated_at.replace("Z", "+00:00"))
|
||||
# Match legacy history semantics: turn_duration is the whole
|
||||
# RunRecord lifetime in integer seconds, including admission
|
||||
# delay. Persist zero for sub-second successful turns.
|
||||
duration = max(0, int((updated - created).total_seconds()))
|
||||
await _persist_run_duration(
|
||||
checkpointer=checkpointer,
|
||||
thread_id=thread_id,
|
||||
run_id=run_id,
|
||||
duration_seconds=duration,
|
||||
)
|
||||
except Exception:
|
||||
logger.debug("Failed to persist run duration for thread %s run %s (non-fatal)", thread_id, run_id)
|
||||
|
||||
if scheduled_goal is not None and not record.ownership_lost and checkpointer is not None:
|
||||
if record.abort_event.is_set() and record.abort_action == "rollback" and not checkpoint_rollback_completed:
|
||||
await _finish_cancellation("rollback")
|
||||
# Scheduled goals are occurrence-scoped. Keep the real run slot
|
||||
# active through cleanup so a peer's next user turn cannot
|
||||
# inherit this goal. Ordinary goals retain terminal-first cleanup.
|
||||
cleanup_task = asyncio.create_task(
|
||||
_clear_scheduled_goal(goal=scheduled_goal, record=record, run_manager=run_manager, bridge=bridge, checkpointer=checkpointer, owned_admission=True),
|
||||
name=f"scheduled-goal-cleanup-{run_id}",
|
||||
)
|
||||
try:
|
||||
deferred_finalization_interrupt = await _await_task_stop_after_host_cancellation(cleanup_task, deferred_finalization_interrupt)
|
||||
except Exception:
|
||||
logger.warning("Could not remove scheduled goal for run %s", run_id, exc_info=True)
|
||||
await run_manager.set_status(run_id, RunStatus.error, error="Scheduled goal cleanup failed", persist=False)
|
||||
|
||||
if not record.ownership_lost and (event_store is not None or scheduled_goal is not None):
|
||||
try:
|
||||
# Even after bounded receipt retries are exhausted, persist the
|
||||
# real worker outcome. Leaving a successful row inflight would
|
||||
# let lease recovery rewrite it as an error with a synthetic
|
||||
# zero receipt.
|
||||
if record.abort_event.is_set():
|
||||
if scheduled_goal is not None:
|
||||
# A cancel may arrive while scheduled cleanup was
|
||||
# awaiting its saver. The local status may still be
|
||||
# staged success; apply the accepted action, including
|
||||
# rollback, before committing any terminal outcome.
|
||||
await _finish_cancellation(record.abort_action, restore_checkpoint=not checkpoint_rollback_completed)
|
||||
await run_manager.persist_current_status(run_id)
|
||||
else:
|
||||
cancel_action = await run_manager.set_status_if_not_cancelled(
|
||||
run_id,
|
||||
record.status,
|
||||
error=record.error,
|
||||
stop_reason=record.stop_reason,
|
||||
goal_verdict=record.goal_verdict,
|
||||
)
|
||||
if cancel_action is not None:
|
||||
await _finish_cancellation(cancel_action)
|
||||
await run_manager.persist_current_status(run_id)
|
||||
except Exception:
|
||||
logger.warning("Failed to persist terminal status for run %s after delivery receipt attempts", run_id, exc_info=True)
|
||||
# The deferred commit has been attempted. A failed write is left to
|
||||
# completion fallback or lease recovery, as before.
|
||||
record.terminal_commit_pending = False
|
||||
|
||||
if not record.ownership_lost and journal is not None and persist_completion:
|
||||
try:
|
||||
# Persist token usage + convenience fields to RunStore
|
||||
completion_data = completion_data or journal.get_completion_data()
|
||||
await run_manager.update_run_completion(run_id, status=record.status.value, **completion_data)
|
||||
except Exception:
|
||||
logger.warning("Failed to persist run completion for %s (non-fatal)", run_id, exc_info=True)
|
||||
|
||||
if scheduled_goal is not None and not record.ownership_lost:
|
||||
if not await run_manager.persist_current_status(run_id):
|
||||
await run_manager._mark_ownership_lost(record, reason="Scheduled run terminal status could not be confirmed after goal cleanup.", require_active=False)
|
||||
# Late cancellation can still drain a rollback after goal
|
||||
# cleanup. Retain admission/heartbeat protection through that
|
||||
# rollback and the actual durable terminal acknowledgement.
|
||||
record.scheduled_goal_cleanup_pending = False
|
||||
|
||||
# A satisfied evaluator is only a candidate until artifact delivery,
|
||||
# receipt persistence and durable cancellation arbitration have ended.
|
||||
# Status writes are best-effort in single-worker mode: confirm the
|
||||
# existing outcome before deleting recoverable goal state.
|
||||
if scheduled_goal is None and goal_completion is not None and record.status == RunStatus.success and not record.abort_event.is_set() and not record.ownership_lost:
|
||||
try:
|
||||
if await run_manager.persist_current_status(run_id):
|
||||
await _clear_completed_goal(
|
||||
candidate=goal_completion,
|
||||
record=record,
|
||||
run_manager=run_manager,
|
||||
bridge=bridge,
|
||||
accessor=accessor,
|
||||
checkpointer=checkpointer,
|
||||
)
|
||||
except Exception:
|
||||
logger.warning("Could not finalize satisfied goal for thread %s", thread_id, exc_info=True)
|
||||
|
||||
if started and not record.ownership_lost and checkpointer is not None and record.status == RunStatus.interrupted and not _is_edit_replay_run(record):
|
||||
try:
|
||||
await run_manager.wait_for_prior_finalizing(thread_id, run_id)
|
||||
if not await run_manager.has_later_started_run(thread_id, run_id):
|
||||
await _ensure_interrupted_title(checkpointer=checkpointer, thread_id=thread_id, app_config=ctx.app_config, graph_input=graph_input)
|
||||
except Exception:
|
||||
logger.debug("Failed to generate interrupted title for thread %s (non-fatal)", thread_id)
|
||||
|
||||
# Sync title from checkpoint to threads_meta.display_name
|
||||
if started and not record.ownership_lost and checkpointer is not None and thread_store is not None:
|
||||
try:
|
||||
ckpt_config = {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}}
|
||||
ckpt_tuple = await checkpointer.aget_tuple(ckpt_config)
|
||||
if ckpt_tuple is not None:
|
||||
ckpt = getattr(ckpt_tuple, "checkpoint", {}) or {}
|
||||
title = ckpt.get("channel_values", {}).get("title")
|
||||
if title:
|
||||
await thread_store.update_display_name(thread_id, title)
|
||||
except Exception:
|
||||
logger.debug("Failed to sync title for thread %s (non-fatal)", thread_id)
|
||||
|
||||
# Update threads_meta status based on run outcome
|
||||
if started and not record.ownership_lost and thread_store is not None:
|
||||
try:
|
||||
final_status = "idle" if record.status == RunStatus.success else record.status.value
|
||||
await thread_store.update_status(thread_id, final_status)
|
||||
except Exception:
|
||||
logger.debug("Failed to update thread_meta status for %s (non-fatal)", thread_id)
|
||||
|
||||
if started and not record.ownership_lost and accessor is not None and record.status == RunStatus.success and not record.abort_event.is_set() and record.human_input_request is None:
|
||||
try:
|
||||
from deerflow.agents.human_input import extract_unanswered_human_input_request
|
||||
|
||||
record.human_input_request = extract_unanswered_human_input_request(await _materialized_checkpoint_messages(accessor, thread_id))
|
||||
except Exception:
|
||||
logger.debug("Failed to inspect trailing human_input_request for run %s (non-fatal)", run_id, exc_info=True)
|
||||
|
||||
if not record.ownership_lost and not getattr(record, "terminal_status_persisted", True):
|
||||
try:
|
||||
await run_manager.persist_current_status(run_id)
|
||||
except Exception:
|
||||
logger.warning("Failed to confirm terminal status for run %s before completion hook", run_id, exc_info=True)
|
||||
|
||||
if not record.ownership_lost and getattr(record, "terminal_status_persisted", True) and ctx.on_run_completed is not None:
|
||||
try:
|
||||
await ctx.on_run_completed(record)
|
||||
except Exception:
|
||||
logger.warning("Run completion hook failed for %s (non-fatal)", run_id, exc_info=True)
|
||||
except BaseException as exc:
|
||||
# A terminal hook must not leave replacement runs blocked or
|
||||
# stream consumers waiting indefinitely.
|
||||
deferred_finalization_interrupt = _defer_finalization_interrupt(
|
||||
deferred_finalization_interrupt,
|
||||
exc,
|
||||
)
|
||||
logger.warning(
|
||||
"Run completion hook interrupted for %s; completing finalization first",
|
||||
run_id,
|
||||
)
|
||||
|
||||
if task_info is not None and task_store is not None:
|
||||
# Keep the finalizing barrier held until stop observers finish, so
|
||||
# a same-thread replacement cannot overlap this task's lifecycle.
|
||||
task_stop = asyncio.create_task(
|
||||
notify_task_stop(
|
||||
extensions,
|
||||
task_store,
|
||||
task_info,
|
||||
lead_task_outcome(
|
||||
aborted=(record.abort_event.is_set() or record.status == RunStatus.interrupted),
|
||||
succeeded=record.status == RunStatus.success,
|
||||
),
|
||||
timeout=_EXTENSION_TASK_NOTIFY_TIMEOUT_SECONDS,
|
||||
),
|
||||
name=f"extension-task-stop-{run_id}",
|
||||
)
|
||||
try:
|
||||
deferred_finalization_interrupt = await _await_task_stop_after_host_cancellation(
|
||||
task_stop,
|
||||
deferred_finalization_interrupt,
|
||||
)
|
||||
except Exception:
|
||||
logger.warning(
|
||||
"Extension task-stop notification failed for run %s (non-fatal)",
|
||||
run_id,
|
||||
exc_info=True,
|
||||
)
|
||||
except BaseException as exc:
|
||||
# Cancellation here must not strand the finalizing barrier or
|
||||
# leave stream consumers waiting for the end frame.
|
||||
deferred_finalization_interrupt = _defer_finalization_interrupt(
|
||||
deferred_finalization_interrupt,
|
||||
exc,
|
||||
)
|
||||
logger.warning(
|
||||
"Extension task-stop notification interrupted for run %s; completing cleanup first",
|
||||
run_id,
|
||||
)
|
||||
if record.finalizing:
|
||||
await run_manager.set_finalizing(run_id, False)
|
||||
|
||||
await bridge.publish_end(run_id)
|
||||
|
||||
if deferred_finalization_interrupt is not None:
|
||||
raise deferred_finalization_interrupt
|
||||
finally:
|
||||
try:
|
||||
if journal is not None:
|
||||
try:
|
||||
await journal.close(flush=not record.ownership_lost)
|
||||
except Exception:
|
||||
logger.warning("Failed to close journal for run %s", run_id, exc_info=True)
|
||||
logger.warning("Failed to record workspace changes for run %s", run_id, exc_info=True)
|
||||
|
||||
# Flush buffered journal events before the terminal receipt. The
|
||||
# receipt uses a run-scoped idempotent write shared with recovery, then
|
||||
# the staged terminal status is persisted. This ordering closes the
|
||||
# crash window where a terminal run could otherwise outlive its receipt.
|
||||
# A fenced worker leaves receipt recovery to the peer that claimed it.
|
||||
if not record.ownership_lost and journal is not None:
|
||||
try:
|
||||
await journal.flush()
|
||||
except Exception:
|
||||
logger.warning("Failed to flush journal for run %s", run_id, exc_info=True)
|
||||
|
||||
if delivery_content is None:
|
||||
if produced_output_paths is None:
|
||||
produced_output_paths = await _produced_output_paths(pre_run_workspace_snapshot, thread_id=thread_id, user_id=workspace_changes_user_id, extra_excluded_dir_names=workspace_excluded_dir_names)
|
||||
delivery_content = _delivery_content_with_outputs(journal.get_delivery_content(), produced_output_paths)
|
||||
receipt_persisted = await _persist_delivery_receipt(event_store, thread_id=thread_id, run_id=run_id, content=delivery_content)
|
||||
if produced_output_paths and record.status == RunStatus.success and not receipt_persisted:
|
||||
await run_manager.set_status(run_id, RunStatus.error, error=_DELIVERY_RECEIPT_FAILED_ERROR, persist=False)
|
||||
|
||||
if not record.ownership_lost and journal is not None and persist_completion:
|
||||
try:
|
||||
# Advance the final completion fields and timestamp without
|
||||
# terminalizing the durable row. That active row continues to
|
||||
# fence peer checkpoint writers through the duration write.
|
||||
completion_data = journal.get_completion_data()
|
||||
await run_manager.update_finalizing_progress(run_id, **completion_data)
|
||||
except Exception:
|
||||
logger.warning("Failed to persist finalizing run progress for %s (non-fatal)", run_id, exc_info=True)
|
||||
|
||||
# Keep the durable run row active through its final duration checkpoint
|
||||
# write. A peer Gateway admits history migration from the durable row,
|
||||
# not this worker's staged terminal status; terminalizing first would
|
||||
# let that migration read an unfinished lifetime and race this write.
|
||||
if started and not record.ownership_lost and checkpointer is not None and record.status == RunStatus.success:
|
||||
try:
|
||||
created = datetime.fromisoformat(record.created_at.replace("Z", "+00:00"))
|
||||
updated = datetime.fromisoformat(record.updated_at.replace("Z", "+00:00"))
|
||||
# Match legacy history semantics: turn_duration is the whole
|
||||
# RunRecord lifetime in integer seconds, including admission
|
||||
# delay. Persist zero for sub-second successful turns.
|
||||
duration = max(0, int((updated - created).total_seconds()))
|
||||
await _persist_run_duration(checkpointer=checkpointer, thread_id=thread_id, run_id=run_id, duration_seconds=duration)
|
||||
except Exception:
|
||||
logger.debug("Failed to persist run duration for thread %s run %s (non-fatal)", thread_id, run_id)
|
||||
|
||||
if scheduled_goal is not None and not record.ownership_lost and checkpointer is not None:
|
||||
if record.abort_event.is_set() and record.abort_action == "rollback" and not checkpoint_rollback_completed:
|
||||
await _finish_cancellation("rollback")
|
||||
# Scheduled goals are occurrence-scoped. Keep the real run slot
|
||||
# active through cleanup so a peer's next user turn cannot
|
||||
# inherit this goal. Ordinary goals retain terminal-first cleanup.
|
||||
cleanup_task = asyncio.create_task(
|
||||
_clear_scheduled_goal(goal=scheduled_goal, record=record, run_manager=run_manager, bridge=bridge, checkpointer=checkpointer, owned_admission=True),
|
||||
name=f"scheduled-goal-cleanup-{run_id}",
|
||||
)
|
||||
try:
|
||||
deferred_finalization_interrupt = await _await_task_stop_after_host_cancellation(cleanup_task, deferred_finalization_interrupt)
|
||||
except Exception:
|
||||
logger.warning("Could not remove scheduled goal for run %s", run_id, exc_info=True)
|
||||
await run_manager.set_status(run_id, RunStatus.error, error="Scheduled goal cleanup failed", persist=False)
|
||||
|
||||
if not record.ownership_lost and (event_store is not None or scheduled_goal is not None):
|
||||
try:
|
||||
# Even after bounded receipt retries are exhausted, persist the
|
||||
# real worker outcome. Leaving a successful row inflight would
|
||||
# let lease recovery rewrite it as an error with a synthetic
|
||||
# zero receipt.
|
||||
if record.abort_event.is_set():
|
||||
if scheduled_goal is not None:
|
||||
# A cancel may arrive while scheduled cleanup was
|
||||
# awaiting its saver. The local status may still be
|
||||
# staged success; apply the accepted action, including
|
||||
# rollback, before committing any terminal outcome.
|
||||
await _finish_cancellation(record.abort_action, restore_checkpoint=not checkpoint_rollback_completed)
|
||||
await run_manager.persist_current_status(run_id)
|
||||
else:
|
||||
cancel_action = await run_manager.set_status_if_not_cancelled(run_id, record.status, error=record.error, stop_reason=record.stop_reason, goal_verdict=record.goal_verdict)
|
||||
if cancel_action is not None:
|
||||
await _finish_cancellation(cancel_action)
|
||||
await run_manager.persist_current_status(run_id)
|
||||
except Exception:
|
||||
logger.warning("Failed to persist terminal status for run %s after delivery receipt attempts", run_id, exc_info=True)
|
||||
# The deferred commit has been attempted. A failed write is left to
|
||||
# completion fallback or lease recovery, as before.
|
||||
record.terminal_commit_pending = False
|
||||
|
||||
if not record.ownership_lost and journal is not None and persist_completion:
|
||||
try:
|
||||
# Persist token usage + convenience fields to RunStore
|
||||
completion_data = completion_data or journal.get_completion_data()
|
||||
await run_manager.update_run_completion(run_id, status=record.status.value, **completion_data)
|
||||
except Exception:
|
||||
logger.warning("Failed to persist run completion for %s (non-fatal)", run_id, exc_info=True)
|
||||
|
||||
if scheduled_goal is not None and not record.ownership_lost:
|
||||
status_confirmed = await run_manager.persist_current_status(run_id)
|
||||
if not status_confirmed:
|
||||
await run_manager._mark_ownership_lost(record, reason="Scheduled run terminal status could not be confirmed after goal cleanup.", require_active=False)
|
||||
# Late cancellation can still drain a rollback after goal
|
||||
# cleanup. Retain admission/heartbeat protection through that
|
||||
# rollback and the actual durable terminal acknowledgement.
|
||||
record.scheduled_goal_cleanup_pending = False
|
||||
|
||||
# A satisfied evaluator is only a candidate until artifact delivery,
|
||||
# receipt persistence and durable cancellation arbitration have ended.
|
||||
# Status writes are best-effort in single-worker mode: confirm the
|
||||
# existing outcome before deleting recoverable goal state.
|
||||
if scheduled_goal is None and goal_completion is not None and record.status == RunStatus.success and not record.abort_event.is_set() and not record.ownership_lost:
|
||||
try:
|
||||
status_confirmed = await run_manager.persist_current_status(run_id)
|
||||
if status_confirmed:
|
||||
await _clear_completed_goal(
|
||||
candidate=goal_completion,
|
||||
record=record,
|
||||
run_manager=run_manager,
|
||||
bridge=bridge,
|
||||
accessor=accessor,
|
||||
checkpointer=checkpointer,
|
||||
)
|
||||
except Exception:
|
||||
logger.warning("Could not finalize satisfied goal for thread %s", thread_id, exc_info=True)
|
||||
|
||||
if started and not record.ownership_lost and checkpointer is not None and record.status == RunStatus.interrupted and not _is_edit_replay_run(record):
|
||||
try:
|
||||
await run_manager.wait_for_prior_finalizing(thread_id, run_id)
|
||||
if not await run_manager.has_later_started_run(thread_id, run_id):
|
||||
await _ensure_interrupted_title(checkpointer=checkpointer, thread_id=thread_id, app_config=ctx.app_config, graph_input=graph_input)
|
||||
except Exception:
|
||||
logger.debug("Failed to generate interrupted title for thread %s (non-fatal)", thread_id)
|
||||
|
||||
# Sync title from checkpoint to threads_meta.display_name
|
||||
if started and not record.ownership_lost and checkpointer is not None and thread_store is not None:
|
||||
try:
|
||||
ckpt_config = {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}}
|
||||
ckpt_tuple = await checkpointer.aget_tuple(ckpt_config)
|
||||
if ckpt_tuple is not None:
|
||||
ckpt = getattr(ckpt_tuple, "checkpoint", {}) or {}
|
||||
title = ckpt.get("channel_values", {}).get("title")
|
||||
if title:
|
||||
await thread_store.update_display_name(thread_id, title)
|
||||
except Exception:
|
||||
logger.debug("Failed to sync title for thread %s (non-fatal)", thread_id)
|
||||
|
||||
# Update threads_meta status based on run outcome
|
||||
if started and not record.ownership_lost and thread_store is not None:
|
||||
try:
|
||||
final_status = "idle" if record.status == RunStatus.success else record.status.value
|
||||
await thread_store.update_status(thread_id, final_status)
|
||||
except Exception:
|
||||
logger.debug("Failed to update thread_meta status for %s (non-fatal)", thread_id)
|
||||
|
||||
if started and not record.ownership_lost and accessor is not None and record.status == RunStatus.success and not record.abort_event.is_set() and record.human_input_request is None:
|
||||
try:
|
||||
from deerflow.agents.human_input import extract_unanswered_human_input_request
|
||||
|
||||
record.human_input_request = extract_unanswered_human_input_request(await _materialized_checkpoint_messages(accessor, thread_id))
|
||||
except Exception:
|
||||
logger.debug("Failed to inspect trailing human_input_request for run %s (non-fatal)", run_id, exc_info=True)
|
||||
|
||||
if not record.ownership_lost and not getattr(record, "terminal_status_persisted", True):
|
||||
try:
|
||||
await run_manager.persist_current_status(run_id)
|
||||
except Exception:
|
||||
logger.warning("Failed to confirm terminal status for run %s before completion hook", run_id, exc_info=True)
|
||||
|
||||
if not record.ownership_lost and getattr(record, "terminal_status_persisted", True) and ctx.on_run_completed is not None:
|
||||
try:
|
||||
# Preserve hook liveness without forwarding cancellation
|
||||
# to the finalizer's durable writes or resource cleanup.
|
||||
async def invoke_completion_hook() -> None:
|
||||
await ctx.on_run_completed(record)
|
||||
|
||||
completion_hook_task = asyncio.create_task(invoke_completion_hook(), name=f"run-completion-hook-{run_id}")
|
||||
if host_interrupted:
|
||||
await asyncio.sleep(0)
|
||||
completion_hook_task.cancel()
|
||||
await completion_hook_task
|
||||
except Exception:
|
||||
logger.warning("Run completion hook failed for %s (non-fatal)", run_id, exc_info=True)
|
||||
except BaseException as exc:
|
||||
# A terminal hook must not leave replacement runs blocked or
|
||||
# stream consumers waiting indefinitely.
|
||||
deferred_finalization_interrupt = _defer_finalization_interrupt(
|
||||
deferred_finalization_interrupt,
|
||||
exc,
|
||||
)
|
||||
logger.warning(
|
||||
"Run completion hook interrupted for %s; completing finalization first",
|
||||
run_id,
|
||||
)
|
||||
finally:
|
||||
completion_hook_task = None
|
||||
|
||||
if task_info is not None and task_store is not None:
|
||||
# Keep the finalizing barrier held until stop observers finish, so
|
||||
# a same-thread replacement cannot overlap this task's lifecycle.
|
||||
task_stop = asyncio.create_task(
|
||||
notify_task_stop(
|
||||
extensions,
|
||||
task_store,
|
||||
task_info,
|
||||
lead_task_outcome(
|
||||
aborted=(record.abort_event.is_set() or record.status == RunStatus.interrupted),
|
||||
succeeded=record.status == RunStatus.success,
|
||||
),
|
||||
timeout=_EXTENSION_TASK_NOTIFY_TIMEOUT_SECONDS,
|
||||
),
|
||||
name=f"extension-task-stop-{run_id}",
|
||||
)
|
||||
try:
|
||||
deferred_finalization_interrupt = await _await_task_stop_after_host_cancellation(
|
||||
task_stop,
|
||||
deferred_finalization_interrupt,
|
||||
)
|
||||
except Exception:
|
||||
logger.warning(
|
||||
"Extension task-stop notification failed for run %s (non-fatal)",
|
||||
run_id,
|
||||
exc_info=True,
|
||||
)
|
||||
except BaseException as exc:
|
||||
# Cancellation here must not strand the finalizing barrier or
|
||||
# leave stream consumers waiting for the end frame.
|
||||
deferred_finalization_interrupt = _defer_finalization_interrupt(
|
||||
deferred_finalization_interrupt,
|
||||
exc,
|
||||
)
|
||||
logger.warning(
|
||||
"Extension task-stop notification interrupted for run %s; completing cleanup first",
|
||||
run_id,
|
||||
)
|
||||
if record.finalizing:
|
||||
await run_manager.set_finalizing(run_id, False)
|
||||
|
||||
await bridge.publish_end(run_id)
|
||||
|
||||
if deferred_finalization_interrupt is not None:
|
||||
raise deferred_finalization_interrupt
|
||||
finally:
|
||||
lease_cleanup_interrupt: BaseException | None = None
|
||||
try:
|
||||
from deerflow.sandbox.lease import release_sandbox_execution_lease_async
|
||||
|
||||
await release_sandbox_execution_lease_async(runtime_ctx)
|
||||
except Exception:
|
||||
logger.warning("Failed to release sandbox execution lease for run %s", run_id, exc_info=True)
|
||||
except BaseException as exc:
|
||||
# release_async completes the underlying cleanup before it
|
||||
# re-raises cancellation. Defer that interruption until the
|
||||
# worker has dropped all other run-scoped references too.
|
||||
lease_cleanup_interrupt = exc
|
||||
logger.warning(
|
||||
"Sandbox execution lease cleanup was interrupted for run %s; completing local cleanup first",
|
||||
run_id,
|
||||
)
|
||||
if journal is not None:
|
||||
try:
|
||||
await journal.close(flush=not record.ownership_lost)
|
||||
except Exception:
|
||||
logger.warning("Failed to close journal for run %s", run_id, exc_info=True)
|
||||
finally:
|
||||
_release_run_scoped_references(
|
||||
runnable_configs,
|
||||
runtime_ctx,
|
||||
journal,
|
||||
)
|
||||
# Drop graph and per-run payload references before the terminal
|
||||
# worker task itself becomes collectable.
|
||||
agent = None
|
||||
agent_result = None
|
||||
accessor = None
|
||||
runtime = None
|
||||
runtime_ctx = None
|
||||
rollback_point = None
|
||||
subagent_events = None
|
||||
goal_evaluator_model = None
|
||||
goal_completion = None
|
||||
task_store = None
|
||||
task_info = None
|
||||
pre_run_workspace_snapshot = None
|
||||
delivery_content = None
|
||||
produced_output_paths = None
|
||||
graph_input = {}
|
||||
lease_cleanup_interrupt: BaseException | None = None
|
||||
try:
|
||||
from deerflow.sandbox.lease import release_sandbox_execution_lease_async
|
||||
|
||||
# Durable finalization and terminal publication may depend on
|
||||
# external backends, but local housekeeping must always run.
|
||||
record.scheduled_goal_cleanup_pending = False
|
||||
_create_contextless_task(bridge.cleanup(run_id, delay=60))
|
||||
# Preserve the existing five-minute grace period for local
|
||||
# join/status paths, then release the terminal record, completed
|
||||
# task, and request payload. Durable run history remains available
|
||||
# through RunStore.
|
||||
_create_contextless_task(run_manager.cleanup(run_id))
|
||||
_schedule_terminal_cycle_collection()
|
||||
await release_sandbox_execution_lease_async(runtime_ctx)
|
||||
except Exception:
|
||||
logger.warning("Failed to release sandbox execution lease for run %s", run_id, exc_info=True)
|
||||
except BaseException as exc:
|
||||
# release_async completes the underlying cleanup before it
|
||||
# re-raises cancellation. Defer that interruption until the
|
||||
# worker has dropped all other run-scoped references too.
|
||||
lease_cleanup_interrupt = exc
|
||||
logger.warning(
|
||||
"Sandbox execution lease cleanup was interrupted for run %s; completing local cleanup first",
|
||||
run_id,
|
||||
)
|
||||
finally:
|
||||
_release_run_scoped_references(
|
||||
runnable_configs,
|
||||
runtime_ctx,
|
||||
journal,
|
||||
)
|
||||
# Drop graph and per-run payload references before the terminal
|
||||
# worker task itself becomes collectable.
|
||||
agent = None
|
||||
agent_result = None
|
||||
accessor = None
|
||||
runtime = None
|
||||
runtime_ctx = None
|
||||
rollback_point = None
|
||||
subagent_events = None
|
||||
goal_evaluator_model = None
|
||||
goal_completion = None
|
||||
task_store = None
|
||||
task_info = None
|
||||
pre_run_workspace_snapshot = None
|
||||
delivery_content = None
|
||||
produced_output_paths = None
|
||||
graph_input = {}
|
||||
|
||||
if lease_cleanup_interrupt is not None:
|
||||
raise lease_cleanup_interrupt
|
||||
# Durable finalization and terminal publication may depend on
|
||||
# external backends, but local housekeeping must always run.
|
||||
record.scheduled_goal_cleanup_pending = False
|
||||
_create_contextless_task(bridge.cleanup(run_id, delay=60))
|
||||
# Preserve the existing five-minute grace period for local
|
||||
# join/status paths, then release the terminal record, completed
|
||||
# task, and request payload. Durable run history remains available
|
||||
# through RunStore.
|
||||
_create_contextless_task(run_manager.cleanup(run_id))
|
||||
_schedule_terminal_cycle_collection()
|
||||
|
||||
if lease_cleanup_interrupt is not None:
|
||||
raise lease_cleanup_interrupt
|
||||
|
||||
# The registered worker remains the owner until every terminal and
|
||||
# resource-cleanup await settles. Never cancel or detach this task.
|
||||
finalizer = asyncio.create_task(finalize(), name=f"run-finalization-{run_id}")
|
||||
await _drain_run_finalization(finalizer, on_host_cancel=interrupt_completion_hook)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
332
backend/tests/test_run_worker_finalization_cancellation.py
Normal file
332
backend/tests/test_run_worker_finalization_cancellation.py
Normal file
@ -0,0 +1,332 @@
|
||||
"""Host cancellation must preserve the terminal lifecycle and owned cleanup."""
|
||||
|
||||
import asyncio
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
|
||||
import deerflow.runtime.runs.worker as worker_module
|
||||
from deerflow.runtime.events.store.memory import MemoryRunEventStore
|
||||
from deerflow.runtime.journal import RunJournal
|
||||
from deerflow.runtime.runs.manager import ConflictError, RunManager
|
||||
from deerflow.runtime.runs.store.memory import MemoryRunStore
|
||||
from deerflow.runtime.runs.worker import RunContext, run_agent
|
||||
from deerflow.runtime.stream_bridge import MemoryStreamBridge
|
||||
|
||||
|
||||
class DummyAgent:
|
||||
async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
|
||||
yield {"type": "task_started", "task_id": "verification-subtask", "description": "offline verification"}
|
||||
|
||||
|
||||
async def start_run(monkeypatch, *, event_store, thread_store=None, run_store=None, bridge=None):
|
||||
run_store = run_store or MemoryRunStore()
|
||||
manager = RunManager(store=run_store)
|
||||
record = await manager.create_or_reject("verify-6914")
|
||||
bridge = bridge or MemoryStreamBridge()
|
||||
completed = AsyncMock()
|
||||
cleanup_tasks = []
|
||||
original_create = worker_module._create_contextless_task
|
||||
|
||||
def track_cleanup(coro):
|
||||
task = original_create(coro)
|
||||
cleanup_tasks.append(task)
|
||||
return task
|
||||
|
||||
monkeypatch.setattr(worker_module, "_create_contextless_task", track_cleanup)
|
||||
worker = asyncio.create_task(
|
||||
run_agent(
|
||||
bridge,
|
||||
manager,
|
||||
record,
|
||||
ctx=RunContext(checkpointer=None, event_store=event_store, thread_store=thread_store, on_run_completed=completed),
|
||||
agent_factory=lambda *, config: DummyAgent(),
|
||||
graph_input={},
|
||||
config={},
|
||||
stream_modes=["custom"],
|
||||
)
|
||||
)
|
||||
record.task = worker
|
||||
return SimpleNamespace(worker=worker, manager=manager, record=record, bridge=bridge, completed=completed, run_store=run_store, event_store=event_store, cleanup_tasks=cleanup_tasks)
|
||||
|
||||
|
||||
async def snapshot(run):
|
||||
record = run.record
|
||||
stored = await run.run_store.get(record.run_id)
|
||||
events = await run.event_store.list_events(record.thread_id, record.run_id)
|
||||
try:
|
||||
await run.manager.create_or_reject(record.thread_id)
|
||||
replacement_admitted = True
|
||||
except ConflictError:
|
||||
replacement_admitted = False
|
||||
return {
|
||||
"cancelled": run.worker.cancelled(),
|
||||
"local_status": record.status.value,
|
||||
"stored_status": stored["status"],
|
||||
"terminal_commit_pending": record.terminal_commit_pending,
|
||||
"finalizing": record.finalizing,
|
||||
"receipt_count": sum(e["event_type"] == "run.delivery" for e in events),
|
||||
"subagent_count": sum(e["event_type"] == "subagent.start" for e in events),
|
||||
"end": run.bridge._streams[record.run_id].ended,
|
||||
"completion_hooks": run.completed.await_count,
|
||||
"replacement_admitted": replacement_admitted,
|
||||
}
|
||||
|
||||
|
||||
async def reap(run):
|
||||
if not run.worker.done():
|
||||
run.worker.cancel()
|
||||
await asyncio.gather(run.worker, return_exceptions=True)
|
||||
for task in run.cleanup_tasks:
|
||||
task.cancel()
|
||||
await asyncio.gather(*run.cleanup_tasks, return_exceptions=True)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("stage", ["receipt", "subagent_flush", "thread_status", "terminal_status", "journal_close", "end"])
|
||||
async def test_cancellation_preserves_terminal_lifecycle(monkeypatch, stage):
|
||||
entered, release, closed = asyncio.Event(), asyncio.Event(), asyncio.Event()
|
||||
|
||||
async def gate():
|
||||
entered.set()
|
||||
await release.wait()
|
||||
|
||||
class GatedEventStore(MemoryRunEventStore):
|
||||
async def put_if_absent(self, **kwargs):
|
||||
if stage == "receipt" and kwargs.get("event_type") == "run.delivery":
|
||||
await gate()
|
||||
return await super().put_if_absent(**kwargs)
|
||||
|
||||
async def put_batch(self, items):
|
||||
if stage == "subagent_flush" and any(i.get("event_type") == "subagent.start" for i in items):
|
||||
await gate()
|
||||
return await super().put_batch(items)
|
||||
|
||||
async def update_status(thread_id, status):
|
||||
if stage == "thread_status" and status != "running":
|
||||
await gate()
|
||||
|
||||
class GatedRunStore(MemoryRunStore):
|
||||
async def update_status(self, run_id, status, **kwargs):
|
||||
if stage == "terminal_status" and status not in ("pending", "running"):
|
||||
await gate()
|
||||
return await super().update_status(run_id, status, **kwargs)
|
||||
|
||||
class GatedBridge(MemoryStreamBridge):
|
||||
async def publish_end(self, run_id):
|
||||
if stage == "end":
|
||||
await gate()
|
||||
await super().publish_end(run_id)
|
||||
|
||||
original_close = RunJournal.close
|
||||
|
||||
async def gated_close(journal, **kwargs):
|
||||
if stage == "journal_close":
|
||||
await gate()
|
||||
await original_close(journal, **kwargs)
|
||||
closed.set()
|
||||
|
||||
monkeypatch.setattr(RunJournal, "close", gated_close)
|
||||
|
||||
run = await start_run(monkeypatch, event_store=GatedEventStore(), thread_store=SimpleNamespace(update_status=update_status), run_store=GatedRunStore(), bridge=GatedBridge())
|
||||
try:
|
||||
await asyncio.wait_for(entered.wait(), 5)
|
||||
run.worker.cancel("verification cancellation")
|
||||
await asyncio.sleep(0)
|
||||
for _ in range(3):
|
||||
run.worker.cancel("repeated finalization cancellation")
|
||||
await asyncio.sleep(0)
|
||||
assert not run.worker.done()
|
||||
assert not run.cleanup_tasks
|
||||
release.set()
|
||||
done, _ = await asyncio.wait({run.worker}, timeout=5)
|
||||
assert done, "worker failed to settle after the operation was released"
|
||||
with pytest.raises(asyncio.CancelledError, match="verification cancellation"):
|
||||
await run.worker
|
||||
assert closed.is_set(), "worker exited before journal cleanup finished"
|
||||
state = await snapshot(run)
|
||||
assert state["cancelled"], state
|
||||
assert state["end"], state
|
||||
assert state["receipt_count"] == 1, state
|
||||
assert state["subagent_count"] == 1, state
|
||||
assert state["stored_status"] == state["local_status"], state
|
||||
assert not state["terminal_commit_pending"], state
|
||||
assert not state["finalizing"], state
|
||||
assert state["completion_hooks"] == 1, state
|
||||
assert state["replacement_admitted"], state
|
||||
finally:
|
||||
release.set()
|
||||
await reap(run)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel_is_preserved_when_guarded_flush_fails(monkeypatch):
|
||||
entered, release = asyncio.Event(), asyncio.Event()
|
||||
original_flush = RunJournal.flush
|
||||
first_flush = True
|
||||
|
||||
async def failing_flush(journal):
|
||||
nonlocal first_flush
|
||||
if first_flush:
|
||||
first_flush = False
|
||||
entered.set()
|
||||
await release.wait()
|
||||
raise OSError("injected finalization flush failure")
|
||||
return await original_flush(journal)
|
||||
|
||||
monkeypatch.setattr(RunJournal, "flush", failing_flush)
|
||||
run = await start_run(monkeypatch, event_store=MemoryRunEventStore())
|
||||
try:
|
||||
await asyncio.wait_for(entered.wait(), 5)
|
||||
run.worker.cancel("original host cancellation")
|
||||
await asyncio.sleep(0)
|
||||
release.set()
|
||||
done, _ = await asyncio.wait({run.worker}, timeout=5)
|
||||
assert done
|
||||
state = await snapshot(run)
|
||||
assert state["cancelled"], state
|
||||
assert state["end"], state
|
||||
finally:
|
||||
release.set()
|
||||
await reap(run)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failed_finalizer_preserves_host_cancellation_as_primary():
|
||||
entered, release = asyncio.Event(), asyncio.Event()
|
||||
failure = OSError("finalizer failed during cancellation drain")
|
||||
|
||||
async def finalize():
|
||||
entered.set()
|
||||
await release.wait()
|
||||
raise failure
|
||||
|
||||
finalizer = asyncio.create_task(finalize())
|
||||
host = asyncio.create_task(worker_module._drain_run_finalization(finalizer))
|
||||
try:
|
||||
await asyncio.wait_for(entered.wait(), 5)
|
||||
host.cancel("first host cancellation")
|
||||
await asyncio.sleep(0)
|
||||
host.cancel("second host cancellation")
|
||||
await asyncio.sleep(0)
|
||||
assert not finalizer.done()
|
||||
release.set()
|
||||
with pytest.raises(asyncio.CancelledError, match="first host cancellation") as caught:
|
||||
await host
|
||||
assert caught.value.__cause__ is failure
|
||||
assert finalizer.done()
|
||||
assert finalizer.exception() is failure
|
||||
finally:
|
||||
release.set()
|
||||
await asyncio.gather(host, finalizer, return_exceptions=True)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_finalizer_cancellation_is_not_mistaken_for_host_cancellation():
|
||||
async def finalize():
|
||||
raise asyncio.CancelledError("finalizer interrupted")
|
||||
|
||||
finalizer = asyncio.create_task(finalize())
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await worker_module._drain_run_finalization(finalizer)
|
||||
assert asyncio.current_task().cancelling() == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("stage", ["thread_status", "journal_close", "end"])
|
||||
async def test_shutdown_drains_terminal_worker_after_durable_commit(monkeypatch, stage):
|
||||
entered, release, shutdown_entered = asyncio.Event(), asyncio.Event(), asyncio.Event()
|
||||
|
||||
async def gate():
|
||||
entered.set()
|
||||
await release.wait()
|
||||
|
||||
async def update_status(thread_id, status):
|
||||
if stage == "thread_status" and status != "running":
|
||||
await gate()
|
||||
|
||||
original_close = RunJournal.close
|
||||
|
||||
async def close(journal, **kwargs):
|
||||
if stage == "journal_close":
|
||||
await gate()
|
||||
await original_close(journal, **kwargs)
|
||||
|
||||
class GatedBridge(MemoryStreamBridge):
|
||||
async def publish_end(self, run_id):
|
||||
if stage == "end":
|
||||
await gate()
|
||||
await super().publish_end(run_id)
|
||||
|
||||
monkeypatch.setattr(RunJournal, "close", close)
|
||||
run = await start_run(monkeypatch, event_store=MemoryRunEventStore(), thread_store=SimpleNamespace(update_status=update_status), bridge=GatedBridge())
|
||||
shutdown = None
|
||||
try:
|
||||
await asyncio.wait_for(entered.wait(), 5)
|
||||
assert (await run.run_store.get(run.record.run_id))["status"] == "success"
|
||||
assert not run.record.terminal_commit_pending
|
||||
original_stop = run.manager.stop_heartbeat
|
||||
|
||||
async def stop_heartbeat(**kwargs):
|
||||
await original_stop(**kwargs)
|
||||
shutdown_entered.set()
|
||||
|
||||
monkeypatch.setattr(run.manager, "stop_heartbeat", stop_heartbeat)
|
||||
# A short drain budget is allowed to expire, but must not cancel or
|
||||
# retire the still-owned finalizer. A subsequent drain can finish it.
|
||||
await run.manager.shutdown(timeout=0.01)
|
||||
assert not run.worker.done()
|
||||
shutdown_entered.clear()
|
||||
shutdown = asyncio.create_task(run.manager.shutdown(timeout=5))
|
||||
await asyncio.wait_for(shutdown_entered.wait(), 5)
|
||||
await asyncio.sleep(0)
|
||||
assert not shutdown.done(), "shutdown skipped a committed run still finalizing"
|
||||
assert not run.cleanup_tasks
|
||||
release.set()
|
||||
await asyncio.wait_for(shutdown, 5)
|
||||
assert run.worker.done()
|
||||
assert not run.worker.cancelled()
|
||||
state = await snapshot(run)
|
||||
assert state["end"] and state["completion_hooks"] == 1, state
|
||||
assert state["stored_status"] == "success", state
|
||||
finally:
|
||||
release.set()
|
||||
if shutdown is not None:
|
||||
await asyncio.gather(shutdown, return_exceptions=True)
|
||||
await reap(run)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_earlier_cancellation_does_not_leave_completion_hook_blocked(monkeypatch):
|
||||
receipt_entered, release_receipt, hook_stopped = asyncio.Event(), asyncio.Event(), asyncio.Event()
|
||||
|
||||
class GatedEventStore(MemoryRunEventStore):
|
||||
async def put_if_absent(self, **kwargs):
|
||||
if kwargs.get("event_type") == "run.delivery":
|
||||
receipt_entered.set()
|
||||
await release_receipt.wait()
|
||||
return await super().put_if_absent(**kwargs)
|
||||
|
||||
async def completion_hook(record):
|
||||
try:
|
||||
await asyncio.Event().wait()
|
||||
finally:
|
||||
hook_stopped.set()
|
||||
|
||||
run = await start_run(monkeypatch, event_store=GatedEventStore())
|
||||
run.completed.side_effect = completion_hook
|
||||
try:
|
||||
await asyncio.wait_for(receipt_entered.wait(), 5)
|
||||
run.worker.cancel("cancellation before the completion hook")
|
||||
await asyncio.sleep(0)
|
||||
release_receipt.set()
|
||||
done, _ = await asyncio.wait({run.worker}, timeout=5)
|
||||
assert done, "an earlier deferred interruption left the completion hook blocked"
|
||||
with pytest.raises(asyncio.CancelledError, match="cancellation before the completion hook"):
|
||||
await run.worker
|
||||
assert hook_stopped.is_set()
|
||||
assert run.bridge._streams[run.record.run_id].ended
|
||||
finally:
|
||||
release_receipt.set()
|
||||
await reap(run)
|
||||
Loading…
x
Reference in New Issue
Block a user