mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-07-28 17:06:05 +00:00
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:
parent
b1984cf4ab
commit
8a78c264b7
@ -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.
|
||||
|
||||
@ -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)
|
||||
|
||||
|
||||
@ -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),
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -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")
|
||||
@ -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))
|
||||
|
||||
@ -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,
|
||||
|
||||
@ -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"
|
||||
|
||||
@ -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"]
|
||||
|
||||
@ -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,
|
||||
|
||||
@ -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,
|
||||
}
|
||||
|
||||
@ -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)
|
||||
|
||||
|
||||
@ -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",
|
||||
|
||||
@ -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
|
||||
|
||||
|
||||
@ -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.
|
||||
|
||||
@ -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.
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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
|
||||
|
||||
|
||||
|
||||
@ -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:
|
||||
|
||||
@ -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()
|
||||
|
||||
@ -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."""
|
||||
|
||||
@ -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(
|
||||
|
||||
@ -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
|
||||
|
||||
|
||||
@ -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
|
||||
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user