diff --git a/backend/AGENTS.md b/backend/AGENTS.md index 7e72d57d0..2c166b4d6 100644 --- a/backend/AGENTS.md +++ b/backend/AGENTS.md @@ -498,6 +498,7 @@ JSONL event stores when `GATEWAY_WORKERS > 1`. - When a persistent `RunStore` is configured, `get()` and `list_by_thread()` hydrate historical runs from the store. In-memory records win for the same `run_id` so task, abort, and stream-control state stays attached to active local runs. - Thread metadata status switches to `running` only after `RunManager.try_start()` succeeds. Pending-cancelled runs therefore skip the old `running` projection, while clients may observe the prior thread status during the short worker-startup window. - `cancel()` returns a :class:`~deerflow.runtime.CancelOutcome` enum: `cancelled` (local cancel), `taken_over` (non-owning worker claimed the run because the owner's lease expired — marks it as `error`), `lease_valid_elsewhere` (owner's lease is still alive — caller should return 409 + `Retry-After`), `not_active_locally` (heartbeat disabled, preserving the old 409 path), `not_cancellable` (terminal state), or `unknown` (not found in memory or store). `create_or_reject(..., multitask_strategy="interrupt"|"rollback")` persists interrupted status through `RunStore.update_status()`, matching normal `set_status()` transitions. +- Interrupt/rollback admission registers the replacement before its best-effort persistence of locally interrupted predecessors. If the admitting caller is cancelled during that post-registration await, `RunManager` drains a shielded replacement cleanup before propagating `CancelledError`, including across repeated cancellation. The cleanup normally persists `interrupted`; if that best-effort transition fails, it retries the active-to-`interrupted` store transition strictly and verifies the result with the replacement's captured owner identity. A concurrent peer terminal transition wins and is synchronized back into the local record rather than being overwritten or deleted. - Store-only hydrated runs are readable history. In multi-worker mode with heartbeat enabled, cancel on a store-only run can take over (mark `error`) when the owner's lease has expired past the grace window; otherwise it fails with 409 + `Retry-After`. In single-worker mode (heartbeat off), store-only runs still return 409. - Startup/orphan reconciliation must claim stale active rows with `RunStore.claim_for_takeover()`, not a plain `update_status()`. The final claim re-checks `status` and lease expiry atomically, so a heartbeat renewal between the candidate scan and the recovery write keeps the run active. - Run admission and independent checkpoint writes are first-class thread operations. `runs.operation_kind` distinguishes user-visible `run` rows from internal `checkpoint_write` reservations, while every active kind shares the existing durable active-thread uniqueness constraint. New operation kinds must go through `RunStore.create_thread_operation_atomic()` and `RunManager.reserve_thread_operation()` rather than adding another lock or metadata marker. Live and lease-less reservations are non-interruptible; an expired leased reservation can be reclaimed immediately by interrupt/rollback admission without waiting for orphan reconciliation. Lease-less rows stay fail-closed because the store cannot distinguish a stale row from a live checkpoint writer in another heartbeat-disabled worker; a rare failed delete therefore requires startup reconciliation, and heartbeat-disabled multi-worker deployment remains unsupported. Reservation bodies are attached to their caller task so loss detected by lease renewal cancels the checkpoint writer before it can continue after takeover; the context manager translates that lease-loss cancellation to `ConflictError` after cleanup so Gateway mutation routes return a retryable 409 instead of dropping the HTTP request. The cleanup scope begins immediately after durable admission, including the await that attaches the caller task, so cancellation cannot strand a locally renewed pending reservation. A failed renewal is revalidated under the manager lock before cancellation; if the reservation completed and unregistered while the store update was in flight, its request task must not be cancelled after the checkpoint write. Reservations are excluded from run history/reporting and from run-only helpers such as `list_by_thread()` and `has_inflight()`, release uses the captured owner rather than ambient user context, and local cleanup still runs when the best-effort store delete fails. `RunStore.create_run_atomic()` remains a deprecated compatibility shim for external stores that only admit normal runs; new stores must implement `create_thread_operation_atomic()` to support internal operation kinds. diff --git a/backend/packages/harness/deerflow/runtime/runs/manager.py b/backend/packages/harness/deerflow/runtime/runs/manager.py index 4d816f8d0..a8924fcab 100644 --- a/backend/packages/harness/deerflow/runtime/runs/manager.py +++ b/backend/packages/harness/deerflow/runtime/runs/manager.py @@ -1092,6 +1092,51 @@ class RunManager: user_id=user_id, ) + async def _close_cancelled_admission(self, record: RunRecord) -> None: + """Terminalize an unseen replacement and confirm its durable state.""" + await self.cancel(record.run_id) + if self._store is None: + return + + stored = await self._call_store_with_retry( + "verify cancelled admission", + record.run_id, + lambda: self._store.get(record.run_id, user_id=record.user_id), + ) + active_statuses = (RunStatus.pending.value, RunStatus.running.value) + if stored is not None and stored.get("status") in active_statuses: + # `_persist_status` is deliberately best-effort. This compensation + # path needs a strict second CAS attempt because the caller never + # receives the record and no worker can attach after it returns. + # A peer terminal transition wins the CAS and is preserved below. + await self._call_store_with_retry( + "terminalize cancelled admission", + record.run_id, + lambda: self._store.update_status(record.run_id, RunStatus.interrupted.value), + ) + stored = await self._call_store_with_retry( + "verify terminal cancelled admission", + record.run_id, + lambda: self._store.get(record.run_id, user_id=record.user_id), + ) + if stored is not None and stored.get("status") in active_statuses: + raise RuntimeError(f"Cancelled admission {record.run_id} remains active in the run store") + + if stored is None: + async with self._lock: + if self._runs.get(record.run_id) is record: + self._runs.pop(record.run_id, None) + self._unindex_run_locked(record.run_id, record.thread_id) + return + + stored_status = RunStatus(stored.get("status") or RunStatus.pending.value) + async with self._lock: + if self._runs.get(record.run_id) is record: + record.status = stored_status + record.error = stored.get("error") + record.stop_reason = stored.get("stop_reason") + record.updated_at = _now_iso() + async def _admit_thread_operation( self, thread_id: str, @@ -1260,9 +1305,29 @@ class RunManager: interrupted_records.append(r) # Outside the lock: persist interrupted status for locally-cancelled - # runs. Store-side claimed rows are already finalised. - for interrupted_record in interrupted_records: - await self._persist_status(interrupted_record, RunStatus.interrupted) + # runs. Store-side claimed rows are already finalised. Cancellation at + # this point happens after the replacement was admitted, so close that + # new run before propagating cancellation to the caller. + try: + for interrupted_record in interrupted_records: + await self._persist_status(interrupted_record, RunStatus.interrupted) + except asyncio.CancelledError: + cleanup = asyncio.create_task(self._close_cancelled_admission(record)) + cleanup.set_name(f"deerflow-close-cancelled-admission-{record.run_id}") + while not cleanup.done(): + try: + await asyncio.shield(cleanup) + except asyncio.CancelledError: + pass + except Exception: + break + try: + cleanup.result() + except asyncio.CancelledError: + logger.error("Cancelled admission cleanup task was itself cancelled for run %s", record.run_id) + except Exception: + logger.exception("Failed to close run %s after admission was cancelled", record.run_id) + raise logger.info("Run created: run_id=%s thread_id=%s", run_id, thread_id) return record diff --git a/backend/tests/test_run_manager.py b/backend/tests/test_run_manager.py index 3ee1e1229..aeeb294a0 100644 --- a/backend/tests/test_run_manager.py +++ b/backend/tests/test_run_manager.py @@ -993,6 +993,245 @@ async def test_create_or_reject_does_not_interrupt_old_run_when_new_run_store_wr assert stored_old["status"] == "running" +@pytest.mark.anyio +@pytest.mark.parametrize("strategy", ["interrupt", "rollback"]) +async def test_create_or_reject_cancellation_after_registration_interrupts_replacement( + monkeypatch: pytest.MonkeyPatch, + strategy: str, +) -> None: + """Cancellation after admission must not leave the replacement active.""" + store = MemoryRunStore() + manager = RunManager(store=store) + old = await manager.create("thread-1") + await manager.set_status(old.run_id, RunStatus.running) + persist_started = asyncio.Event() + release_persist = asyncio.Event() + original_persist_status = manager._persist_status + + async def blocking_persist_status(record: Any, status: RunStatus, **kwargs: Any) -> bool: + persist_started.set() + await asyncio.wait_for(release_persist.wait(), timeout=1) + return await original_persist_status(record, status, **kwargs) + + monkeypatch.setattr(manager, "_persist_status", blocking_persist_status) + create_task = asyncio.create_task(manager.create_or_reject("thread-1", multitask_strategy=strategy)) + await asyncio.wait_for(persist_started.wait(), timeout=1) + create_task.cancel() + release_persist.set() + + with pytest.raises(asyncio.CancelledError): + _ = await create_task + + records = await manager.list_by_thread("thread-1") + replacement = next(record for record in records if record.run_id != old.run_id) + stored_replacement = await store.get(replacement.run_id) + assert not await manager.has_inflight("thread-1") + assert replacement.status == RunStatus.interrupted + assert replacement.abort_event.is_set() + assert stored_replacement is not None + assert stored_replacement["status"] == RunStatus.interrupted.value + + +@pytest.mark.anyio +async def test_create_or_reject_repeated_cancellation_drains_replacement_cleanup( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Repeated cancellation must not abandon the durable cleanup task.""" + store = MemoryRunStore() + manager = RunManager(store=store) + old = await manager.create("thread-1") + await manager.set_status(old.run_id, RunStatus.running) + old_persist_started = asyncio.Event() + release_old_persist = asyncio.Event() + replacement_persist_started = asyncio.Event() + release_replacement_persist = asyncio.Event() + original_persist_status = manager._persist_status + + async def staged_persist_status(record: Any, status: RunStatus, **kwargs: Any) -> bool: + if record.run_id == old.run_id: + old_persist_started.set() + await asyncio.wait_for(release_old_persist.wait(), timeout=1) + else: + replacement_persist_started.set() + await asyncio.wait_for(release_replacement_persist.wait(), timeout=1) + return await original_persist_status(record, status, **kwargs) + + monkeypatch.setattr(manager, "_persist_status", staged_persist_status) + create_task = asyncio.create_task(manager.create_or_reject("thread-1", multitask_strategy="interrupt")) + await asyncio.wait_for(old_persist_started.wait(), timeout=1) + replacement = next(record for record in manager._runs.values() if record.run_id != old.run_id) + + create_task.cancel() + release_old_persist.set() + await asyncio.wait_for(replacement_persist_started.wait(), timeout=1) + create_task.cancel() + await asyncio.sleep(0) + assert not create_task.done() + + release_replacement_persist.set() + with pytest.raises(asyncio.CancelledError): + _ = await asyncio.wait_for(create_task, timeout=1) + + stored_replacement = await store.get(replacement.run_id) + assert replacement.status == RunStatus.interrupted + assert stored_replacement is not None + assert stored_replacement["status"] == RunStatus.interrupted.value + + +@pytest.mark.anyio +async def test_create_or_reject_retries_replacement_when_cancel_status_cannot_persist( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A failed best-effort update must get a strict durable retry.""" + + class FailFirstReplacementInterruptStore(MemoryRunStore): + failed = False + + async def get(self, run_id: str, *, user_id: str | None = None) -> dict[str, Any] | None: + raw = await super().get(run_id) + if raw is not None and raw.get("status") == RunStatus.pending.value and raw.get("user_id") is not None and user_id != raw.get("user_id"): + raise RuntimeError("replacement lookup was not owner-scoped") + return await super().get(run_id, user_id=user_id) + + async def update_status(self, run_id: str, status: str, **kwargs: Any) -> bool: + row = await super().get(run_id) + if not self.failed and status == RunStatus.interrupted.value and row is not None and row.get("status") == RunStatus.pending.value: + self.failed = True + raise RuntimeError("replacement status write failed") + return await super().update_status(run_id, status, **kwargs) + + store = FailFirstReplacementInterruptStore() + manager = RunManager(store=store) + old = await manager.create("thread-1", user_id="owner-1") + await manager.set_status(old.run_id, RunStatus.running) + old_persist_started = asyncio.Event() + release_old_persist = asyncio.Event() + original_persist_status = manager._persist_status + + async def block_old_persist(record: Any, status: RunStatus, **kwargs: Any) -> bool: + if record.run_id != old.run_id: + return await original_persist_status(record, status, **kwargs) + old_persist_started.set() + await asyncio.wait_for(release_old_persist.wait(), timeout=1) + return await original_persist_status(record, status, **kwargs) + + monkeypatch.setattr(manager, "_persist_status", block_old_persist) + create_task = asyncio.create_task(manager.create_or_reject("thread-1", multitask_strategy="interrupt", user_id="owner-1")) + await asyncio.wait_for(old_persist_started.wait(), timeout=1) + replacement = next(record for record in manager._runs.values() if record.run_id != old.run_id) + create_task.cancel() + release_old_persist.set() + + with pytest.raises(asyncio.CancelledError): + _ = await asyncio.wait_for(create_task, timeout=1) + + stored_replacement = await store.get(replacement.run_id, user_id="owner-1") + assert stored_replacement is not None + assert stored_replacement["status"] == RunStatus.interrupted.value + assert replacement.status == RunStatus.interrupted + assert not await manager.has_inflight("thread-1") + + +@pytest.mark.anyio +async def test_create_or_reject_cleanup_failure_preserves_caller_cancellation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Cleanup IO failure must not replace the caller's CancelledError.""" + + class FailingReplacementCleanupStore(MemoryRunStore): + replacement_update_failed = False + + async def update_status(self, run_id: str, status: str, **kwargs: Any) -> bool: + row = await super().get(run_id) + if status == RunStatus.interrupted.value and row is not None and row.get("status") == RunStatus.pending.value: + self.replacement_update_failed = True + raise RuntimeError("replacement status unavailable") + return await super().update_status(run_id, status, **kwargs) + + async def get(self, run_id: str, *, user_id: str | None = None) -> dict[str, Any] | None: + row = await super().get(run_id, user_id=user_id) + if self.replacement_update_failed and row is not None and row.get("status") == RunStatus.pending.value: + raise RuntimeError("replacement verification unavailable") + return row + + store = FailingReplacementCleanupStore() + manager = RunManager(store=store) + old = await manager.create("thread-1") + await manager.set_status(old.run_id, RunStatus.running) + old_persist_started = asyncio.Event() + release_old_persist = asyncio.Event() + original_persist_status = manager._persist_status + + async def block_old_persist(record: Any, status: RunStatus, **kwargs: Any) -> bool: + if record.run_id != old.run_id: + return await original_persist_status(record, status, **kwargs) + old_persist_started.set() + await asyncio.wait_for(release_old_persist.wait(), timeout=1) + return await original_persist_status(record, status, **kwargs) + + monkeypatch.setattr(manager, "_persist_status", block_old_persist) + create_task = asyncio.create_task(manager.create_or_reject("thread-1", multitask_strategy="interrupt")) + await asyncio.wait_for(old_persist_started.wait(), timeout=1) + create_task.cancel() + release_old_persist.set() + + with pytest.raises(asyncio.CancelledError): + _ = await asyncio.wait_for(create_task, timeout=1) + + +@pytest.mark.anyio +async def test_create_or_reject_preserves_peer_terminal_status_during_cancel_retry( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A peer terminal transition must win the strict cancellation retry.""" + + class PeerWinsReplacementInterruptStore(MemoryRunStore): + replacement_attempts = 0 + + async def update_status(self, run_id: str, status: str, **kwargs: Any) -> bool: + row = await self.get(run_id) + if status == RunStatus.interrupted.value and row is not None and row.get("status") == RunStatus.pending.value: + self.replacement_attempts += 1 + if self.replacement_attempts == 1: + raise RuntimeError("replacement status write failed") + await super().update_status(run_id, RunStatus.error.value, error="peer takeover") + return False + return await super().update_status(run_id, status, **kwargs) + + store = PeerWinsReplacementInterruptStore() + manager = RunManager(store=store) + old = await manager.create("thread-1") + await manager.set_status(old.run_id, RunStatus.running) + old_persist_started = asyncio.Event() + release_old_persist = asyncio.Event() + original_persist_status = manager._persist_status + + async def block_old_persist(record: Any, status: RunStatus, **kwargs: Any) -> bool: + if record.run_id != old.run_id: + return await original_persist_status(record, status, **kwargs) + old_persist_started.set() + await asyncio.wait_for(release_old_persist.wait(), timeout=1) + return await original_persist_status(record, status, **kwargs) + + monkeypatch.setattr(manager, "_persist_status", block_old_persist) + create_task = asyncio.create_task(manager.create_or_reject("thread-1", multitask_strategy="interrupt")) + await asyncio.wait_for(old_persist_started.wait(), timeout=1) + replacement = next(record for record in manager._runs.values() if record.run_id != old.run_id) + create_task.cancel() + release_old_persist.set() + + with pytest.raises(asyncio.CancelledError): + _ = await asyncio.wait_for(create_task, timeout=1) + + stored_replacement = await store.get(replacement.run_id) + assert stored_replacement is not None + assert stored_replacement["status"] == RunStatus.error.value + assert stored_replacement["error"] == "peer takeover" + assert replacement.status == RunStatus.error + assert replacement.error == "peer takeover" + assert not await manager.has_inflight("thread-1") + + @pytest.mark.anyio async def test_create_or_reject_rollback_persists_interrupted_status_to_store(): """rollback strategy should persist interrupted status for old runs."""