From b53c1ae0e0589c07bdd4ce99c146d0b67a0b4a83 Mon Sep 17 00:00:00 2001 From: heart-scalpel Date: Tue, 14 Jul 2026 07:37:59 +0800 Subject: [PATCH] fix(runs): cancel degrades to lease takeover for multi-worker (#4064) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(runs): cancel degrades to lease takeover for multi-worker Work item 4 of the multi-worker ownership epic (https://github.com/bytedance/deer-flow/issues/3948). Problem: POST /runs/{run_id}/cancel landing on a non-owning worker returns 409 — the cancel button silently fails under GATEWAY_WORKERS>1 with no sticky routing. cancel() required the current worker to hold the in-memory task/abort_event, which any non-owner pod cannot satisfy. Changes: - RunManager.cancel() returns CancelOutcome enum (cancelled / taken_over / lease_valid_elsewhere / not_active_locally / not_cancellable / unknown) instead of bool, so the router can map each outcome to the right HTTP response. - New store primitive claim_for_takeover(): a single atomic conditional UPDATE that marks a run as error only when status IN (pending, running) AND (lease IS NULL OR lease < now - grace). Closes the stale-read / concurrent-heartbeat race — if the owner renews between our read and write, the UPDATE matches 0 rows and we surface lease_valid_elsewhere. - HTTP cancel + stream-join endpoints route on CancelOutcome: cancelled -> 202 (or 204 with wait=true); taken_over -> 202 immediately (no SSE streaming — the run is terminal on another worker, streaming would hang); lease_valid_elsewhere -> 409 + Retry-After header computed from lease_expires_at + grace_seconds. - RunManager.grace_seconds exposed as a public property; the router no longer reaches into _run_ownership_config. - _is_lease_expired extracted to a module-level function, shared by RunManager.cancel() and MemoryRunStore.claim_for_takeover(). - GATEWAY_WORKERS=1 + heartbeat_enabled=false is zero-regression: the non-local path short-circuits to not_active_locally, preserving the original 409 behaviour the existing tests pin. Tests: 12 new (5 store primitive + 4 cancel-takeover unit + 3 HTTP including a regression guard verifying POST /stream?action=interrupt on a dead-owner run returns 202 instead of hanging on SSE). 244 directly-related tests pass; 36/36 blocking-IO gate pass. Co-Authored-By: Claude Opus 4.7 * fix(runs): guard update_status and self-terminate on takeover Two defenses close a split-brain window where the original owner could overwrite a peer's takeover status: - update_status (SQL + memory store) now guards on status IN ('pending','running'). When takeover already set the row to 'error', the owner's final status write matches 0 rows and is dropped. - _persist_status: when update_status returns False, check whether the row exists before attempting recovery via put(). If the row exists (takeover by another worker), skip recovery instead of blindly upserting over the takeover. - Heartbeat _renew_leases: when update_lease returns False (row no longer pending/running or owner changed), cancel the local task so wasted CPU is bounded to the next heartbeat tick (~10s) instead of the full task lifetime. Also fix three reviewer feedback items: - Re-fetch the store row when cancel() returns lease_valid_elsewhere, so Retry-After uses the owner's freshly-renewed lease instead of a stale value from request start. - Fallback 'unknown' in takeover error message when owner_worker_id is NULL (pre-ownership data). - Remove dead else-10 branch from grace_seconds property (unreachable — all callers are downstream of the heartbeat_enabled guard). Co-Authored-By: Claude Opus 4.7 * test(runs): pin split-brain defences from update_status guard + heartbeat Three tests lock down the takeover authoritativeness so a late-running owner cannot overwrite a peer's claim: - update_status must reject writes when the store row is already terminal (taken over by another worker). - _persist_status must skip row-recovery via put() when the row exists but has been taken over. - Heartbeat _renew_leases must cancel the local task when update_lease returns False (row claimed by another worker). Co-Authored-By: Claude Opus 4.7 * fix(runs): precise outcome + log when local cancel loses to peer takeover Two reviewer precision nits on the split-brain defence: - _persist_status: branch the skip-reason log on existing["status"]. error → WARNING "peer takeover" (anomalous); interrupted/success → INFO "local cancel/completion race" (expected when user hits stop as the run finishes). Stops noisy false-positive takeover warnings in operator logs. - cancel() local path: when _persist_status returns False, re-check the store. If a peer's claim_for_takeover flipped the row to error between our in-memory cancel and the guarded update_status, surface taken_over instead of cancelled so the client sees a status consistent with the store. Test: test_cancel_returns_taken_over_when_peer_claims_during_local_cancel pins the race outcome. Co-Authored-By: Claude Opus 4.7 * fix(runs): widen update_status guard, de-duplicate lease helpers, add coverage Round 3 of reviewer feedback: - Widen update_status guard to status IN ('pending','running','interrupted'). The original guard blocked interrupted→error (the rollback finalize path), losing the "Rolled back by user" message. interrupted is now permitted while error/success stay locked — takeover protection unchanged. - claim_for_takeover False now re-reads the store row to distinguish causes: owner renewed lease → lease_valid_elsewhere; row went terminal → not_cancellable; another worker already took it over → taken_over. - Extract _raise_lease_valid_elsewhere() helper to de-duplicate the 409+Retry-After block shared across cancel_run and stream_existing_run. - Extract _lease_expired_or_null() in persistence/run/sql.py to de-duplicate the lease-expiry SQL WHERE clause shared by claim_for_takeover and list_inflight_with_expired_lease. - 11 new tests: 5 SQL-layer claim_for_takeover (expired/valid/NULL/ terminal/nonexistent), 3 _compute_retry_after unit (NULL/unparseable/ normal), 2 claim re-read precision (terminal/takeover), 1 stream endpoint 409+Retry-After. Not addressed (non-blocking, reviewer agreed): - The 2–3 store.gets in the takeover cold path: optimizing the API to accept a pre-fetched record would couple the router to the manager more tightly than justified by the perf gain. - The lease-expiry inline loop in MemoryRunStore.list_inflight_with_- expired_lease pre-computes cutoff once for all rows; switching to the shared _is_lease_expired helper would recompute datetime.now() per row with no real benefit. 260 related tests pass; 36/36 blocking-IO gate pass; ruff clean. Co-Authored-By: Claude Opus 4.7 * fix(runs): de-duplicate lease-expiry helper, restore defensive fallback Address final round of review feedback: - Extract is_lease_expired to deerflow.utils.time (no _ prefix, public utility). Manager and MemoryRunStore now import from the same place instead of the store reaching backward into the manager for a private function. - Restore defensive else-10 fallback in grace_seconds property (removed in an earlier round). The guard is unreachable for current callers but protects future ones from AttributeError. - Comment the transient in-memory interrupted vs store error state when a local cancel is superseded by a peer takeover. - Comment the max(1, ...) floor in _compute_retry_after — the floor is a lower bound, not a poll interval; clients should apply jitter. Co-Authored-By: Claude Opus 4.7 --------- Co-authored-by: Claude Opus 4.7 Co-authored-by: rayhpeng --- backend/AGENTS.md | 4 +- backend/app/gateway/routers/thread_runs.py | 93 ++- .../harness/deerflow/persistence/run/sql.py | 38 +- .../harness/deerflow/runtime/__init__.py | 3 +- .../harness/deerflow/runtime/runs/__init__.py | 3 +- .../harness/deerflow/runtime/runs/manager.py | 223 ++++++- .../deerflow/runtime/runs/store/base.py | 22 + .../deerflow/runtime/runs/store/memory.py | 41 +- .../packages/harness/deerflow/utils/time.py | 24 +- backend/tests/test_cancel_run_idempotent.py | 20 +- .../tests/test_multi_worker_run_ownership.py | 574 +++++++++++++++++- backend/tests/test_run_manager.py | 10 +- backend/tests/test_run_repository.py | 70 ++- 13 files changed, 1046 insertions(+), 79 deletions(-) diff --git a/backend/AGENTS.md b/backend/AGENTS.md index e6247ec62..b078d2b5b 100644 --- a/backend/AGENTS.md +++ b/backend/AGENTS.md @@ -341,8 +341,8 @@ metadata only. **RunManager / RunStore contract**: - `RunManager.get()` is async; direct callers must `await` it. - 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. -- `cancel()` and `create_or_reject(..., multitask_strategy="interrupt"|"rollback")` persist interrupted status through `RunStore.update_status()`, matching normal `set_status()` transitions. -- Store-only hydrated runs are readable history. If the current worker has no in-memory task/control state for that run, cancellation APIs can return 409 because this worker cannot stop the task. +- `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. +- 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. - `POST /wait` (both thread-scoped and `/api/runs/wait`) drains the stream bridge via `wait_for_run_completion()` instead of bare `await record.task`, so it honours the run's `on_disconnect` setting and cancels the background run on real client disconnect rather than returning a stale checkpoint (issue #3265). - Redis `StreamBridge` keys use a rolling retained-buffer TTL (`stream_bridge.stream_ttl_seconds`, refreshed on `publish()` / `publish_end()`) as a leak safety net, not as a run timeout. Startup orphan recovery publishes `END_SENTINEL` and schedules stream cleanup for recovered runs; malformed `Last-Event-ID` reconnect values live-tail new Redis events rather than replaying the retained buffer. Do not broaden this into a shared-database multi-pod reaper without adding worker ownership/liveness first. - Thread-scoped run creation accepts `checkpoint` / `checkpoint_id`; Gateway validates the checkpoint belongs to the request thread before writing `checkpoint_id` / `checkpoint_ns` into `config.configurable` for LangGraph branching. diff --git a/backend/app/gateway/routers/thread_runs.py b/backend/app/gateway/routers/thread_runs.py index 41282bac9..4cbdb17d3 100644 --- a/backend/app/gateway/routers/thread_runs.py +++ b/backend/app/gateway/routers/thread_runs.py @@ -13,6 +13,7 @@ from __future__ import annotations import asyncio import logging +from datetime import UTC, datetime from typing import Any, Literal from fastapi import APIRouter, HTTPException, Query, Request @@ -24,7 +25,7 @@ from app.gateway.authz import require_permission from app.gateway.deps import get_checkpointer, get_current_user, get_feedback_repo, get_run_event_store, get_run_manager, get_run_store, get_stream_bridge from app.gateway.pagination import trim_run_message_page from app.gateway.services import sse_consumer, start_run, wait_for_run_completion -from deerflow.runtime import RunRecord, RunStatus, serialize_channel_values_for_api +from deerflow.runtime import CancelOutcome, RunRecord, RunStatus, serialize_channel_values_for_api from deerflow.utils.messages import ORIGINAL_USER_CONTENT_KEY, get_original_user_content_text, message_to_text from deerflow.workspace_changes import get_workspace_changes_response @@ -145,6 +146,54 @@ def _cancel_conflict_detail(run_id: str, record: RunRecord) -> str: return f"Run {run_id} is not cancellable (status: {record.status.value})" +def _compute_retry_after(lease_expires_at: str | None, grace_seconds: int) -> int | None: + """Return seconds until the lease expires + grace, for ``Retry-After``. + + Returns ``None`` when the lease is NULL or unparseable so the caller + can decide whether to send a generic 409 without the header. + + The ``max(1, ...)`` floor means a lease just about to expire yields + ``Retry-After: 1``. This is a lower bound, not a recommended poll + interval — clients that honour this header should apply minimum + backoff / jitter rather than retrying every second. + """ + if lease_expires_at is None: + return None + try: + dt = datetime.fromisoformat(lease_expires_at) + if dt.tzinfo is None: + dt = dt.replace(tzinfo=UTC) + except (ValueError, TypeError): + return None + remaining = (dt - datetime.now(UTC)).total_seconds() + grace_seconds + return max(1, int(remaining)) + + +async def _raise_lease_valid_elsewhere( + run_id: str, + run_mgr, # RunManager (avoid import for testability) + record: RunRecord, +) -> None: + """Re-fetch the lease and raise HTTP 409 + Retry-After. + + ``record.lease_expires_at`` may be stale (fetched at request start while + the owner renewed between our read and the conditional UPDATE). Re-read + from the store to get the fresh value so ``Retry-After`` is accurate. + """ + fresh = await run_mgr.get(run_id) + if fresh is not None: + record = fresh + retry_after = _compute_retry_after(record.lease_expires_at, run_mgr.grace_seconds) + headers: dict[str, str] = {} + if retry_after is not None: + headers["Retry-After"] = str(retry_after) + raise HTTPException( + status_code=409, + detail=f"Run {run_id} is active on another worker; retry after lease expiry.", + headers=headers, + ) + + def _record_to_response(record: RunRecord) -> RunResponse: return RunResponse( run_id=record.run_id, @@ -512,24 +561,34 @@ async def cancel_run( - action=rollback: Stop execution, revert to pre-run checkpoint state - wait=true: Block until the run fully stops, return 204 - 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. """ run_mgr = get_run_manager(request) record = await run_mgr.get(run_id) if record is None or record.thread_id != thread_id: raise HTTPException(status_code=404, detail=f"Run {run_id} not found") - cancelled = await run_mgr.cancel(run_id, action=action) - if not cancelled: - raise HTTPException(status_code=409, detail=_cancel_conflict_detail(run_id, record)) + outcome = await run_mgr.cancel(run_id, action=action) - if wait and record.task is not None: - try: - await record.task - except asyncio.CancelledError: - pass - return Response(status_code=204) + # Success paths — the run was either cancelled locally or taken over + # from a dead worker. + if outcome in (CancelOutcome.cancelled, CancelOutcome.taken_over): + if wait and record.task is not None: + try: + await record.task + except asyncio.CancelledError: + pass + return Response(status_code=204) + return Response(status_code=202) - return Response(status_code=202) + if outcome == CancelOutcome.lease_valid_elsewhere: + await _raise_lease_valid_elsewhere(run_id, run_mgr, record) + + # not_cancellable, not_active_locally, unknown + raise HTTPException(status_code=409, detail=_cancel_conflict_detail(run_id, record)) @router.get("/{thread_id}/runs/{run_id}/join") @@ -586,8 +645,16 @@ async def stream_existing_run( # Cancel if an action was requested (stop-button / interrupt flow) if action is not None: - cancelled = await run_mgr.cancel(run_id, action=action) - if not cancelled: + outcome = await run_mgr.cancel(run_id, action=action) + if outcome == CancelOutcome.taken_over: + # The run was on another worker and is now marked ``error`` in the + # store. There is no local stream to drain — return immediately so + # 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 == 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 wait and record.task is not None: try: diff --git a/backend/packages/harness/deerflow/persistence/run/sql.py b/backend/packages/harness/deerflow/persistence/run/sql.py index 0f1bcf66e..3cb41aef1 100644 --- a/backend/packages/harness/deerflow/persistence/run/sql.py +++ b/backend/packages/harness/deerflow/persistence/run/sql.py @@ -20,6 +20,11 @@ from deerflow.runtime.user_context import AUTO, _AutoSentinel, resolve_user_id from deerflow.utils.time import coerce_iso +def _lease_expired_or_null(lease_col, cutoff: datetime): + """SQLAlchemy filter: True when the lease is NULL or has expired past *cutoff*.""" + return or_(lease_col.is_(None), lease_col < cutoff) + + class RunRepository(RunStore): def __init__(self, session_factory: async_sessionmaker[AsyncSession]) -> None: self._sf = session_factory @@ -165,8 +170,13 @@ class RunRepository(RunStore): values: dict[str, Any] = {"status": status, "updated_at": datetime.now(UTC)} if error is not None: values["error"] = error + # Guard: only transition rows that are still active. ``interrupted`` is + # included because the rollback path goes ``running → interrupted`` + # (cancel acknowledged) then ``interrupted → error`` (task finalize). + # ``error`` and ``success`` remain locked so a peer's takeover (or a + # completed run) cannot be overwritten by a late writer. async with self._sf() as session: - result = await session.execute(update(RunRow).where(RunRow.run_id == run_id).values(**values)) + result = await session.execute(update(RunRow).where(RunRow.run_id == run_id, RunRow.status.in_(("pending", "running", "interrupted"))).values(**values)) await session.commit() return result.rowcount != 0 @@ -404,6 +414,27 @@ class RunRepository(RunStore): await session.commit() return result.rowcount != 0 + async def claim_for_takeover( + self, + run_id: str, + *, + grace_seconds: int, + error: str, + ) -> bool: + cutoff = datetime.now(UTC) - timedelta(seconds=grace_seconds) + async with self._sf() as session: + result = await session.execute( + update(RunRow) + .where( + RunRow.run_id == run_id, + RunRow.status.in_(("pending", "running")), + _lease_expired_or_null(RunRow.lease_expires_at, cutoff), + ) + .values(status="error", error=error, updated_at=datetime.now(UTC)) + ) + await session.commit() + return result.rowcount != 0 + async def list_inflight_with_expired_lease( self, *, @@ -422,10 +453,7 @@ class RunRepository(RunStore): .where( RunRow.status.in_(("pending", "running")), RunRow.created_at <= before_dt, - or_( - RunRow.lease_expires_at.is_(None), - RunRow.lease_expires_at < cutoff, - ), + _lease_expired_or_null(RunRow.lease_expires_at, cutoff), ) .order_by(RunRow.created_at.asc()) ) diff --git a/backend/packages/harness/deerflow/runtime/__init__.py b/backend/packages/harness/deerflow/runtime/__init__.py index 849581005..d87090748 100644 --- a/backend/packages/harness/deerflow/runtime/__init__.py +++ b/backend/packages/harness/deerflow/runtime/__init__.py @@ -6,7 +6,7 @@ directly from ``deerflow.runtime``. """ from .checkpointer import checkpointer_context, get_checkpointer, make_checkpointer, reset_checkpointer -from .runs import ConflictError, DisconnectMode, RunContext, RunManager, RunRecord, RunStatus, UnsupportedStrategyError, run_agent +from .runs import CancelOutcome, ConflictError, DisconnectMode, RunContext, RunManager, RunRecord, RunStatus, UnsupportedStrategyError, run_agent from .serialization import serialize, serialize_channel_values, serialize_channel_values_for_api, serialize_lc_object, serialize_messages_tuple, strip_data_url_image_blocks from .store import get_store, make_store, reset_store, store_context @@ -22,6 +22,7 @@ __all__ = [ "make_checkpointer", "reset_checkpointer", # runs + "CancelOutcome", "ConflictError", "DisconnectMode", "RunContext", diff --git a/backend/packages/harness/deerflow/runtime/runs/__init__.py b/backend/packages/harness/deerflow/runtime/runs/__init__.py index 9faa30c17..5d2a2661f 100644 --- a/backend/packages/harness/deerflow/runtime/runs/__init__.py +++ b/backend/packages/harness/deerflow/runtime/runs/__init__.py @@ -1,10 +1,11 @@ """Run lifecycle management for LangGraph Platform API compatibility.""" -from .manager import ConflictError, RunManager, RunRecord, UnsupportedStrategyError +from .manager import CancelOutcome, ConflictError, RunManager, RunRecord, UnsupportedStrategyError from .schemas import DisconnectMode, RunStatus from .worker import RunContext, run_agent __all__ = [ + "CancelOutcome", "ConflictError", "DisconnectMode", "RunContext", diff --git a/backend/packages/harness/deerflow/runtime/runs/manager.py b/backend/packages/harness/deerflow/runtime/runs/manager.py index 1715cfda2..5571001ed 100644 --- a/backend/packages/harness/deerflow/runtime/runs/manager.py +++ b/backend/packages/harness/deerflow/runtime/runs/manager.py @@ -10,10 +10,12 @@ import uuid from collections.abc import Awaitable, Callable from dataclasses import dataclass, field from datetime import UTC, datetime, timedelta +from enum import StrEnum from typing import TYPE_CHECKING, Any from sqlalchemy.exc import IntegrityError as SAIntegrityError +from deerflow.utils.time import is_lease_expired from deerflow.utils.time import now_iso as _now_iso from .schemas import DisconnectMode, RunStatus @@ -340,6 +342,29 @@ class RunManager: lambda: self._store.update_status(record.run_id, status.value, error=error), ) if updated is False: + # ``update_status`` is now guarded by ``status IN ('pending','running')``. + # False can mean either: + # (a) the row was never persisted (initial ``put()`` failed) → recreate. + # (b) the row is terminal — either a peer takeover (``error``) + # or a local cancel/completion race (``interrupted`` / + # ``success``). The log severity branches on which. + existing = await self._store.get(record.run_id) + if existing is not None: + existing_status = existing.get("status") + if existing_status == "error": + logger.warning( + "Run %s status update to %s skipped: store row already at error (peer takeover)", + record.run_id, + status.value, + ) + else: + logger.info( + "Run %s status update to %s skipped: store row already at %s (local cancel/completion race)", + record.run_id, + status.value, + existing_status, + ) + return False return await self._persist_snapshot_to_store(record.run_id, row_recovery_payload) return True except Exception: @@ -662,41 +687,151 @@ class RunManager: await self._persist_model_name(run_id, model_name) logger.info("Run %s model_name=%s", run_id, model_name) - async def cancel(self, run_id: str, *, action: str = "interrupt") -> bool: + async def cancel(self, run_id: str, *, action: str = "interrupt") -> CancelOutcome: """Request cancellation of a run. + When the call lands on the owning worker the run is cancelled + locally as before (in-memory abort + status persisted to store). + + When the call lands on a non-owning worker in a multi-worker + deployment with heartbeat enabled: + + - **Lease expired** — the run's lease has passed the grace + threshold, so this worker takes ownership and marks it as + ``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. + + In single-worker mode (``heartbeat_enabled=False``) store-only + hydrated runs that aren't in-memory return ``not_active_locally``, + preserving the original 409 behaviour. + Args: run_id: The run ID to cancel. - action: "interrupt" keeps checkpoint, "rollback" reverts to pre-run state. + action: ``"interrupt"`` keeps checkpoint, ``"rollback"`` + reverts to pre-run state. - Sets the abort event with the action reason and cancels the asyncio task. - Returns ``True`` if cancellation was initiated **or** the run was already - interrupted (idempotent — a second cancel is a no-op success). - Returns ``False`` only when the run is unknown to this worker or has - reached a terminal state other than interrupted (completed, failed, etc.). + Returns: + A :class:`CancelOutcome` enum describing what happened. """ + # ------------------------------------------------------------------ + # Local path — this worker owns the run in-memory. + # ------------------------------------------------------------------ async with self._lock: record = self._runs.get(run_id) - if record is None: - return False - if record.status == RunStatus.interrupted: - return True # idempotent — already cancelled on this worker - if record.status not in (RunStatus.pending, RunStatus.running): - return False - 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: - record.task.cancel() - record.status = RunStatus.interrupted - record.updated_at = _now_iso() - await self._persist_status(record, RunStatus.interrupted) - logger.info("Run %s cancelled (action=%s)", run_id, action) - return True + if record is not None: + if record.status == RunStatus.interrupted: + return CancelOutcome.cancelled # idempotent + if record.status not in (RunStatus.pending, RunStatus.running): + return CancelOutcome.not_cancellable + 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: + record.task.cancel() + record.status = RunStatus.interrupted + record.updated_at = _now_iso() + + # Persist outside the lock so store calls don't block other mutations. + if record is not None: + persisted = await self._persist_status(record, RunStatus.interrupted) + if not persisted and self._store is not None: + # ``_persist_status`` already fetched ``existing`` internally; + # re-check the store to see if a peer takeover flipped the + # row to ``error`` between our in-memory cancel and the + # guarded ``update_status``. If so, surface ``taken_over`` + # so the client sees a status consistent with the store. + try: + existing = await self._store.get(run_id) + except Exception: + existing = None + if existing is not None and existing.get("status") == "error": + # The in-memory ``record.status`` is still ``interrupted`` + # (set under the lock above) while the store row is now + # ``error``. This transient staleness is harmless: the + # ``_persist_status`` guard prevents the late finalisation + # write from overwriting the takeover, and the store is the + # authoritative source for subsequent reads. + logger.info("Run %s local cancel superseded by peer takeover", run_id) + return CancelOutcome.taken_over + logger.info("Run %s cancelled (action=%s)", run_id, action) + return CancelOutcome.cancelled + + # ------------------------------------------------------------------ + # Non-local path — no in-memory record, must consult the store. + # ------------------------------------------------------------------ + + if not self.heartbeat_enabled: + return CancelOutcome.not_active_locally + + if self._store is None: + return CancelOutcome.unknown + + try: + row = await self._store.get(run_id) + except Exception: + logger.warning("Failed to fetch run %s from store during cancel", run_id, exc_info=True) + return CancelOutcome.unknown + + if row is None: + return CancelOutcome.unknown + + store_status = row.get("status") + if store_status not in ("pending", "running"): + return CancelOutcome.not_cancellable + + grace_seconds = self.grace_seconds + 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 + + 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: + taken = await self._call_store_with_retry( + "claim_for_takeover", + run_id, + lambda: self._store.claim_for_takeover( + run_id, + grace_seconds=grace_seconds, + error=take_over_msg, + ), + ) + except Exception: + logger.warning("Take-over claim for run %s failed with exception", run_id, exc_info=True) + return CancelOutcome.unknown + + if taken: + logger.warning("Run %s taken over by worker %s (action=%s)", run_id, self._worker_id, action) + return CancelOutcome.taken_over + + # The conditional UPDATE matched 0 rows. Two causes: + # (a) the owner renewed the lease → lease_valid_elsewhere. + # (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. + # Re-read to distinguish. + try: + fresh = await self._store.get(run_id) + except Exception: + fresh = None + if fresh is None: + return CancelOutcome.unknown + fresh_status = fresh.get("status") + if fresh_status not in ("pending", "running"): + if fresh_status == "error": + 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 def _compute_lease_expires_at(self) -> str | None: - """Compute the lease expiration timestamp for a new run. + """Return the lease expiry ISO timestamp for a freshly created run. Returns ``None`` when heartbeat is disabled (single-worker mode) so reconciliation treats crashed runs as orphans (NULL lease) and @@ -967,6 +1102,17 @@ class RunManager: return False return self._run_ownership_config.heartbeat_enabled + @property + def grace_seconds(self) -> int: + """Return the configured grace seconds. + + All current callers are downstream of ``heartbeat_enabled``, which + is False whenever ``_run_ownership_config`` is None. The fallback + matches the Pydantic model default and is defensive against future + callers that might reach this property without that guard. + """ + return self._run_ownership_config.grace_seconds if self._run_ownership_config else 10 + async def start_heartbeat(self) -> None: """Start the background lease-renewal task. @@ -1084,6 +1230,22 @@ class RunManager: # fields). Re-acquiring ``self._lock`` here would # serialise against unrelated run mutations for no gain. record.lease_expires_at = new_expiry + else: + # ``update_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 + # finalisation. + logger.warning( + "Run %s lease renewal failed (status=%s,owner=%s) – worker likely taken over; aborting local task", + run_id, + record.status.value, + record.owner_worker_id, + ) + record.abort_event.set() + task_active = record.task is not None and not record.task.done() + if task_active: + record.task.cancel() except Exception: logger.warning("Failed to renew lease for run %s", run_id, exc_info=True) @@ -1198,6 +1360,17 @@ class RunManager: logger.info("Drained %d in-flight run(s) on shutdown (%d settled within %.1fs)", len(inflight), len(inflight) - len(pending), timeout) +class CancelOutcome(StrEnum): + """Result of a :meth:`RunManager.cancel` call.""" + + cancelled = "cancelled" + taken_over = "taken_over" + lease_valid_elsewhere = "lease_valid_elsewhere" + not_cancellable = "not_cancellable" + not_active_locally = "not_active_locally" + unknown = "unknown" + + class ConflictError(Exception): """Raised when multitask_strategy=reject and thread has inflight runs.""" diff --git a/backend/packages/harness/deerflow/runtime/runs/store/base.py b/backend/packages/harness/deerflow/runtime/runs/store/base.py index 1ccc1a8cc..0089640b2 100644 --- a/backend/packages/harness/deerflow/runtime/runs/store/base.py +++ b/backend/packages/harness/deerflow/runtime/runs/store/base.py @@ -156,6 +156,28 @@ class RunStore(abc.ABC): """Renew the lease on an active run. Returns ``False`` when no row matched.""" pass + @abc.abstractmethod + async def claim_for_takeover( + self, + run_id: str, + *, + grace_seconds: int, + error: str, + ) -> bool: + """Atomically mark an expired-lease active run as ``error``. + + Only rows whose lease has expired past *grace_seconds* (or whose + lease is NULL — pre-ownership data) are updated. The conditional + WHERE closes the race between the caller's stale read of the lease + and a concurrent heartbeat renewal by the owning worker. + + Returns ``False`` when: + - the run is no longer ``pending`` / ``running``, + - the lease is still valid (owner heartbeat is alive), or + - the row doesn't exist. + """ + pass + @abc.abstractmethod async def list_inflight_with_expired_lease( self, diff --git a/backend/packages/harness/deerflow/runtime/runs/store/memory.py b/backend/packages/harness/deerflow/runtime/runs/store/memory.py index a2b98fe86..d597db922 100644 --- a/backend/packages/harness/deerflow/runtime/runs/store/memory.py +++ b/backend/packages/harness/deerflow/runtime/runs/store/memory.py @@ -88,13 +88,18 @@ class MemoryRunStore(RunStore): return results[:limit] async def update_status(self, run_id, status, *, error=None): - if run_id in self._runs: - self._runs[run_id]["status"] = status - if error is not None: - self._runs[run_id]["error"] = error - self._runs[run_id]["updated_at"] = datetime.now(UTC).isoformat() - return True - return False + run = self._runs.get(run_id) + if run is None: + return False + # Guard: only transition rows that are still active. ``interrupted`` + # is included for the rollback path (``interrupted → error`` finalize). + if run["status"] not in ("pending", "running", "interrupted"): + return False + run["status"] = status + if error is not None: + run["error"] = error + run["updated_at"] = datetime.now(UTC).isoformat() + return True async def update_model_name(self, run_id, model_name): if run_id in self._runs: @@ -194,6 +199,28 @@ class MemoryRunStore(RunStore): run["updated_at"] = datetime.now(UTC).isoformat() return True + async def claim_for_takeover( + self, + run_id: str, + *, + grace_seconds: int, + error: str, + ) -> bool: + from deerflow.utils.time import is_lease_expired + + run = self._runs.get(run_id) + if run is None: + return False + if run["status"] not in ("pending", "running"): + return False + lease = run.get("lease_expires_at") + if not is_lease_expired(lease, grace_seconds=grace_seconds): + return False + run["status"] = "error" + run["error"] = error + run["updated_at"] = datetime.now(UTC).isoformat() + return True + async def list_inflight_with_expired_lease( self, *, diff --git a/backend/packages/harness/deerflow/utils/time.py b/backend/packages/harness/deerflow/utils/time.py index 307a4b6b0..3a4eb079b 100644 --- a/backend/packages/harness/deerflow/utils/time.py +++ b/backend/packages/harness/deerflow/utils/time.py @@ -15,9 +15,29 @@ records that historically stored ``str(time.time())`` floats. from __future__ import annotations import re -from datetime import UTC, datetime +from datetime import UTC, datetime, timedelta + +__all__ = ["coerce_iso", "is_lease_expired", "now_iso"] + + +def is_lease_expired(lease_expires_at: str | None, *, grace_seconds: int) -> bool: + """Return ``True`` when *lease_expires_at* has elapsed past grace. + + A NULL lease (pre-ownership data) is always considered expired so + take-over (cancel from a non-owning worker) can reclaim it in the + same way reconciliation does. Unparseable timestamps are also + treated as expired (defence in depth). + """ + if lease_expires_at is None: + return True + try: + dt = datetime.fromisoformat(lease_expires_at) + if dt.tzinfo is None: + dt = dt.replace(tzinfo=UTC) + except (ValueError, TypeError): + return True + return dt < datetime.now(UTC) - timedelta(seconds=grace_seconds) -__all__ = ["coerce_iso", "now_iso"] _UNIX_TIMESTAMP_PATTERN = re.compile(r"^\d{10}(?:\.\d+)?$") """Matches the unix-timestamp string shape historically written by diff --git a/backend/tests/test_cancel_run_idempotent.py b/backend/tests/test_cancel_run_idempotent.py index 0bf2548d1..4a193b224 100644 --- a/backend/tests/test_cancel_run_idempotent.py +++ b/backend/tests/test_cancel_run_idempotent.py @@ -14,7 +14,7 @@ from _router_auth_helpers import make_authed_test_app from fastapi.testclient import TestClient from app.gateway.routers import thread_runs -from deerflow.runtime import RunManager, RunStatus +from deerflow.runtime import CancelOutcome, RunManager, RunStatus THREAD_ID = "thread-cancel-test" @@ -49,22 +49,22 @@ def _create_interrupted_run(mgr: RunManager) -> str: class TestRunManagerCancelIdempotency: - def test_cancel_returns_true_for_already_interrupted_run(self): - """cancel() must return True when the run is already interrupted.""" + def test_cancel_returns_cancelled_for_already_interrupted_run(self): + """cancel() must return CancelledOutcome.cancelled when the run is already interrupted.""" async def run(): mgr = RunManager() record = await mgr.create(THREAD_ID) await mgr.set_status(record.run_id, RunStatus.running) first = await mgr.cancel(record.run_id) - assert first is True + assert first == CancelOutcome.cancelled second = await mgr.cancel(record.run_id) - assert second is True # idempotent + assert second == CancelOutcome.cancelled # idempotent asyncio.run(run()) - def test_cancel_returns_false_for_successful_run(self): - """cancel() must still return False for runs that completed successfully.""" + def test_cancel_returns_not_cancellable_for_successful_run(self): + """cancel() must return not_cancellable for runs that completed successfully.""" async def run(): mgr = RunManager() @@ -72,15 +72,15 @@ class TestRunManagerCancelIdempotency: await mgr.set_status(record.run_id, RunStatus.running) await mgr.set_status(record.run_id, RunStatus.success) result = await mgr.cancel(record.run_id) - assert result is False + assert result == CancelOutcome.not_cancellable asyncio.run(run()) - def test_cancel_returns_false_for_unknown_run(self): + def test_cancel_returns_not_active_locally_for_unknown_run(self): async def run(): mgr = RunManager() result = await mgr.cancel("nonexistent-run-id") - assert result is False + assert result == CancelOutcome.not_active_locally asyncio.run(run()) diff --git a/backend/tests/test_multi_worker_run_ownership.py b/backend/tests/test_multi_worker_run_ownership.py index 1d16f182c..6f2e985b0 100644 --- a/backend/tests/test_multi_worker_run_ownership.py +++ b/backend/tests/test_multi_worker_run_ownership.py @@ -20,7 +20,7 @@ import pytest from deerflow.config.run_ownership_config import RunOwnershipConfig from deerflow.runtime import RunManager, RunStatus -from deerflow.runtime.runs.manager import ConflictError, _generate_worker_id +from deerflow.runtime.runs.manager import CancelOutcome, ConflictError, _generate_worker_id from deerflow.runtime.runs.store.memory import MemoryRunStore # --------------------------------------------------------------------------- @@ -467,30 +467,30 @@ async def test_cancel_local_run_succeeds(): await manager.set_status(record.run_id, RunStatus.running) result = await manager.cancel(record.run_id) - assert result is True + assert result == CancelOutcome.cancelled assert record.status == RunStatus.interrupted @pytest.mark.anyio async def test_cancel_unknown_run_returns_false(): - """Cancel must return False for a run not known to this worker.""" + """Cancel must return not_active_locally for a run not known to this worker (heartbeat off).""" store = MemoryRunStore() manager = _make_manager(store=store) result = await manager.cancel("nonexistent-run") - assert result is False + assert result == CancelOutcome.not_active_locally @pytest.mark.anyio async def test_cancel_idempotent(): - """Cancel must return True when the run is already interrupted.""" + """Cancel must return cancelled when the run is already interrupted.""" store = MemoryRunStore() manager = _make_manager(store=store) record = await manager.create("thread-1") await manager.set_status(record.run_id, RunStatus.interrupted) result = await manager.cancel(record.run_id) - assert result is True + assert result == CancelOutcome.cancelled # --------------------------------------------------------------------------- @@ -965,3 +965,565 @@ async def test_list_inflight_with_expired_lease_null_lease_always_reclaimed(): results = await store.list_inflight_with_expired_lease(grace_seconds=grace) result_ids = {r["run_id"] for r in results} assert "null-run" in result_ids + + +# --------------------------------------------------------------------------- +# claim_for_takeover — store primitive +# --------------------------------------------------------------------------- + + +@pytest.mark.anyio +async def test_claim_for_takeover_succeeds_with_expired_lease(): + """claim_for_takeover must succeed when the lease has passed the grace window.""" + store = MemoryRunStore() + grace = 10 + expired_lease = (datetime.now(UTC) - timedelta(seconds=grace + 5)).isoformat() + await store.put("run-1", thread_id="t1", status="running", created_at=datetime.now(UTC).isoformat(), owner_worker_id="w-a", lease_expires_at=expired_lease) + + ok = await store.claim_for_takeover("run-1", grace_seconds=grace, error="claimed") + assert ok is True + + row = await store.get("run-1") + assert row is not None + assert row["status"] == "error" + assert row["error"] == "claimed" + + +@pytest.mark.anyio +async def test_claim_for_takeover_fails_with_valid_lease(): + """claim_for_takeover must return False when the lease is still valid.""" + store = MemoryRunStore() + grace = 10 + valid_lease = (datetime.now(UTC) + timedelta(seconds=60)).isoformat() + await store.put("run-1", thread_id="t1", status="running", created_at=datetime.now(UTC).isoformat(), owner_worker_id="w-a", lease_expires_at=valid_lease) + + ok = await store.claim_for_takeover("run-1", grace_seconds=grace, error="claimed") + assert ok is False + + row = await store.get("run-1") + assert row is not None + assert row["status"] == "running" + + +@pytest.mark.anyio +async def test_claim_for_takeover_succeeds_with_null_lease(): + """NULL-lease rows (pre-ownership data) must be claimable.""" + store = MemoryRunStore() + await store.put("run-null", thread_id="t1", status="running", created_at=datetime.now(UTC).isoformat()) + + ok = await store.claim_for_takeover("run-null", grace_seconds=10, error="claimed") + assert ok is True + + row = await store.get("run-null") + assert row["status"] == "error" + + +@pytest.mark.anyio +async def test_claim_for_takeover_fails_on_terminal_status(): + """claim_for_takeover must return False for already-terminal runs.""" + store = MemoryRunStore() + await store.put("run-done", thread_id="t1", status="success", created_at=datetime.now(UTC).isoformat()) + + ok = await store.claim_for_takeover("run-done", grace_seconds=10, error="claimed") + assert ok is False + + +@pytest.mark.anyio +async def test_claim_for_takeover_fails_for_nonexistent_run(): + """claim_for_takeover must return False when the run doesn't exist.""" + store = MemoryRunStore() + ok = await store.claim_for_takeover("no-such-run", grace_seconds=10, error="claimed") + assert ok is False + + +# --------------------------------------------------------------------------- +# cancel() cross-worker takeover — work item 4 +# --------------------------------------------------------------------------- + + +@pytest.mark.anyio +async def test_cancel_takeover_from_crashed_worker(): + """cancel must take over (mark error) when lease is expired and owner is another worker.""" + store = MemoryRunStore() + grace = 10 + expired_lease = (datetime.now(UTC) - timedelta(seconds=grace + 5)).isoformat() + await store.put("run-expired", thread_id="t1", status="running", created_at=datetime.now(UTC).isoformat(), owner_worker_id="dead-worker", lease_expires_at=expired_lease) + + manager = _make_manager(store=store, run_ownership_config=_lease_config(heartbeat_enabled=True, grace_seconds=grace)) + outcome = await manager.cancel("run-expired") + assert outcome == CancelOutcome.taken_over + + row = await store.get("run-expired") + assert row is not None + assert row["status"] == "error" + + +@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.""" + store = MemoryRunStore() + grace = 10 + valid_lease = (datetime.now(UTC) + timedelta(seconds=60)).isoformat() + await store.put("run-alive", thread_id="t1", status="running", created_at=datetime.now(UTC).isoformat(), owner_worker_id="alive-worker", lease_expires_at=valid_lease) + + 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 + + row = await store.get("run-alive") + assert row is not None + assert row["status"] == "running" # untouched + + +@pytest.mark.anyio +async def test_cancel_returns_unknown_when_no_store(): + """cancel must return unknown when there's no store and the run is not in memory.""" + manager = _make_manager(run_ownership_config=_lease_config(heartbeat_enabled=True)) + outcome = await manager.cancel("no-such-run") + assert outcome == CancelOutcome.unknown + + +@pytest.mark.anyio +async def test_cancel_returns_not_active_locally_when_heartbeat_disabled(): + """With heartbeat disabled, store-only runs must not be cancellable (old 409 path).""" + store = MemoryRunStore() + await store.put("store-only", thread_id="t1", status="running", created_at=datetime.now(UTC).isoformat()) + + manager = _make_manager(store=store, run_ownership_config=_lease_config(heartbeat_enabled=False)) + outcome = await manager.cancel("store-only") + assert outcome == CancelOutcome.not_active_locally + + +@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.""" + store = MemoryRunStore() + grace = 10 + expired_lease = (datetime.now(UTC) - timedelta(seconds=grace + 5)).isoformat() + await store.put("run-race", thread_id="t1", status="running", created_at=datetime.now(UTC).isoformat(), owner_worker_id="w-a", lease_expires_at=expired_lease) + + # Simulate the race: right before claim_for_takeover writes, another + # heartbeat renews the lease. We monkey-patch claim_for_takeover to + # simulate the lease having been renewed. + original = store.claim_for_takeover + + async def race_lost(run_id, *, grace_seconds, error): + # Simulate a heartbeat renewal between the read and the write + run = store._runs.get(run_id) + if run and run["status"] in ("pending", "running"): + run["lease_expires_at"] = (datetime.now(UTC) + timedelta(seconds=60)).isoformat() + return await original(run_id, grace_seconds=grace_seconds, error=error) + + store.claim_for_takeover = race_lost + 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 + + +@pytest.mark.anyio +async def test_cancel_takeover_respects_grace_seconds(): + """Cancel must not take over when the lease is within the grace window.""" + store = MemoryRunStore() + grace = 10 + # Lease expired, but only by 3s — still within the 10s grace window + just_expired = (datetime.now(UTC) - timedelta(seconds=3)).isoformat() + await store.put("run-grace", thread_id="t1", status="running", created_at=datetime.now(UTC).isoformat(), owner_worker_id="w-a", lease_expires_at=just_expired) + + 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 + + +@pytest.mark.anyio +async def test_cancel_not_cancellable_for_store_terminal_run(): + """cancel must return not_cancellable when the store run is already in a terminal state.""" + store = MemoryRunStore() + await store.put("run-done", thread_id="t1", status="success", created_at=datetime.now(UTC).isoformat()) + + manager = _make_manager(store=store, run_ownership_config=_lease_config(heartbeat_enabled=True)) + outcome = await manager.cancel("run-done") + assert outcome == CancelOutcome.not_cancellable + + +# --------------------------------------------------------------------------- +# HTTP-level — cancel endpoint cross-worker responses +# --------------------------------------------------------------------------- + + +def _make_cancel_test_app(mgr: RunManager): + """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 + + from app.gateway.routers import thread_runs + from deerflow.runtime import MemoryStreamBridge + + app = make_authed_test_app() + app.include_router(thread_runs.router) + app.state.run_manager = mgr + app.state.stream_bridge = 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.""" + store = MemoryRunStore() + grace = 10 + valid_lease = (datetime.now(UTC) + timedelta(seconds=60)).isoformat() + asyncio.run( + store.put( + "run-alive", + 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, grace_seconds=grace)) + 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 + + # Store row must be untouched + row = asyncio.run(store.get("run-alive")) + assert row["status"] == "running" + + +def test_http_cancel_non_owner_expired_lease_returns_202_takeover(): + """POST /cancel on a non-owning worker with an expired lease must return 202 (takeover).""" + store = MemoryRunStore() + grace = 10 + expired_lease = (datetime.now(UTC) - timedelta(seconds=grace + 30)).isoformat() + asyncio.run( + store.put( + "run-dead", + thread_id="t1", + status="running", + created_at=datetime.now(UTC).isoformat(), + owner_worker_id="dead-worker", + lease_expires_at=expired_lease, + ) + ) + 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-dead/cancel") + assert resp.status_code == 202 + + # Store row must be marked error + row = asyncio.run(store.get("run-dead")) + assert row["status"] == "error" + + +def test_http_stream_action_interrupt_takeover_returns_202_not_hang(): + """POST /stream?action=interrupt on a dead-owner run must return 202 immediately, not hang on SSE.""" + store = MemoryRunStore() + grace = 10 + expired_lease = (datetime.now(UTC) - timedelta(seconds=grace + 30)).isoformat() + asyncio.run( + store.put( + "run-dead-stream", + thread_id="t1", + status="running", + created_at=datetime.now(UTC).isoformat(), + owner_worker_id="dead-worker", + lease_expires_at=expired_lease, + ) + ) + mgr = _make_manager(store=store, run_ownership_config=_lease_config(heartbeat_enabled=True, grace_seconds=grace)) + client = _make_cancel_test_app(mgr) + + # This must NOT hang — the takeover path returns 202 before reaching StreamingResponse. + resp = client.post("/api/threads/t1/runs/run-dead-stream/stream", params={"action": "interrupt"}) + assert resp.status_code == 202 + + row = asyncio.run(store.get("run-dead-stream")) + assert row["status"] == "error" + + +# --------------------------------------------------------------------------- +# Split-brain defences — update_status guard + heartbeat self-termination +# --------------------------------------------------------------------------- + + +@pytest.mark.anyio +async def test_update_status_rejects_terminal_row(): + """update_status must return False when the store row is already terminal + (error/success), so a late writer cannot overwrite a peer's takeover or + a completed run. interrupted is NOT terminal — the rollback path needs + ``interrupted → error`` to finalize.""" + store = MemoryRunStore() + # error (takeover) must stay locked + await store.put("run-err", thread_id="t1", status="error", created_at=datetime.now(UTC).isoformat()) + assert await store.update_status("run-err", "success") is False + assert (await store.get("run-err"))["status"] == "error" + + # success must stay locked + await store.put("run-ok", thread_id="t1", status="success", created_at=datetime.now(UTC).isoformat()) + assert await store.update_status("run-ok", "error") is False + assert (await store.get("run-ok"))["status"] == "success" + + # interrupted → error MUST pass (rollback finalize path) + await store.put("run-rb", thread_id="t1", status="interrupted", created_at=datetime.now(UTC).isoformat()) + assert await store.update_status("run-rb", "error", error="Rolled back by user") is True + row = await store.get("run-rb") + assert row["status"] == "error" + assert row["error"] == "Rolled back by user" + + +@pytest.mark.anyio +async def test_persist_status_skips_recovery_when_row_taken_over(): + """_persist_status must not recreate a row that was taken over by another worker. + + When update_status returns False, the recovery path checks whether the + row still exists. A row that exists but is terminal (taken over) must + be left alone — calling put() would overwrite the takeover.""" + store = MemoryRunStore() + mgr = RunManager(store=store, run_ownership_config=_lease_config(heartbeat_enabled=True)) + + # Simulate: this worker created and started a run, but a peer took it over. + record = await mgr.create("thread-1") + await mgr.set_status(record.run_id, RunStatus.running) + # Peer takeover: directly flip the store row to error + await store.update_status(record.run_id, "error") + # Now simulate the original owner's task finishing and trying to write success + ok = await mgr._persist_status(record, RunStatus.success) + assert ok is False # skipped recovery, row already exists and is terminal + row = await store.get(record.run_id) + assert row["status"] == "error" # not overwritten + + +@pytest.mark.anyio +async def test_heartbeat_cancels_task_on_lease_loss(): + """Heartbeat must cancel the local asyncio task when update_lease returns False. + + If the store row was claimed by another worker (status no longer + pending/running, or owner changed), the heartbeat tick must abort the + local task so wasted CPU is bounded to ~10s instead of the full task + lifetime.""" + store = MemoryRunStore() + mgr = RunManager(store=store, run_ownership_config=_lease_config(heartbeat_enabled=True, lease_seconds=30)) + + # Create a run that this worker owns + record = await mgr.create("thread-1") + await mgr.set_status(record.run_id, RunStatus.running) + + # Spawn a dummy task so cancel has something to stop + loop = asyncio.get_running_loop() + record.task = loop.create_task(asyncio.sleep(3600)) + + # Simulate takeover: directly flip the store row to error + await store.update_status(record.run_id, "error") + + # Run a single heartbeat tick — it should see update_lease return False + # and cancel the task + await mgr._renew_leases() + + # Let the event loop process the cancellation (task.cancel() schedules, + # doesn't await). + await asyncio.sleep(0) + assert record.task.cancelled() + + +@pytest.mark.anyio +async def test_cancel_returns_taken_over_when_peer_claims_during_local_cancel(): + """When a peer's claim_for_takeover flips the row to error between this + worker's in-memory cancel and the guarded update_status, cancel() must + surface taken_over (not cancelled) so the client sees a status consistent + with the store.""" + store = MemoryRunStore() + mgr = RunManager(store=store, run_ownership_config=_lease_config(heartbeat_enabled=True)) + + record = await mgr.create("thread-1") + await mgr.set_status(record.run_id, RunStatus.running) + + # Wrap update_status so that the first call (from cancel's _persist_status) + # is rejected as if a peer already marked the row error. This simulates + # the race: in-memory cancel succeeds, but store write is blocked. + original = store.update_status + + async def race_update(run_id, status, *, error=None): + # Simulate peer takeover: flip to error before our write lands + run = store._runs.get(run_id) + if run and run["status"] == "running" and status == "interrupted": + run["status"] = "error" + run["error"] = "peer takeover" + run["updated_at"] = datetime.now(UTC).isoformat() + return False # our write was blocked + return await original(run_id, status, error=error) + + store.update_status = race_update + + outcome = await mgr.cancel(record.run_id) + assert outcome == CancelOutcome.taken_over + + # Store row must reflect the takeover, not the local cancel + row = await store.get(record.run_id) + assert row["status"] == "error" + + +@pytest.mark.anyio +async def test_cancel_action_rollback_finalizes_to_error_in_store(): + """action=rollback must end up as error in the store with the + "Rolled back by user" message preserved. + + Regression guard: the update_status guard was originally + ``status IN ('pending','running')`` which blocked the rollback path's + ``interrupted → error`` transition — the store stayed interrupted and + the rollback message was lost. + """ + store = MemoryRunStore() + mgr = RunManager(store=store, run_ownership_config=_lease_config(heartbeat_enabled=True)) + + record = await mgr.create("thread-1") + await mgr.set_status(record.run_id, RunStatus.running) + + # Step 1: cancel(action=rollback) flips running → interrupted + outcome = await mgr.cancel(record.run_id, action="rollback") + assert outcome == CancelOutcome.cancelled + row = await store.get(record.run_id) + assert row["status"] == "interrupted" + + # Step 2: worker.py finalize path — task raises CancelledError, then + # set_status(error, "Rolled back by user"). The widened guard + # (interrupted is in the whitelist) must let this through. + await mgr.set_status(record.run_id, RunStatus.error, error="Rolled back by user") + row = await store.get(record.run_id) + assert row["status"] == "error" + assert row["error"] == "Rolled back by user" + + +# --------------------------------------------------------------------------- +# cancel() claim_for_takeover False → re-read precision +# --------------------------------------------------------------------------- + + +@pytest.mark.anyio +async def test_cancel_claim_lost_to_terminal_returns_not_cancellable(): + """When cancel() reads the run as active but claim_for_takeover returns + False because the row went terminal (run finished) between the read and + the conditional UPDATE, the re-read must surface not_cancellable.""" + store = MemoryRunStore() + mgr = _make_manager(store=store, run_ownership_config=_lease_config(heartbeat_enabled=True, grace_seconds=10)) + + # Seed as running so cancel()'s first read passes the status guard. + expired = (datetime.now(UTC) - timedelta(seconds=60)).isoformat() + await store.put( + "run-race", + thread_id="t1", + status="running", + owner_worker_id="w-a", + lease_expires_at=expired, + created_at=datetime.now(UTC).isoformat(), + ) + + # Wrap claim_for_takeover: flip the row to success just before the + # conditional UPDATE so it matches 0 rows. + original = store.claim_for_takeover + + async def race_claim(run_id, *, grace_seconds, error): + store._runs[run_id]["status"] = "success" + return await original(run_id, grace_seconds=grace_seconds, error=error) + + store.claim_for_takeover = race_claim + + outcome = await mgr.cancel("run-race") + assert outcome == CancelOutcome.not_cancellable + + +@pytest.mark.anyio +async def test_cancel_claim_lost_to_takeover_returns_taken_over(): + """When cancel() reads the run as active but claim_for_takeover returns + False because another worker already took it over (row is error), the + re-read must surface taken_over.""" + store = MemoryRunStore() + mgr = _make_manager(store=store, run_ownership_config=_lease_config(heartbeat_enabled=True, grace_seconds=10)) + + expired = (datetime.now(UTC) - timedelta(seconds=60)).isoformat() + await store.put( + "run-race", + thread_id="t1", + status="running", + owner_worker_id="w-a", + lease_expires_at=expired, + created_at=datetime.now(UTC).isoformat(), + ) + + # Wrap claim_for_takeover: flip the row to error before the conditional + # UPDATE so it matches 0 rows (peer already took it over). + original = store.claim_for_takeover + + async def race_takeover(run_id, *, grace_seconds, error): + store._runs[run_id]["status"] = "error" + store._runs[run_id]["error"] = "peer claim" + return await original(run_id, grace_seconds=grace_seconds, error=error) + + store.claim_for_takeover = race_takeover + + outcome = await mgr.cancel("run-race") + assert outcome == CancelOutcome.taken_over + + +# --------------------------------------------------------------------------- +# _compute_retry_after unit tests +# --------------------------------------------------------------------------- + + +def test_compute_retry_after_null_lease_returns_none(): + from app.gateway.routers.thread_runs import _compute_retry_after + + assert _compute_retry_after(None, 10) is None + + +def test_compute_retry_after_unparseable_returns_none(): + from app.gateway.routers.thread_runs import _compute_retry_after + + assert _compute_retry_after("not-a-date", 10) is None + + +def test_compute_retry_after_normal(): + from app.gateway.routers.thread_runs import _compute_retry_after + + future = (datetime.now(UTC) + timedelta(seconds=45)).isoformat() + val = _compute_retry_after(future, 10) + 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 diff --git a/backend/tests/test_run_manager.py b/backend/tests/test_run_manager.py index b1c5dfd2f..33d213003 100644 --- a/backend/tests/test_run_manager.py +++ b/backend/tests/test_run_manager.py @@ -10,7 +10,7 @@ import pytest from sqlalchemy.exc import DatabaseError as SQLAlchemyDatabaseError from deerflow.runtime import DisconnectMode, RunManager, RunStatus -from deerflow.runtime.runs.manager import ConflictError, PersistenceRetryPolicy +from deerflow.runtime.runs.manager import CancelOutcome, ConflictError, PersistenceRetryPolicy from deerflow.runtime.runs.store.memory import MemoryRunStore ISO_RE = re.compile(r"^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}") @@ -151,7 +151,7 @@ async def test_cancel(manager: RunManager): await manager.set_status(record.run_id, RunStatus.running) cancelled = await manager.cancel(record.run_id) - assert cancelled is True + assert cancelled == CancelOutcome.cancelled assert record.abort_event.is_set() assert record.status == RunStatus.interrupted @@ -167,7 +167,7 @@ async def test_cancel_persists_interrupted_status_to_store(): cancelled = await manager.cancel(record.run_id) stored = await store.get(record.run_id) - assert cancelled is True + assert cancelled == CancelOutcome.cancelled assert stored is not None assert stored["status"] == "interrupted" @@ -323,12 +323,12 @@ async def test_reconcile_orphaned_inflight_runs_skips_rows_when_error_status_is_ @pytest.mark.anyio async def test_cancel_not_inflight(manager: RunManager): - """Cancelling a completed run should return False.""" + """Cancelling a completed run should return not_cancellable.""" record = await manager.create("thread-1") await manager.set_status(record.run_id, RunStatus.success) cancelled = await manager.cancel(record.run_id) - assert cancelled is False + assert cancelled == CancelOutcome.not_cancellable @pytest.mark.anyio diff --git a/backend/tests/test_run_repository.py b/backend/tests/test_run_repository.py index f9c975a72..31785e272 100644 --- a/backend/tests/test_run_repository.py +++ b/backend/tests/test_run_repository.py @@ -3,11 +3,13 @@ Uses a temp SQLite DB to test ORM-backed CRUD operations. """ +from datetime import UTC, datetime, timedelta + import pytest from sqlalchemy.dialects import postgresql from deerflow.persistence.run import RunRepository -from deerflow.runtime import RunManager, RunStatus +from deerflow.runtime import CancelOutcome, RunManager, RunStatus from deerflow.runtime.runs.manager import ConflictError from deerflow.runtime.runs.store.base import RunStore @@ -66,6 +68,9 @@ class _CustomRunStoreWithoutProgress(RunStore): async def create_run_atomic(self, *args, **kwargs): return {}, [] + async def claim_for_takeover(self, *args, **kwargs): + return False + @pytest.mark.anyio async def test_update_run_progress_defaults_to_noop_for_custom_store(): @@ -578,7 +583,7 @@ class TestRunRepository: cancelled = await manager.cancel(record.run_id) row = await repo.get(record.run_id) - assert cancelled is True + assert cancelled == CancelOutcome.cancelled assert row is not None assert row["status"] == "interrupted" await _cleanup() @@ -818,3 +823,64 @@ class TestRunRepository: ) await _cleanup() + + # ------------------------------------------------------------------ + # claim_for_takeover SQL path + # ------------------------------------------------------------------ + + @pytest.mark.anyio + async def test_claim_for_takeover_succeeds_with_expired_lease(self, tmp_path): + repo = await _make_repo(tmp_path) + grace = 10 + expired = (datetime.now(UTC) - timedelta(seconds=grace + 5)).isoformat() + await repo.put("run-1", thread_id="t1", status="running", owner_worker_id="w-a", lease_expires_at=expired, created_at=datetime.now(UTC).isoformat()) + + ok = await repo.claim_for_takeover("run-1", grace_seconds=grace, error="claimed") + assert ok is True + + row = await repo.get("run-1") + assert row["status"] == "error" + assert row["error"] == "claimed" + await _cleanup() + + @pytest.mark.anyio + async def test_claim_for_takeover_fails_on_valid_lease(self, tmp_path): + repo = await _make_repo(tmp_path) + grace = 10 + valid = (datetime.now(UTC) + timedelta(seconds=30)).isoformat() + await repo.put("run-1", thread_id="t1", status="running", owner_worker_id="w-a", lease_expires_at=valid, created_at=datetime.now(UTC).isoformat()) + + ok = await repo.claim_for_takeover("run-1", grace_seconds=grace, error="claimed") + assert ok is False + + row = await repo.get("run-1") + assert row["status"] == "running" + await _cleanup() + + @pytest.mark.anyio + async def test_claim_for_takeover_succeeds_with_null_lease(self, tmp_path): + repo = await _make_repo(tmp_path) + await repo.put("run-null", thread_id="t1", status="running", created_at=datetime.now(UTC).isoformat()) + + ok = await repo.claim_for_takeover("run-null", grace_seconds=10, error="claimed") + assert ok is True + + row = await repo.get("run-null") + assert row["status"] == "error" + await _cleanup() + + @pytest.mark.anyio + async def test_claim_for_takeover_fails_on_terminal_row(self, tmp_path): + repo = await _make_repo(tmp_path) + await repo.put("run-done", thread_id="t1", status="success", created_at=datetime.now(UTC).isoformat()) + + ok = await repo.claim_for_takeover("run-done", grace_seconds=10, error="claimed") + assert ok is False + await _cleanup() + + @pytest.mark.anyio + async def test_claim_for_takeover_nonexistent_run(self, tmp_path): + repo = await _make_repo(tmp_path) + ok = await repo.claim_for_takeover("no-such-run", grace_seconds=10, error="claimed") + assert ok is False + await _cleanup()