fix(runtime): cancel runs across live gateway workers (#4500)

* docs(runtime): design cross-worker cancellation

* fix(runtime): cancel runs across gateway workers

* fix(runtime): harden cross-worker cancellation races

让取消请求与 owner 终态写入通过持久化 CAS 决定先后,保证首次取消 action 在不同 worker 路由下保持一致。\n\n将 heartbeat 收敛为续租后仅发送本地中止信号,并补齐完成竞态与路由重试的回归用例。

* docs(runtime): drop implementation plan from PR

移除仅用于实现过程的跨 worker 取消设计记录,保留 README 和 backend/AGENTS.md 中面向最终行为的文档。

* fix(runtime): preserve local cancel fallback

* test(runtime): adapt worker run manager fakes

* docs(runtime): fix run cancel migration registry

---------

Co-authored-by: MiaoRuidx <12540796+MiaoRuidx@users.noreply.github.com>
This commit is contained in:
MiaoRuidx 2026-07-28 21:45:20 +08:00 committed by GitHub
parent b1984cf4ab
commit 8a78c264b7
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
24 changed files with 1107 additions and 158 deletions

View File

@ -286,6 +286,8 @@ DeerFlow still uses `Forwarded` / `X-Forwarded-*` headers to recover the browser
> [!IMPORTANT]
> The Gateway still owns active run tasks in process, so production defaults to a single Gateway worker (`GATEWAY_WORKERS=1`). Multi-worker deployments require Postgres, the Redis stream bridge (`stream_bridge.type: redis`), `run_ownership.heartbeat_enabled: true`, and `run_events.backend: db`; process-local memory/JSONL event stores cannot enforce singleton delivery receipts across workers. The bridge shares SSE delivery and bounded `Last-Event-ID` replay across workers. When a valid reconnect cursor has been trimmed, or a subscriber that already established an empty-stream wait falls behind before its first delivery, Memory and Redis emit a machine-readable SSE `gap` event instead of silently returning a partial replay; the Web UI reloads durable thread/event state and resumes from the retained tail. Lease reconciliation marks runs from dead workers as errors, persists their delivery receipts, publishes the terminal stream marker, schedules retained-stream cleanup, and updates the affected thread status. SSE and `/wait` consumers also refresh durable status on heartbeats as a fallback if terminal publication fails. Malformed Redis reconnect IDs live-tail new events instead of replaying the retained buffer, and the rolling retained-buffer TTL (`stream_ttl_seconds`) remains a cleanup safety net rather than a run timeout. IM channel state and other process-local services still need their own multi-worker coordination.
>
> Run cancellation may land on any Gateway worker. A non-owning worker now persists the interrupt or rollback request for the live owner, which observes it during lease renewal and performs the normal cancellation flow; load-balancer routing alone no longer produces a 409. The first accepted action wins even if a retry lands on the owner, and accepted cancellation competes atomically with owner completion. Dead owners still follow lease takeover and orphan recovery. Cancellation latency is therefore bounded by the lease heartbeat interval.
>
> With lease heartbeat enabled, a transient RunStore renewal error is retried only until the last confirmed lease expires; the stale worker then cancels local execution and suppresses checkpoint, completion-hook, delivery-receipt, and thread-status finalization. A remote tool side effect already in flight may still be outside local cancellation.
>
> Reconciliation uses an atomic takeover claim that re-checks the lease after candidate selection, so a successful owner renewal wins over orphan recovery and only one reconciler can report a run as recovered. When multiple Gateway workers share the Docker/AIO or E2B sandbox backend, also configure `sandbox.ownership.type: redis`; E2B uses the leases during background startup and periodic reconciliation so duplicate/orphan cleanup cannot terminate a live peer's sandbox.

View File

@ -508,10 +508,10 @@ JSONL event stores when `GATEWAY_WORKERS > 1`.
- Edit-and-rerun visibility is derived from edit replay runs (`metadata.replay_kind="edit"` plus `regenerate_from_run_id`) by `RunManager.list_edit_replay_visibility()`: the newest attempt for each source run is authoritative. Pending/running/success attempts hide the original source run; failed, timed-out, or interrupted attempts hide only the failed attempt so the original conversation reappears.
- When a persistent `RunStore` is configured, `get()` and `list_by_thread()` hydrate historical runs from the store. In-memory records win for the same `run_id` so task, abort, and stream-control state stays attached to active local runs.
- Thread metadata status switches to `running` only after `RunManager.try_start()` succeeds. Pending-cancelled runs therefore skip the old `running` projection, while clients may observe the prior thread status during the short worker-startup window.
- `cancel()` returns a :class:`~deerflow.runtime.CancelOutcome` enum: `cancelled` (local cancel), `taken_over` (non-owning worker claimed the run because the owner's lease expired — marks it as `error`), `lease_valid_elsewhere` (owner's lease is still alive — caller should return 409 + `Retry-After`), `not_active_locally` (heartbeat disabled, preserving the old 409 path), `not_cancellable` (terminal state), or `unknown` (not found in memory or store). `create_or_reject(..., multitask_strategy="interrupt"|"rollback")` persists interrupted status through `RunStore.update_status()`, matching normal `set_status()` transitions.
- `cancel()` returns a :class:`~deerflow.runtime.CancelOutcome` enum: `cancelled` (local cancel), `requested` (the non-owning worker durably recorded the first cancellation action for the live owner), `taken_over` (non-owning worker claimed the run because the owner's lease expired — marks it as `error`), `lease_valid_elsewhere` (legacy/custom store lacks the durable request primitive — caller retains the safe 409 + `Retry-After` fallback), `not_active_locally` (heartbeat disabled, preserving the old 409 path), `not_cancellable` (terminal state), or `unknown` (not found in memory or store). `create_or_reject(..., multitask_strategy="interrupt"|"rollback")` persists interrupted status through `RunStore.update_status()`, matching normal `set_status()` transitions.
- Interrupt/rollback admission registers the replacement before its best-effort persistence of locally interrupted predecessors. If the admitting caller is cancelled during that post-registration await, `RunManager` drains a shielded replacement cleanup before propagating `CancelledError`, including across repeated cancellation. The cleanup normally persists `interrupted`; if that best-effort transition fails, it retries the active-to-`interrupted` store transition strictly and verifies the result with the replacement's captured owner identity. A concurrent peer terminal transition wins and is synchronized back into the local record rather than being overwritten or deleted.
- Store-only hydrated runs are readable history. In multi-worker mode with heartbeat enabled, cancel on a store-only run can take over (mark `error`) when the owner's lease has expired past the grace window; otherwise it fails with 409 + `Retry-After`. In single-worker mode (heartbeat off), store-only runs still return 409.
- A local worker's `RunRecord.lease_expires_at` is the last durably confirmed ownership deadline. `_renew_leases()` bounds each renewal attempt by that deadline: transient store exceptions remain retryable while it is valid, but an exception or blocked call that reaches expiry sets the process-local `ownership_lost` fence, raises `abort_event`, and cancels the run task. Fenced workers do not perform subsequent journal/delivery-receipt, progress/completion/status, checkpoint/thread-metadata, or `on_run_completed` writes; the peer recovery path owns the terminal receipt. `RunStore.update_run_completion()` also refuses to replace a different terminal status, closing the peer-takeover/late-finalization race. `grace_seconds` delays peer reclamation for clock skew but is not extra execution time for an owner that can no longer confirm its lease. Already-committed remote tool side effects remain outside this local cancellation boundary.
- Store-only hydrated runs are readable history. In multi-worker mode with heartbeat enabled, cancel on a store-only run records `runs.cancel_action` / `cancel_requested_at` while the owner's lease is live; the first action wins even if a retry later lands on the owner. `RunStore.request_cancel()` and owner completion through `finalize_if_not_cancelled()` are competing active-row CAS operations, so an accepted cancel cannot be overwritten by a later success. `RunStore.renew_lease()` renews and observes the request atomically in the SQL implementation. The owner then executes the normal process-local interrupt/rollback and terminal stream path without transferring the lease. An expired owner is still taken over and marked `error`. `wait=true` and cancel-then-stream use the shared bridge to observe owner finalization; a non-standard process-local bridge returns accepted 202 instead of subscribing to an unreachable stream. In single-worker mode (heartbeat off), store-only runs still return 409.
- A local worker's `RunRecord.lease_expires_at` is the last durably confirmed ownership deadline. `_renew_leases()` bounds each renewal attempt by that deadline: transient store exceptions remain retryable while it is valid, but an exception or blocked call that reaches expiry sets the process-local `ownership_lost` fence, raises `abort_event`, and cancels the run task. Successful renewals collect durable cancellation actions; after all local renewals have been attempted, heartbeat only signals the corresponding process-local tasks, leaving status writes and rollback cleanup to the worker finalization path. Fenced workers do not perform subsequent journal/delivery-receipt, progress/completion/status, checkpoint/thread-metadata, or `on_run_completed` writes; the peer recovery path owns the terminal receipt. `RunStore.update_run_completion()` also refuses to replace a different terminal status, closing the peer-takeover/late-finalization race. `grace_seconds` delays peer reclamation for clock skew but is not extra execution time for an owner that can no longer confirm its lease. Already-committed remote tool side effects remain outside this local cancellation boundary.
- Startup/orphan reconciliation must claim stale active rows with `RunStore.claim_for_takeover()`, not a plain `update_status()`. The final claim re-checks `status` and lease expiry atomically, so a heartbeat renewal between the candidate scan and the recovery write keeps the run active.
- Run admission and independent checkpoint writes are first-class thread operations. `runs.operation_kind` distinguishes user-visible `run` rows from internal `checkpoint_write` reservations, while every active kind shares the existing durable active-thread uniqueness constraint. New operation kinds must go through `RunStore.create_thread_operation_atomic()` and `RunManager.reserve_thread_operation()` rather than adding another lock or metadata marker. Live and lease-less reservations are non-interruptible; an expired leased reservation can be reclaimed immediately by interrupt/rollback admission without waiting for orphan reconciliation. Lease-less rows stay fail-closed because the store cannot distinguish a stale row from a live checkpoint writer in another heartbeat-disabled worker; a rare failed delete therefore requires startup reconciliation, and heartbeat-disabled multi-worker deployment remains unsupported. Reservation bodies are attached to their caller task so loss detected by lease renewal cancels the checkpoint writer before it can continue after takeover; the context manager translates that lease-loss cancellation to `ConflictError` after cleanup so Gateway mutation routes return a retryable 409 instead of dropping the HTTP request. The cleanup scope begins immediately after durable admission, including the await that attaches the caller task, so cancellation cannot strand a locally renewed pending reservation. A failed renewal is revalidated under the manager lock before cancellation; if the reservation completed and unregistered while the store update was in flight, its request task must not be cancelled after the checkpoint write. Reservations are excluded from run history/reporting and from run-only helpers such as `list_by_thread()` and `has_inflight()`, release uses the captured owner rather than ambient user context, and local cleanup still runs when the best-effort store delete fails. `RunStore.create_run_atomic()` remains a deprecated compatibility shim for external stores that only admit normal runs; new stores must implement `create_thread_operation_atomic()` to support internal operation kinds.
- 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.
@ -921,6 +921,7 @@ This invokes `alembic revision --autogenerate` against the live ORM models. Revi
- `migrations/versions/0004_run_ownership.py``runs` multi-worker ownership + the `uq_runs_thread_active` partial unique index, with a `_dedupe_active_runs_per_thread()` pre-step so `CREATE UNIQUE INDEX` cannot fail on a field DB that already has duplicate active rows per thread
- `migrations/versions/0007_scheduled_run_active_index.py` — the `uq_scheduled_task_run_active` partial unique index (at most one queued/running `scheduled_task_runs` row per `task_id`), with a `_dedupe_active_scheduled_runs_per_task()` pre-step (keeps the newest active row per task, supersedes the rest to `interrupted` with an explanatory `error` + `finished_at`) mirroring 0004; chains after `0006_agents`
- `migrations/versions/0008_thread_operation_kind.py` — adds `runs.operation_kind` for durable non-run thread reservations; chains after `0007_scheduled_run_active_index`
- `migrations/versions/0010_run_cancel_request.py` — adds the nullable `runs.cancel_action` / `cancel_requested_at` handoff used by non-owning workers; chains after `0009_webhook_dedupe`
- `persistence/bootstrap.py``bootstrap_schema(engine, backend=...)`, the three-branch decision + locking
- Tests: `tests/test_persistence_bootstrap.py` (branches), `tests/test_persistence_bootstrap_concurrency.py` (concurrency), `tests/test_persistence_bootstrap_regression.py` (issue #3682), `tests/test_persistence_migrations_env.py` (filter), `tests/blocking_io/test_persistence_bootstrap.py` (asyncio.to_thread anchor), `tests/test_migration_0004_run_ownership_dedupe.py` + `tests/test_migration_0007_scheduled_run_active_dedupe.py` (dedupe-before-unique-index pre-steps)

View File

@ -839,8 +839,8 @@ async def cancel_run(
- wait=false: Return immediately with 202
In multi-worker deployments, a cancel landing on a non-owning worker
can take over the run when the owner's lease has expired. When the
lease is still valid a 409 + ``Retry-After`` header is returned.
durably notifies the owner when its lease is live, or takes over and
terminalizes the run when that lease has expired.
"""
run_mgr = get_run_manager(request)
record = await run_mgr.get(run_id)
@ -849,15 +849,30 @@ async def cancel_run(
outcome = await run_mgr.cancel(run_id, action=action)
# Success paths — the run was either cancelled locally or taken over
# from a dead worker.
if outcome in (CancelOutcome.cancelled, CancelOutcome.taken_over):
# Success paths — the run was cancelled locally, durably requested from
# a live owner, or taken over from a dead worker.
if outcome in (
CancelOutcome.cancelled,
CancelOutcome.requested,
CancelOutcome.taken_over,
):
if wait and record.task is not None:
try:
await record.task
except asyncio.CancelledError:
pass
return Response(status_code=204)
if wait and outcome == CancelOutcome.requested:
bridge = get_stream_bridge(request)
if record.store_only and bridge.supports_cross_process:
completed = await wait_for_run_completion(
bridge,
record,
request,
run_mgr,
)
if completed:
return Response(status_code=204)
return Response(status_code=202)
if outcome == CancelOutcome.lease_valid_elsewhere:
@ -928,16 +943,29 @@ async def stream_existing_run(
# the client doesn't hang on an SSE subscription this worker can
# never serve.
return Response(status_code=202)
if outcome != CancelOutcome.cancelled:
if outcome not in (CancelOutcome.cancelled, CancelOutcome.requested):
if outcome == CancelOutcome.lease_valid_elsewhere:
await _raise_lease_valid_elsewhere(run_id, run_mgr, record)
raise HTTPException(status_code=409, detail=_cancel_conflict_detail(run_id, record))
if outcome == CancelOutcome.requested and record.store_only and not bridge.supports_cross_process:
# The request is durable, but this bridge cannot observe the
# owner's stream. Returning 202 is safer than hanging forever on
# a process-local subscription.
return Response(status_code=202)
if wait and record.task is not None:
try:
await record.task
except (asyncio.CancelledError, Exception):
pass
return Response(status_code=204)
if wait and outcome == CancelOutcome.requested:
completed = await wait_for_run_completion(
bridge,
record,
request,
run_mgr,
)
return Response(status_code=204 if completed else 202)
return StreamingResponse(
sse_consumer(bridge, record, request, run_mgr),

View File

@ -1283,10 +1283,10 @@ async def sse_consumer(
yield format_sse(entry.event, entry.data, event_id=entry.id or None)
finally:
# store_only records are cross-worker runs hydrated from the RunStore; this
# worker holds no in-memory task/abort state for them, so run_mgr.cancel()
# cannot stop the task (it would 409). Skip on_disconnect cancellation for
# those and only act on runs this worker actually owns.
# store_only records are cross-worker observation handles. An explicit
# cancel-then-stream action has already persisted its request before
# subscribing; a plain join disconnect must not invent a new
# cancellation request. Only apply on_disconnect to locally-owned runs.
if not gap_emitted and not record.store_only and record.status in (RunStatus.pending, RunStatus.running):
if record.on_disconnect == DisconnectMode.cancel:
await run_mgr.cancel(record.run_id)

View File

@ -0,0 +1,37 @@
"""durable cross-worker run cancellation requests.
Revision ID: 0010_run_cancel_request
Revises: 0009_webhook_dedupe
Create Date: 2026-07-27
"""
from __future__ import annotations
from collections.abc import Sequence
import sqlalchemy as sa
revision: str = "0010_run_cancel_request"
down_revision: str | Sequence[str] | None = "0009_webhook_dedupe"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def upgrade() -> None:
from deerflow.persistence.migrations._helpers import safe_add_column
safe_add_column(
"runs",
sa.Column("cancel_action", sa.String(length=20), nullable=True),
)
safe_add_column(
"runs",
sa.Column("cancel_requested_at", sa.DateTime(timezone=True), nullable=True),
)
def downgrade() -> None:
from deerflow.persistence.migrations._helpers import safe_drop_column
safe_drop_column("runs", "cancel_requested_at")
safe_drop_column("runs", "cancel_action")

View File

@ -49,6 +49,10 @@ class RunRow(Base):
# Multi-worker run ownership
owner_worker_id: Mapped[str | None] = mapped_column(String(128), nullable=True)
lease_expires_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
# A non-owning worker records cancellation here; the owner consumes it
# while renewing its lease. The first action wins.
cancel_action: Mapped[str | None] = mapped_column(String(20), nullable=True)
cancel_requested_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(UTC))
updated_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(UTC), onupdate=lambda: datetime.now(UTC))

View File

@ -11,11 +11,15 @@ import json
from datetime import UTC, datetime, timedelta
from typing import Any
from sqlalchemy import or_, select, update
from sqlalchemy import case, or_, select, update
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
from deerflow.persistence.run.model import RunRow
from deerflow.runtime.runs.store.base import RunStore
from deerflow.runtime.runs.store.base import (
LeaseRenewal,
RunStore,
StatusFinalization,
)
from deerflow.runtime.user_context import AUTO, _AutoSentinel, resolve_user_id
from deerflow.utils.time import coerce_iso
@ -77,7 +81,7 @@ class RunRepository(RunStore):
# Convert datetime to ISO string for consistency with MemoryRunStore.
# SQLite drops tzinfo on read despite ``DateTime(timezone=True)`` —
# ``coerce_iso`` normalizes naive datetimes as UTC.
for key in ("created_at", "updated_at", "lease_expires_at"):
for key in ("created_at", "updated_at", "lease_expires_at", "cancel_requested_at"):
val = d.get(key)
if isinstance(val, datetime):
d[key] = coerce_iso(val)
@ -512,6 +516,105 @@ class RunRepository(RunStore):
await session.commit()
return result.rowcount != 0
async def renew_lease(
self,
run_id: str,
*,
owner_worker_id: str,
lease_expires_at: str,
) -> LeaseRenewal:
"""Renew the owner lease and read cancellation intent atomically."""
lease_dt = datetime.fromisoformat(lease_expires_at)
async with self._sf() as session:
result = await session.execute(
update(RunRow)
.where(
RunRow.run_id == run_id,
RunRow.owner_worker_id == owner_worker_id,
RunRow.status.in_(("pending", "running")),
)
.values(
lease_expires_at=lease_dt,
updated_at=datetime.now(UTC),
)
.returning(RunRow.run_id, RunRow.cancel_action)
)
row = result.first()
await session.commit()
if row is None:
return LeaseRenewal(renewed=False)
return LeaseRenewal(renewed=True, cancel_action=row.cancel_action)
async def request_cancel(self, run_id: str, *, action: str) -> str | None:
"""Atomically persist the first cancellation action on an active run."""
if action not in ("interrupt", "rollback"):
raise ValueError(f"Unsupported cancellation action: {action}")
now = datetime.now(UTC)
async with self._sf() as session:
result = await session.execute(
update(RunRow)
.where(
RunRow.run_id == run_id,
RunRow.status.in_(("pending", "running")),
)
.values(
cancel_action=case(
(RunRow.cancel_action.is_(None), action),
else_=RunRow.cancel_action,
),
cancel_requested_at=case(
(RunRow.cancel_requested_at.is_(None), now),
else_=RunRow.cancel_requested_at,
),
updated_at=now,
)
.returning(RunRow.cancel_action)
)
row = result.first()
await session.commit()
return row.cancel_action if row is not None else None
async def finalize_if_not_cancelled(
self,
run_id: str,
*,
status: str,
error: str | None = None,
stop_reason: str | None = None,
) -> StatusFinalization:
"""Atomically let completion win only before cancellation."""
values: dict[str, Any] = {
"status": status,
"updated_at": datetime.now(UTC),
}
if error is not None:
values["error"] = error
if stop_reason is not None:
values["stop_reason"] = stop_reason
async with self._sf() as session:
result = await session.execute(
update(RunRow)
.where(
RunRow.run_id == run_id,
RunRow.status.in_(("pending", "running")),
RunRow.cancel_action.is_(None),
)
.values(**values)
.returning(RunRow.run_id)
)
if result.first() is not None:
await session.commit()
return StatusFinalization(finalized=True)
current = await session.execute(select(RunRow.cancel_action).where(RunRow.run_id == run_id))
cancel_action = current.scalar_one_or_none()
await session.commit()
return StatusFinalization(
finalized=False,
cancel_action=cancel_action,
)
async def claim_for_takeover(
self,
run_id: str,

View File

@ -925,6 +925,65 @@ class RunManager:
)
return persisted
async def set_status_if_not_cancelled(
self,
run_id: str,
status: RunStatus,
*,
error: str | None = None,
stop_reason: str | None = None,
persist: bool = True,
) -> str | None:
"""Set a terminal status unless a durable cancellation won first."""
if not persist or not self.heartbeat_enabled or self._store is None:
await self.set_status(
run_id,
status,
error=error,
stop_reason=stop_reason,
persist=persist,
)
return None
try:
result = await self._call_store_with_retry(
"finalize_if_not_cancelled",
run_id,
lambda: self._store.finalize_if_not_cancelled(
run_id,
status=status.value,
error=error,
stop_reason=stop_reason,
),
)
except Exception:
async with self._lock:
record = self._runs.get(run_id)
if record is not None:
await self._mark_ownership_lost(
record,
reason=("The durable store could not confirm whether cancellation or completion won."),
require_active=False,
)
return None
if result.cancel_action is not None:
async with self._lock:
record = self._runs.get(run_id)
if record is not None:
record.abort_action = result.cancel_action
record.abort_event.set()
return result.cancel_action
await self.set_status(
run_id,
status,
error=error,
stop_reason=stop_reason,
persist=not result.finalized,
)
return None
async def _ensure_delivery_receipt(self, record: RunRecord) -> bool:
"""Idempotently persist a zero-delivery receipt during recovery."""
if self._event_store is None:
@ -1037,6 +1096,94 @@ class RunManager:
await self._persist_model_name(run_id, model_name)
logger.info("Run %s model_name=%s", run_id, model_name)
async def _request_durable_cancel(
self,
run_id: str,
*,
action: str,
) -> tuple[CancelOutcome, str | None]:
"""Record cancellation and return the first action that won."""
if self._store is None:
return CancelOutcome.unknown, None
try:
winning_action = await self._call_store_with_retry(
"request_cancel",
run_id,
lambda: self._store.request_cancel(run_id, action=action),
)
except NotImplementedError:
# Keep third-party stores that predate durable cancellation on the
# old safe behavior instead of pretending the owner was notified.
logger.info(
"Run store does not support cross-worker cancellation for run %s",
run_id,
)
return CancelOutcome.lease_valid_elsewhere, None
except Exception:
logger.warning(
"Failed to persist cancellation request for run %s",
run_id,
exc_info=True,
)
return CancelOutcome.unknown, None
if winning_action is not None:
logger.info(
"Run %s cancellation requested (requested=%s,winner=%s)",
run_id,
action,
winning_action,
)
return CancelOutcome.requested, winning_action
# Completion may have won the race between the caller's read and the
# guarded cancellation UPDATE. Re-read so the API reports that precise
# terminal result rather than claiming the request was accepted.
try:
fresh = await self._store.get(run_id)
except Exception:
fresh = None
if fresh is None:
return CancelOutcome.unknown, None
if fresh.get("status") not in ("pending", "running"):
return CancelOutcome.not_cancellable, None
# A legacy/partial store implementation may decline the request while
# the owner is still live. Preserve the former lease-conflict signal.
return CancelOutcome.lease_valid_elsewhere, None
async def _request_remote_cancel(
self,
run_id: str,
*,
action: str,
) -> CancelOutcome:
"""Record cancellation for a run whose task belongs to another worker."""
outcome, _ = await self._request_durable_cancel(
run_id,
action=action,
)
return outcome
async def _signal_local_cancel(
self,
run_id: str,
*,
action: str,
) -> None:
"""Set process-local abort state without status persistence or cleanup."""
async with self._lock:
record = self._runs.get(run_id)
if record is None or record.status not in (RunStatus.pending, RunStatus.running) or record.abort_event.is_set():
return
record.abort_action = action
record.abort_event.set()
task_active = record.task is not None and not record.task.done()
record.finalizing = task_active
if task_active and record.status == RunStatus.running:
record.task.cancel()
logger.info("Run %s cancellation signalled locally (action=%s)", run_id, action)
async def cancel(self, run_id: str, *, action: str = "interrupt") -> CancelOutcome:
"""Request cancellation of a run.
@ -1051,9 +1198,9 @@ class RunManager:
``error``. The owning worker is assumed dead (its heartbeat
stopped renewing).
- **Lease still valid** returns ``lease_valid_elsewhere`` so
the caller can return HTTP 409 + ``Retry-After`` to tell the
client when to retry.
- **Lease still valid** durably records the cancellation action.
The owner observes it on its next heartbeat and performs the same
local abort/finalization path as a directly-routed request.
In single-worker mode (``heartbeat_enabled=False``) store-only
hydrated runs that aren't in-memory return ``not_active_locally``,
@ -1075,8 +1222,33 @@ class RunManager:
if record is not None:
if record.status == RunStatus.interrupted:
return CancelOutcome.cancelled # idempotent
if record.status not in (RunStatus.pending, RunStatus.running):
if record.status not in (RunStatus.pending, RunStatus.running) and (not self.heartbeat_enabled or self._store is None):
return CancelOutcome.not_cancellable
durable_cancel_won = False
if record is not None and self.heartbeat_enabled and self._store is not None:
outcome, winning_action = await self._request_durable_cancel(
run_id,
action=action,
)
if outcome == CancelOutcome.requested:
action = winning_action or action
durable_cancel_won = True
elif outcome == CancelOutcome.unknown:
logger.warning(
"Proceeding with local cancellation for run %s after durable cancel persistence failed",
run_id,
)
elif outcome != CancelOutcome.lease_valid_elsewhere:
return outcome
async with self._lock:
record = self._runs.get(run_id)
if record is not None:
if record.status == RunStatus.interrupted or record.abort_event.is_set():
return CancelOutcome.cancelled
if record.status not in (RunStatus.pending, RunStatus.running):
return CancelOutcome.cancelled if durable_cancel_won else CancelOutcome.not_cancellable
record.abort_action = action
record.abort_event.set()
task_active = record.task is not None and not record.task.done()
@ -1111,6 +1283,9 @@ class RunManager:
logger.info("Run %s cancelled (action=%s)", run_id, action)
return CancelOutcome.cancelled
if durable_cancel_won:
return CancelOutcome.cancelled
# ------------------------------------------------------------------
# Non-local path — no in-memory record, must consult the store.
# ------------------------------------------------------------------
@ -1131,6 +1306,8 @@ class RunManager:
return CancelOutcome.unknown
store_status = row.get("status")
if store_status == "interrupted":
return CancelOutcome.requested
if store_status not in ("pending", "running"):
return CancelOutcome.not_cancellable
@ -1138,7 +1315,7 @@ class RunManager:
lease_expires_at: str | None = row.get("lease_expires_at")
if not is_lease_expired(lease_expires_at, grace_seconds=grace_seconds):
return CancelOutcome.lease_valid_elsewhere
return await self._request_remote_cancel(run_id, action=action)
take_over_msg = f"Run reclaimed by worker {self._worker_id}: the owning worker ({row.get('owner_worker_id') or 'unknown'}) stopped renewing its lease and is presumed dead."
try:
@ -1160,7 +1337,7 @@ class RunManager:
return CancelOutcome.taken_over
# The conditional UPDATE matched 0 rows. Two causes:
# (a) the owner renewed the lease → lease_valid_elsewhere.
# (a) the owner renewed the lease → persist a cancellation request.
# (b) the row went terminal between our read and the claim
# (run finished, or another worker already took it over)
# → not_cancellable or taken_over.
@ -1177,8 +1354,9 @@ class RunManager:
logger.info("Run %s takeover lost to another worker already at error", run_id)
return CancelOutcome.taken_over
return CancelOutcome.not_cancellable
# Row is still active — lease must have been renewed by the owner.
return CancelOutcome.lease_valid_elsewhere
# Row is still active — lease was renewed by the owner while the
# takeover raced. Notify that owner instead of exposing routing as 409.
return await self._request_remote_cancel(run_id, action=action)
def _compute_lease_expires_at(self) -> str | None:
"""Return the lease expiry ISO timestamp for a freshly created run.
@ -1788,6 +1966,7 @@ class RunManager:
if self._store is None or self._run_ownership_config is None:
return
lease_seconds = self._run_ownership_config.lease_seconds
cancellations: list[tuple[str, str]] = []
async with self._lock:
# Renew any pending/running run owned by this worker unless its
@ -1815,16 +1994,16 @@ class RunManager:
new_expiry = (datetime.now(UTC) + timedelta(seconds=lease_seconds)).isoformat()
try:
async with asyncio.timeout(remaining):
updated = await self._call_store_with_retry(
"update_lease",
renewal = await self._call_store_with_retry(
"renew_lease",
run_id,
lambda: self._store.update_lease(
lambda: self._store.renew_lease(
run_id,
owner_worker_id=self._worker_id,
lease_expires_at=new_expiry,
),
)
if updated:
if renewal.renewed:
if confirmed_deadline <= datetime.now(UTC):
await self._mark_ownership_lost(
record,
@ -1838,8 +2017,18 @@ class RunManager:
# fields). Re-acquiring ``self._lock`` here would
# serialise against unrelated run mutations for no gain.
record.lease_expires_at = new_expiry
if renewal.cancel_action is not None:
action = renewal.cancel_action
if action not in ("interrupt", "rollback"):
logger.warning(
"Run %s has invalid durable cancel action %r; using interrupt",
run_id,
action,
)
action = "interrupt"
cancellations.append((run_id, action))
else:
# ``update_lease`` returned False — the row was claimed
# ``renew_lease`` returned False — the row was claimed
# by another worker (status is no longer pending/running,
# or ``owner_worker_id`` changed). Stop the local task so
# we don't waste CPU or overwrite the takeover status on
@ -1870,6 +2059,15 @@ class RunManager:
exc_info=True,
)
# Keep cancellation status writes and cleanup out of the sole renewal
# loop. After every local lease has had a chance to renew, only signal
# the owning worker task; that task performs normal terminal handling.
for run_id, action in cancellations:
await self._signal_local_cancel(
run_id,
action=action,
)
async def _reconcile_orphans_periodic(self) -> None:
"""Sweep for expired leases owned by dead peers.
@ -2040,6 +2238,7 @@ class CancelOutcome(StrEnum):
"""Result of a :meth:`RunManager.cancel` call."""
cancelled = "cancelled"
requested = "requested"
taken_over = "taken_over"
lease_valid_elsewhere = "lease_valid_elsewhere"
not_cancellable = "not_cancellable"

View File

@ -1,4 +1,4 @@
from deerflow.runtime.runs.store.base import RunStore
from deerflow.runtime.runs.store.base import LeaseRenewal, RunStore
from deerflow.runtime.runs.store.memory import MemoryRunStore
__all__ = ["MemoryRunStore", "RunStore"]
__all__ = ["LeaseRenewal", "MemoryRunStore", "RunStore"]

View File

@ -21,6 +21,26 @@ class EditReplayVisibility:
hidden_attempt_run_ids: set[str] = field(default_factory=set)
@dataclass(frozen=True)
class LeaseRenewal:
"""Result of renewing a run lease.
``cancel_action`` carries a durable cancellation request to the owning
worker without transferring lease ownership.
"""
renewed: bool
cancel_action: str | None = None
@dataclass(frozen=True)
class StatusFinalization:
"""Result of completing a run only if cancellation has not won."""
finalized: bool
cancel_action: str | None = None
class RunStore(abc.ABC):
@abc.abstractmethod
async def put(
@ -217,6 +237,57 @@ class RunStore(abc.ABC):
"""Renew the lease on an active run. Returns ``False`` when no row matched."""
pass
async def renew_lease(
self,
run_id: str,
*,
owner_worker_id: str,
lease_expires_at: str,
) -> LeaseRenewal:
"""Renew ownership and return any durable cancellation request.
The default wraps the legacy ``update_lease`` method and returns no
cancellation action, so third-party stores remain source-compatible
without adding a background read. Stores that support multi-process
cancellation must override this method to renew and observe the
request atomically.
"""
renewed = await self.update_lease(
run_id,
owner_worker_id=owner_worker_id,
lease_expires_at=lease_expires_at,
)
return LeaseRenewal(renewed=renewed)
async def request_cancel(self, run_id: str, *, action: str) -> str | None:
"""Persist the first cancellation action for an active run.
Implementations must update only ``pending`` or ``running`` rows and
return the winning action, or ``None`` when no active row matched.
"""
raise NotImplementedError
async def finalize_if_not_cancelled(
self,
run_id: str,
*,
status: str,
error: str | None = None,
stop_reason: str | None = None,
) -> StatusFinalization:
"""Atomically finalize an active run unless cancellation won.
The compatibility default is safe for stores that do not implement
durable cancellation.
"""
updated = await self.update_status(
run_id,
status,
error=error,
stop_reason=stop_reason,
)
return StatusFinalization(finalized=updated is not False)
@abc.abstractmethod
async def claim_for_takeover(
self,

View File

@ -8,7 +8,7 @@ from __future__ import annotations
from datetime import UTC, datetime, timedelta
from typing import Any
from deerflow.runtime.runs.store.base import RunStore
from deerflow.runtime.runs.store.base import LeaseRenewal, RunStore, StatusFinalization
class MemoryRunStore(RunStore):
@ -52,6 +52,7 @@ class MemoryRunStore(RunStore):
lease_expires_at=None,
):
now = datetime.now(UTC).isoformat()
existing = self._runs.get(run_id)
self._runs[run_id] = {
"run_id": run_id,
"thread_id": thread_id,
@ -69,6 +70,10 @@ class MemoryRunStore(RunStore):
"updated_at": now,
"owner_worker_id": owner_worker_id,
"lease_expires_at": lease_expires_at,
# ``put`` is an idempotent snapshot write. Preserve a cancellation
# request that may have raced a retry of an earlier snapshot.
"cancel_action": existing.get("cancel_action") if existing else None,
"cancel_requested_at": existing.get("cancel_requested_at") if existing else None,
}
self._index_run(run_id, thread_id)
@ -254,6 +259,66 @@ class MemoryRunStore(RunStore):
run["updated_at"] = datetime.now(UTC).isoformat()
return True
async def renew_lease(
self,
run_id: str,
*,
owner_worker_id: str,
lease_expires_at: str,
) -> LeaseRenewal:
# Delegate through ``update_lease`` so lightweight subclasses and tests
# that override the legacy primitive keep the same behavior.
renewed = await self.update_lease(
run_id,
owner_worker_id=owner_worker_id,
lease_expires_at=lease_expires_at,
)
if not renewed:
return LeaseRenewal(renewed=False)
run = self._runs.get(run_id)
return LeaseRenewal(
renewed=True,
cancel_action=run.get("cancel_action") if run is not None else None,
)
async def request_cancel(self, run_id: str, *, action: str) -> str | None:
if action not in ("interrupt", "rollback"):
raise ValueError(f"Unsupported cancellation action: {action}")
run = self._runs.get(run_id)
if run is None or run["status"] not in ("pending", "running"):
return None
if run.get("cancel_action") is None:
run["cancel_action"] = action
run["cancel_requested_at"] = datetime.now(UTC).isoformat()
run["updated_at"] = datetime.now(UTC).isoformat()
return run["cancel_action"]
async def finalize_if_not_cancelled(
self,
run_id: str,
*,
status: str,
error: str | None = None,
stop_reason: str | None = None,
) -> StatusFinalization:
run = self._runs.get(run_id)
if run is None:
return StatusFinalization(finalized=False)
if run.get("cancel_action") is not None:
return StatusFinalization(
finalized=False,
cancel_action=run["cancel_action"],
)
if run["status"] not in ("pending", "running"):
return StatusFinalization(finalized=False)
run["status"] = status
if error is not None:
run["error"] = error
if stop_reason is not None:
run["stop_reason"] = stop_reason
run["updated_at"] = datetime.now(UTC).isoformat()
return StatusFinalization(finalized=True)
async def claim_for_takeover(
self,
run_id: str,
@ -409,6 +474,8 @@ class MemoryRunStore(RunStore):
"error": None,
"owner_worker_id": owner_worker_id,
"lease_expires_at": lease_expires_at,
"cancel_action": None,
"cancel_requested_at": None,
"created_at": created_at or now,
"updated_at": now,
}

View File

@ -548,6 +548,50 @@ async def run_agent(
subagent_events: _SubagentEventBuffer | None = None
started = False
async def _finish_cancellation(
action: str,
*,
restore_checkpoint: bool = True,
) -> None:
nonlocal checkpoint_rollback_completed
await run_manager.set_finalizing(run_id, True)
if action == "rollback":
await run_manager.set_status(
run_id,
RunStatus.error,
error="Rolled back by user",
**terminal_status_kwargs,
)
if not restore_checkpoint:
return
try:
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,
)
logger.info(
"Run %s rolled back to pre-run checkpoint %s",
run_id,
pre_run_checkpoint_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,
**terminal_status_kwargs,
)
logger.info("Run %s was cancelled", run_id)
try:
normalized_stream_modes = normalize_stream_modes(stream_modes)
requested_modes: set[str] = set(normalized_stream_modes)
@ -576,6 +620,11 @@ async def run_agent(
start_outcome = await run_manager.try_start(run_id)
if start_outcome is not RunStartOutcome.started:
if record.abort_event.is_set():
await _finish_cancellation(
record.abort_action,
restore_checkpoint=False,
)
return
started = True
@ -881,45 +930,21 @@ async def run_agent(
# 8. Final status
if record.abort_event.is_set():
await run_manager.set_finalizing(run_id, True)
action = record.abort_action
if action == "rollback":
await run_manager.set_status(
run_id,
RunStatus.error,
error="Rolled back by user",
**terminal_status_kwargs,
)
try:
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,
)
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,
**terminal_status_kwargs,
)
await _finish_cancellation(record.abort_action)
elif llm_error_fallback_message or (journal is not None and journal.had_llm_error_fallback):
error_msg = llm_error_fallback_message
if error_msg is None and journal is not None:
error_msg = journal.llm_error_fallback_message
error_msg = error_msg or "LLM provider failed after retries"
await _ensure_finalizing_before_edit_failure(run_manager, record)
await run_manager.set_status(
cancel_action = await run_manager.set_status_if_not_cancelled(
run_id,
RunStatus.error,
error=error_msg,
**terminal_status_kwargs,
)
if cancel_action is not None:
await _finish_cancellation(cancel_action)
else:
runtime_context = runtime.context if isinstance(runtime.context, dict) else None
# Guard middlewares that hard-stop a run by stripping tool_calls
@ -947,63 +972,40 @@ async def run_agent(
produced_output_paths,
)
delivery_error = _delivery_error(delivery_content)
await run_manager.set_status(
cancel_action = await run_manager.set_status_if_not_cancelled(
run_id,
RunStatus.error if delivery_error else RunStatus.success,
error=delivery_error,
stop_reason=stop_reason,
**terminal_status_kwargs,
)
if cancel_action is not None:
await _finish_cancellation(cancel_action)
except asyncio.CancelledError:
await run_manager.set_finalizing(run_id, True)
action = record.abort_action
if action == "rollback":
await run_manager.set_status(
run_id,
RunStatus.error,
error="Rolled back by user",
**terminal_status_kwargs,
)
try:
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,
)
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 _ensure_finalizing_before_edit_failure(run_manager, record)
await run_manager.set_status(
run_id,
RunStatus.interrupted,
**terminal_status_kwargs,
)
logger.info("Run %s was cancelled", run_id)
await _finish_cancellation(record.abort_action)
except Exception as exc:
error_msg = f"{exc}"
logger.exception("Run %s failed: %s", run_id, error_msg)
await _ensure_finalizing_before_edit_failure(run_manager, record)
await run_manager.set_status(
cancel_action = await run_manager.set_status_if_not_cancelled(
run_id,
RunStatus.error,
error=error_msg,
**terminal_status_kwargs,
)
await bridge.publish(
run_id,
"error",
{
"message": error_msg,
"name": type(exc).__name__,
},
)
if cancel_action is not None:
await _finish_cancellation(cancel_action)
else:
await bridge.publish(
run_id,
"error",
{
"message": error_msg,
"name": type(exc).__name__,
},
)
finally:
if record.ownership_lost:
@ -1092,7 +1094,18 @@ async def run_agent(
# real worker outcome. Leaving a successful row inflight would
# let lease recovery rewrite it as an error with a synthetic
# zero receipt.
await run_manager.persist_current_status(run_id)
if record.abort_event.is_set():
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,
)
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)

View File

@ -2258,10 +2258,17 @@ async def test_run_agent_full_mode_rejects_delta_before_graph_invocation():
publish_end=AsyncMock(),
cleanup=AsyncMock(),
)
set_status = AsyncMock()
async def set_status_if_not_cancelled(*args, **kwargs):
await set_status(*args, **kwargs)
return None
run_manager = SimpleNamespace(
try_start=AsyncMock(return_value=RunStartOutcome.started),
wait_for_prior_finalizing=AsyncMock(),
set_status=AsyncMock(),
set_status=set_status,
set_status_if_not_cancelled=AsyncMock(side_effect=set_status_if_not_cancelled),
)
record = RunRecord(
run_id="run-checkpoint-mode",
@ -2330,10 +2337,17 @@ async def test_run_agent_full_mode_checks_selected_checkpoint_before_graph():
publish_end=AsyncMock(),
cleanup=AsyncMock(),
)
set_status = AsyncMock()
async def set_status_if_not_cancelled(*args, **kwargs):
await set_status(*args, **kwargs)
return None
run_manager = SimpleNamespace(
try_start=AsyncMock(return_value=RunStartOutcome.started),
wait_for_prior_finalizing=AsyncMock(),
set_status=AsyncMock(),
set_status=set_status,
set_status_if_not_cancelled=AsyncMock(side_effect=set_status_if_not_cancelled),
)
record = RunRecord(
run_id="run-selected-checkpoint-mode",

View File

@ -613,6 +613,10 @@ async def test_run_agent_does_not_stream_continuation_after_abort(monkeypatch):
async def set_status(self, _run_id, status, **_kwargs):
record.status = status
async def set_status_if_not_cancelled(self, _run_id, status, **kwargs):
await self.set_status(_run_id, status, **kwargs)
return None
async def update_model_name(self, *_args, **_kwargs):
return None
@ -694,6 +698,10 @@ async def test_run_agent_reuses_goal_evaluator_model_for_goal_loop(monkeypatch):
async def set_status(self, _run_id, status, **_kwargs):
record.status = status
async def set_status_if_not_cancelled(self, _run_id, status, **kwargs):
await self.set_status(_run_id, status, **kwargs)
return None
async def update_model_name(self, *_args, **_kwargs):
return None
@ -874,6 +882,10 @@ async def test_run_agent_strips_branch_checkpoint_for_goal_continuation(monkeypa
async def set_status(self, _run_id, status, **_kwargs):
record.status = status
async def set_status_if_not_cancelled(self, _run_id, status, **kwargs):
await self.set_status(_run_id, status, **kwargs)
return None
async def update_model_name(self, *_args, **_kwargs):
return None

View File

@ -157,7 +157,7 @@ async def test_migration_dedupes_duplicate_active_rows_before_unique_index(tmp_p
with sqlite3.connect(db_path) as raw:
version_row = raw.execute("SELECT version_num FROM alembic_version").fetchone()
# Bootstrap upgrades through the later revisions after 0004.
assert version_row[0] == "0009_webhook_dedupe"
assert version_row[0] == "0010_run_cancel_request"
# Sanity: the invariant the index enforces is now true — at most one
# active row per thread.

View File

@ -169,7 +169,7 @@ async def test_migration_supersedes_duplicate_active_runs_before_unique_index(tm
with sqlite3.connect(db_path) as raw:
version_row = raw.execute("SELECT version_num FROM alembic_version").fetchone()
assert version_row[0] == "0009_webhook_dedupe"
assert version_row[0] == "0010_run_cancel_request"
# Sanity: the invariant the index enforces now holds — at most one
# active row per task_id.

View File

@ -1650,8 +1650,8 @@ async def test_cancel_takeover_from_crashed_worker():
@pytest.mark.anyio
async def test_cancel_refuses_active_lease_from_other_worker():
"""cancel must return lease_valid_elsewhere when the run is owned by another worker with a valid lease."""
async def test_cancel_requests_active_lease_from_other_worker():
"""A cancel routed to a peer must durably notify the live owner."""
store = MemoryRunStore()
grace = 10
valid_lease = (datetime.now(UTC) + timedelta(seconds=60)).isoformat()
@ -1659,11 +1659,168 @@ async def test_cancel_refuses_active_lease_from_other_worker():
manager = _make_manager(store=store, run_ownership_config=_lease_config(heartbeat_enabled=True, grace_seconds=grace))
outcome = await manager.cancel("run-alive")
assert outcome == CancelOutcome.lease_valid_elsewhere
assert outcome == CancelOutcome.requested
row = await store.get("run-alive")
assert row is not None
assert row["status"] == "running" # untouched
assert row["status"] == "running"
assert row["cancel_action"] == "interrupt"
assert row["cancel_requested_at"] is not None
@pytest.mark.anyio
async def test_non_owner_cancel_is_observed_by_owner_heartbeat():
"""Heartbeat signals the owner task without performing terminal writes."""
store = MemoryRunStore()
config = _lease_config(heartbeat_enabled=True, lease_seconds=30)
owner = _make_manager(store=store, worker_id="worker-a", run_ownership_config=config)
peer = _make_manager(store=store, worker_id="worker-b", run_ownership_config=config)
record = await owner.create_or_reject("thread-1")
await owner.set_status(record.run_id, RunStatus.running)
record.task = asyncio.create_task(asyncio.sleep(3600))
try:
assert await peer.cancel(record.run_id, action="rollback") == CancelOutcome.requested
# Match local idempotency: once accepted, a later action cannot change
# whether the owner rolls back or merely interrupts.
assert await peer.cancel(record.run_id, action="interrupt") == CancelOutcome.requested
await owner._renew_leases()
await asyncio.sleep(0)
assert record.abort_event.is_set()
assert record.abort_action == "rollback"
assert record.status == RunStatus.running
assert record.task.cancelled()
stored = await store.get(record.run_id)
assert stored is not None
assert stored["status"] == "running"
assert stored["cancel_action"] == "rollback"
# The worker's existing cancellation path, not heartbeat, owns the
# terminal status write and any rollback cleanup.
await owner.set_status(
record.run_id,
RunStatus.error,
error="Rolled back by user",
)
next_run = await peer.create_or_reject("thread-1")
assert next_run.status == RunStatus.pending
finally:
if not record.task.done():
record.task.cancel()
with pytest.raises(asyncio.CancelledError):
await record.task
@pytest.mark.anyio
async def test_first_cancel_action_wins_when_retry_lands_on_owner():
"""Routing a retry to the owner must not replace the durable first action."""
store = MemoryRunStore()
config = _lease_config(heartbeat_enabled=True, lease_seconds=30)
owner = _make_manager(store=store, worker_id="worker-a", run_ownership_config=config)
peer = _make_manager(store=store, worker_id="worker-b", run_ownership_config=config)
record = await owner.create_or_reject("thread-1")
await owner.set_status(record.run_id, RunStatus.running)
record.task = asyncio.create_task(asyncio.sleep(3600))
try:
assert await peer.cancel(record.run_id, action="rollback") == CancelOutcome.requested
assert await owner.cancel(record.run_id, action="interrupt") == CancelOutcome.cancelled
await asyncio.sleep(0)
assert record.abort_action == "rollback"
assert record.task.cancelled()
stored = await store.get(record.run_id)
assert stored is not None
assert stored["cancel_action"] == "rollback"
finally:
if not record.task.done():
record.task.cancel()
with pytest.raises(asyncio.CancelledError):
await record.task
@pytest.mark.anyio
async def test_local_owner_cancel_falls_back_when_durable_request_fails():
"""A local owner can still abort its own task when durable cancel persistence fails."""
class FailingCancelStore(MemoryRunStore):
async def request_cancel(self, run_id: str, *, action: str) -> str | None:
raise RuntimeError("store unavailable")
store = FailingCancelStore()
manager = _make_manager(
store=store,
worker_id="worker-a",
run_ownership_config=_lease_config(heartbeat_enabled=True, lease_seconds=30),
)
record = await manager.create_or_reject("thread-1")
await manager.set_status(record.run_id, RunStatus.running)
record.task = asyncio.create_task(asyncio.sleep(3600))
try:
assert await manager.cancel(record.run_id, action="rollback") == CancelOutcome.cancelled
assert record.abort_event.is_set()
assert record.abort_action == "rollback"
with pytest.raises(asyncio.CancelledError):
await record.task
stored = await store.get(record.run_id)
assert stored is not None
assert stored["status"] == "interrupted"
assert stored["cancel_action"] is None
finally:
if not record.task.done():
record.task.cancel()
with pytest.raises(asyncio.CancelledError):
await record.task
@pytest.mark.anyio
async def test_owner_cancel_retry_on_peer_is_accepted():
"""A peer must treat the owner's interrupted status as an accepted cancel."""
store = MemoryRunStore()
config = _lease_config(heartbeat_enabled=True, lease_seconds=30)
owner = _make_manager(store=store, worker_id="worker-a", run_ownership_config=config)
peer = _make_manager(store=store, worker_id="worker-b", run_ownership_config=config)
record = await owner.create_or_reject("thread-1")
await owner.set_status(record.run_id, RunStatus.running)
assert await owner.cancel(record.run_id, action="rollback") == CancelOutcome.cancelled
assert await peer.cancel(record.run_id, action="interrupt") == CancelOutcome.requested
stored = await store.get(record.run_id)
assert stored is not None
assert stored["status"] == "interrupted"
assert stored["cancel_action"] == "rollback"
@pytest.mark.anyio
async def test_owner_cancel_uses_store_while_terminal_status_is_staged_locally():
"""A staged local terminal status must not bypass the durable cancel CAS."""
store = MemoryRunStore()
manager = _make_manager(
store=store,
run_ownership_config=_lease_config(heartbeat_enabled=True),
)
record = await manager.create_or_reject("thread-1")
await manager.set_status(record.run_id, RunStatus.running)
# Event-store finalization stages success in memory before persisting it.
record.status = RunStatus.success
assert await manager.cancel(record.run_id, action="rollback") == CancelOutcome.cancelled
stored = await store.get(record.run_id)
assert stored is not None
assert stored["status"] == "running"
assert stored["cancel_action"] == "rollback"
@pytest.mark.anyio
@ -1687,7 +1844,7 @@ async def test_cancel_returns_not_active_locally_when_heartbeat_disabled():
@pytest.mark.anyio
async def test_cancel_takeover_race_owner_renewed_lease():
"""When the owner heartbeats between our read and the conditional UPDATE, cancel must return lease_valid_elsewhere."""
"""When takeover loses to a renewal, cancellation must notify the live owner."""
store = MemoryRunStore()
grace = 10
expired_lease = (datetime.now(UTC) - timedelta(seconds=grace + 5)).isoformat()
@ -1709,12 +1866,13 @@ async def test_cancel_takeover_race_owner_renewed_lease():
manager = _make_manager(store=store, run_ownership_config=_lease_config(heartbeat_enabled=True, grace_seconds=grace))
outcome = await manager.cancel("run-race")
assert outcome == CancelOutcome.lease_valid_elsewhere
assert outcome == CancelOutcome.requested
assert (await store.get("run-race"))["cancel_action"] == "interrupt"
@pytest.mark.anyio
async def test_cancel_takeover_respects_grace_seconds():
"""Cancel must not take over when the lease is within the grace window."""
"""Within the grace window, cancellation must notify rather than take over."""
store = MemoryRunStore()
grace = 10
# Lease expired, but only by 3s — still within the 10s grace window
@ -1723,7 +1881,8 @@ async def test_cancel_takeover_respects_grace_seconds():
manager = _make_manager(store=store, run_ownership_config=_lease_config(heartbeat_enabled=True, grace_seconds=grace))
outcome = await manager.cancel("run-grace")
assert outcome == CancelOutcome.lease_valid_elsewhere
assert outcome == CancelOutcome.requested
assert (await store.get("run-grace"))["cancel_action"] == "interrupt"
@pytest.mark.anyio
@ -1742,7 +1901,28 @@ async def test_cancel_not_cancellable_for_store_terminal_run():
# ---------------------------------------------------------------------------
def _make_cancel_test_app(mgr: RunManager):
class _EndingCrossProcessBridge:
supports_cross_process = True
async def publish(self, run_id, event, data):
return None
async def publish_end(self, run_id):
return None
def subscribe(self, run_id, *, last_event_id=None, heartbeat_interval=15.0):
from deerflow.runtime import END_SENTINEL
async def events():
yield END_SENTINEL
return events()
async def cleanup(self, run_id, *, delay=0):
return None
def _make_cancel_test_app(mgr: RunManager, *, bridge=None):
"""Build a TestClient wired with the thread_runs router + memory bridge."""
from _router_auth_helpers import make_authed_test_app
from fastapi.testclient import TestClient
@ -1753,12 +1933,12 @@ def _make_cancel_test_app(mgr: RunManager):
app = make_authed_test_app()
app.include_router(thread_runs.router)
app.state.run_manager = mgr
app.state.stream_bridge = MemoryStreamBridge()
app.state.stream_bridge = bridge or MemoryStreamBridge()
return TestClient(app, raise_server_exceptions=False)
def test_http_cancel_non_owner_valid_lease_returns_409_with_retry_after():
"""POST /cancel on a non-owning worker with a valid lease must return 409 + Retry-After."""
def test_http_cancel_non_owner_valid_lease_returns_202():
"""POST /cancel must not fail solely because routing chose a non-owner."""
store = MemoryRunStore()
grace = 10
valid_lease = (datetime.now(UTC) + timedelta(seconds=60)).isoformat()
@ -1776,15 +1956,105 @@ def test_http_cancel_non_owner_valid_lease_returns_409_with_retry_after():
client = _make_cancel_test_app(mgr)
resp = client.post("/api/threads/t1/runs/run-alive/cancel")
assert resp.status_code == 409
assert "Retry-After" in resp.headers
# Retry-After = remaining lease (≈60s) + grace (10s) = ≈70s
retry_after = int(resp.headers["Retry-After"])
assert 50 <= retry_after <= 75
assert resp.status_code == 202
assert "Retry-After" not in resp.headers
# Store row must be untouched
# The owner remains fenced, while the cancellation request is durable.
row = asyncio.run(store.get("run-alive"))
assert row["status"] == "running"
assert row["cancel_action"] == "interrupt"
def test_http_stream_action_non_owner_without_shared_bridge_returns_202():
"""A peer cancel is accepted without subscribing to an unreachable local stream."""
store = MemoryRunStore()
valid_lease = (datetime.now(UTC) + timedelta(seconds=60)).isoformat()
asyncio.run(
store.put(
"run-alive-stream",
thread_id="t1",
status="running",
created_at=datetime.now(UTC).isoformat(),
owner_worker_id="alive-worker",
lease_expires_at=valid_lease,
)
)
mgr = _make_manager(
store=store,
run_ownership_config=_lease_config(heartbeat_enabled=True),
)
client = _make_cancel_test_app(mgr)
resp = client.post(
"/api/threads/t1/runs/run-alive-stream/stream",
params={"action": "interrupt"},
)
assert resp.status_code == 202
row = asyncio.run(store.get("run-alive-stream"))
assert row["status"] == "running"
assert row["cancel_action"] == "interrupt"
def test_http_cancel_non_owner_wait_uses_shared_bridge():
"""wait=true observes remote owner finalization through the shared bridge."""
store = MemoryRunStore()
valid_lease = (datetime.now(UTC) + timedelta(seconds=60)).isoformat()
asyncio.run(
store.put(
"run-alive-wait",
thread_id="t1",
status="running",
created_at=datetime.now(UTC).isoformat(),
owner_worker_id="alive-worker",
lease_expires_at=valid_lease,
)
)
mgr = _make_manager(
store=store,
run_ownership_config=_lease_config(heartbeat_enabled=True),
)
client = _make_cancel_test_app(mgr, bridge=_EndingCrossProcessBridge())
resp = client.post(
"/api/threads/t1/runs/run-alive-wait/cancel",
params={"action": "rollback", "wait": "true"},
)
assert resp.status_code == 204
row = asyncio.run(store.get("run-alive-wait"))
assert row["cancel_action"] == "rollback"
def test_http_stream_action_non_owner_uses_shared_bridge():
"""The SDK stop path drains the remote owner's shared stream after accept."""
store = MemoryRunStore()
valid_lease = (datetime.now(UTC) + timedelta(seconds=60)).isoformat()
asyncio.run(
store.put(
"run-alive-shared-stream",
thread_id="t1",
status="running",
created_at=datetime.now(UTC).isoformat(),
owner_worker_id="alive-worker",
lease_expires_at=valid_lease,
)
)
mgr = _make_manager(
store=store,
run_ownership_config=_lease_config(heartbeat_enabled=True),
)
client = _make_cancel_test_app(mgr, bridge=_EndingCrossProcessBridge())
resp = client.post(
"/api/threads/t1/runs/run-alive-shared-stream/stream",
params={"action": "interrupt"},
)
assert resp.status_code == 200
assert "event: end" in resp.text
row = asyncio.run(store.get("run-alive-shared-stream"))
assert row["cancel_action"] == "interrupt"
def test_http_cancel_non_owner_expired_lease_returns_202_takeover():
@ -2171,34 +2441,3 @@ def test_compute_retry_after_normal():
assert val is not None
# lease_expires_at is ~45s from now + grace_seconds 10 = ~55, within reason
assert 40 <= val <= 65
# ---------------------------------------------------------------------------
# HTTP — stream endpoint cross-worker 409
# ---------------------------------------------------------------------------
def test_http_stream_action_interrupt_non_owner_returns_409_with_retry_after():
"""POST /stream?action=interrupt on a non-owner with valid lease must
return 409 + Retry-After, not hang on SSE."""
store = MemoryRunStore()
grace = 10
valid_lease = (datetime.now(UTC) + timedelta(seconds=60)).isoformat()
asyncio.run(
store.put(
"run-alive-stream",
thread_id="t1",
status="running",
owner_worker_id="alive-worker",
lease_expires_at=valid_lease,
created_at=datetime.now(UTC).isoformat(),
)
)
mgr = _make_manager(store=store, run_ownership_config=_lease_config(heartbeat_enabled=True, grace_seconds=grace))
client = _make_cancel_test_app(mgr)
resp = client.post("/api/threads/t1/runs/run-alive-stream/stream", params={"action": "interrupt"})
assert resp.status_code == 409
assert "Retry-After" in resp.headers
retry_after = int(resp.headers["Retry-After"])
assert 50 <= retry_after <= 75

View File

@ -48,7 +48,7 @@ from deerflow.persistence.migrations._helpers import _normalize_default
asyncio_test = pytest.mark.asyncio
HEAD = "0009_webhook_dedupe"
HEAD = "0010_run_cancel_request"
BASELINE = "0001_baseline"
@ -147,6 +147,8 @@ async def test_empty_branch_creates_all_and_stamps_head(tmp_path: Path) -> None:
}:
assert required in tables, f"missing table: {required}"
assert "token_usage_by_model" in await _runs_columns(engine)
assert "cancel_action" in await _runs_columns(engine)
assert "cancel_requested_at" in await _runs_columns(engine)
operation_kind = await _runs_column_meta(engine, "operation_kind")
assert operation_kind["nullable"] is False
assert await _alembic_version(engine) == HEAD
@ -849,7 +851,7 @@ class TestDecideState:
# ---------------------------------------------------------------------------
def test_head_revision_is_token_usage_revision() -> None:
def test_head_revision_is_expected() -> None:
assert _get_head_revision() == HEAD

View File

@ -28,7 +28,7 @@ from deerflow.persistence.bootstrap import bootstrap_schema
pytestmark = pytest.mark.asyncio
HEAD = "0009_webhook_dedupe"
HEAD = "0010_run_cancel_request"
def _url(tmp_path: Path) -> str:

View File

@ -76,7 +76,7 @@ async def test_legacy_database_recovers_token_usage_column(tmp_path: Path) -> No
cols = {row[1] for row in raw.execute("PRAGMA table_info(runs)").fetchall()}
assert "token_usage_by_model" in cols
version_row = raw.execute("SELECT version_num FROM alembic_version").fetchone()
assert version_row[0] == "0009_webhook_dedupe"
assert version_row[0] == "0010_run_cancel_request"
# And the read path that originally 500'd must now succeed.
sf = get_session_factory()
@ -116,6 +116,6 @@ async def test_legacy_database_with_manual_alter_still_bootstraps(tmp_path: Path
# No duplicate column -- list, not set, to catch dupes.
assert cols.count("token_usage_by_model") == 1
version_row = raw.execute("SELECT version_num FROM alembic_version").fetchone()
assert version_row[0] == "0009_webhook_dedupe"
assert version_row[0] == "0010_run_cancel_request"
finally:
await close_engine()

View File

@ -992,6 +992,84 @@ class TestRunRepository:
assert row["status"] == "running"
await _cleanup()
@pytest.mark.anyio
async def test_request_cancel_is_returned_by_owner_lease_renewal(self, tmp_path):
"""The SQL store must atomically carry the first cancel action to the owner."""
repo = await _make_repo(tmp_path)
lease = (datetime.now(UTC) + timedelta(seconds=30)).isoformat()
await repo.put(
"run-1",
thread_id="t1",
status="running",
owner_worker_id="worker-a",
lease_expires_at=lease,
)
assert await repo.request_cancel("run-1", action="rollback") == "rollback"
assert await repo.request_cancel("run-1", action="interrupt") == "rollback"
renewal = await repo.renew_lease(
"run-1",
owner_worker_id="worker-a",
lease_expires_at=(datetime.now(UTC) + timedelta(seconds=60)).isoformat(),
)
assert renewal.renewed is True
assert renewal.cancel_action == "rollback"
row = await repo.get("run-1")
assert row["cancel_action"] == "rollback"
assert row["cancel_requested_at"] is not None
await _cleanup()
@pytest.mark.anyio
async def test_request_cancel_rejects_terminal_run(self, tmp_path):
repo = await _make_repo(tmp_path)
await repo.put("run-1", thread_id="t1", status="success")
assert await repo.request_cancel("run-1", action="interrupt") is None
await _cleanup()
@pytest.mark.anyio
async def test_cancel_request_wins_before_owner_completion(self, tmp_path):
repo = await _make_repo(tmp_path)
await repo.put(
"run-1",
thread_id="t1",
status="running",
owner_worker_id="worker-a",
)
assert await repo.request_cancel("run-1", action="rollback") == "rollback"
result = await repo.finalize_if_not_cancelled(
"run-1",
status="success",
)
assert result.finalized is False
assert result.cancel_action == "rollback"
assert (await repo.get("run-1"))["status"] == "running"
await _cleanup()
@pytest.mark.anyio
async def test_owner_completion_wins_before_cancel_request(self, tmp_path):
repo = await _make_repo(tmp_path)
await repo.put(
"run-1",
thread_id="t1",
status="running",
owner_worker_id="worker-a",
)
result = await repo.finalize_if_not_cancelled(
"run-1",
status="success",
)
assert result.finalized is True
assert await repo.request_cancel("run-1", action="rollback") is None
assert (await repo.get("run-1"))["status"] == "success"
await _cleanup()
@pytest.mark.anyio
async def test_reconciliation_skips_run_renewed_after_scan(self, tmp_path):
"""The SQL takeover CAS must reject a candidate renewed after its scan."""

View File

@ -16,10 +16,13 @@ from langgraph.graph.message import add_messages
from langgraph.types import Overwrite
from deerflow.agents.thread_state import merge_artifacts, merge_message_writes
from deerflow.config.run_ownership_config import RunOwnershipConfig
from deerflow.runtime.checkpoint_state import CheckpointStateAccessor
from deerflow.runtime.context_keys import CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY
from deerflow.runtime.events.store.memory import MemoryRunEventStore
from deerflow.runtime.runs.manager import CancelOutcome, ConflictError, RunManager
from deerflow.runtime.runs.schemas import RunStatus
from deerflow.runtime.runs.store.memory import MemoryRunStore
from deerflow.runtime.runs.worker import (
RollbackPoint,
RunContext,
@ -87,6 +90,74 @@ async def test_pending_cancel_stops_waiting_for_prior_finalization():
bridge.publish_end.assert_awaited_once_with(record.run_id)
@pytest.mark.anyio
async def test_remote_cancel_wins_when_graph_finishes_before_owner_heartbeat():
store = MemoryRunStore()
ownership = RunOwnershipConfig(
heartbeat_enabled=True,
lease_seconds=30,
grace_seconds=10,
)
owner = RunManager(
store=store,
worker_id="worker-a",
run_ownership_config=ownership,
)
peer = RunManager(
store=store,
worker_id="worker-b",
run_ownership_config=ownership,
)
record = await owner.create_or_reject("thread-cancel-race")
bridge = SimpleNamespace(
publish=AsyncMock(),
publish_end=AsyncMock(),
cleanup=AsyncMock(),
)
started = asyncio.Event()
release = asyncio.Event()
class _FinishingAgent:
async def astream(
self,
graph_input,
config=None,
stream_mode=None,
subgraphs=False,
):
del graph_input, config, stream_mode, subgraphs
started.set()
await release.wait()
yield {"messages": []}
task = asyncio.create_task(
run_agent(
bridge,
owner,
record,
ctx=RunContext(
checkpointer=None,
event_store=MemoryRunEventStore(),
),
agent_factory=lambda **_kwargs: _FinishingAgent(),
graph_input={},
config={},
)
)
record.task = task
await asyncio.wait_for(started.wait(), timeout=1)
assert await peer.cancel(record.run_id, action="rollback") == CancelOutcome.requested
release.set()
await asyncio.wait_for(task, timeout=1)
stored = await store.get(record.run_id)
assert stored is not None
assert stored["status"] == "error"
assert stored["error"] == "Rolled back by user"
assert record.status == RunStatus.error
def _make_rollback_point(*, checkpoint_id="ckpt-1", messages=("before",), pending_writes=()):
materialized_messages = tuple(messages)
return RollbackPoint(

View File

@ -57,6 +57,10 @@ class _FakeRunManager:
async def set_status(self, *_args, **_kwargs) -> None:
return None
async def set_status_if_not_cancelled(self, *_args, **_kwargs) -> None:
await self.set_status(*_args, **_kwargs)
return None
async def update_model_name(self, *_args, **_kwargs) -> None:
return None

View File

@ -326,6 +326,10 @@ class _IntegrationRunManager:
async def set_status(self, _run_id, status, **_kwargs):
self._record.status = status
async def set_status_if_not_cancelled(self, _run_id, status, **kwargs):
await self.set_status(_run_id, status, **kwargs)
return None
async def update_model_name(self, *_args, **_kwargs):
return None