mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-07-21 10:15:47 +00:00
fix(runs): cancel degrades to lease takeover for multi-worker (#4064)
* 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 <noreply@anthropic.com> * 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 <noreply@anthropic.com> * 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 <noreply@anthropic.com> * 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 <noreply@anthropic.com> * 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 <noreply@anthropic.com> * 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 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com> Co-authored-by: rayhpeng <rayhpeng@gmail.com>
This commit is contained in:
heart-scalpel
co-authored by
Claude Opus 4.7
rayhpeng
Claude Opus 4.7
rayhpeng
parent
3e7baba39a
commit
b53c1ae0e0
+2
-2
@@ -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.
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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())
|
||||
)
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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."""
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
*,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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())
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user