From 090e80c1ddade25a26efc2de5087136ebcb1daae Mon Sep 17 00:00:00 2001 From: Huixin615 Date: Mon, 27 Jul 2026 07:25:34 +0800 Subject: [PATCH] fix(runtime): fail-stop runs when lease ownership cannot be confirmed (#4431) * fix(runtime): fail-stop runs after lease expiry * test(runtime): cover late successful lease renewal --------- Co-authored-by: Willem Jiang --- CHANGELOG.md | 11 +- README.md | 2 + backend/AGENTS.md | 1 + .../harness/deerflow/persistence/run/sql.py | 17 +- .../harness/deerflow/runtime/runs/manager.py | 204 ++++++++++++++-- .../deerflow/runtime/runs/store/base.py | 4 +- .../deerflow/runtime/runs/store/memory.py | 23 +- .../harness/deerflow/runtime/runs/worker.py | 29 ++- .../tests/test_multi_worker_run_ownership.py | 225 ++++++++++++++++++ backend/tests/test_run_repository.py | 15 ++ backend/tests/test_run_worker_delivery.py | 40 +++- backend/tests/test_run_worker_rollback.py | 44 +++- config.example.yaml | 5 + 13 files changed, 567 insertions(+), 53 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 4c2ffb08d..361ddbbb4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -250,10 +250,12 @@ This section accumulates work toward the **2.1.0** milestone facts list is empty; stop the busy-spin in the debounced update queue; and flush the memory queue on graceful shutdown to prevent loss. ([#3719], [#4074], [#4076], [#4034], [#4217], [#3993], [#3992], [#4073], [#4181]) -- **runs:** Close multi-worker ownership gaps in run atomicity; degrade cancel - to lease takeover for multi-worker; keep `create_thread` idempotent when the - insert loses a race; read `stop_reason` from runtime context; and persist run - duration in checkpoints for history reads. ([#4003], [#4064], [#3800], [#4188], +- **runs:** Close multi-worker ownership gaps in run atomicity; fail-stop local + execution when lease renewal cannot be confirmed before its deadline and + fence late completion writes after peer takeover; degrade cancel to lease + takeover for multi-worker; keep `create_thread` idempotent when the insert + loses a race; read `stop_reason` from runtime context; and persist run duration + in checkpoints for history reads. ([#4003], [#4064], [#4414], [#3800], [#4188], [#4118]) - **runtime:** Serialize SQLite event-store writes to prevent per-thread sequence collisions; skip hidden human messages in the journal; and drop the @@ -1121,4 +1123,5 @@ with **180 merged pull requests** since the first 2.0 milestone tag. [#4287]: https://github.com/bytedance/deer-flow/pull/4287 [#4288]: https://github.com/bytedance/deer-flow/pull/4288 [#4324]: https://github.com/bytedance/deer-flow/issues/4324 +[#4414]: https://github.com/bytedance/deer-flow/issues/4414 [#4424]: https://github.com/bytedance/deer-flow/issues/4424 diff --git a/README.md b/README.md index b081c2102..87dace770 100644 --- a/README.md +++ b/README.md @@ -277,6 +277,8 @@ DeerFlow still uses `Forwarded` / `X-Forwarded-*` headers to recover the browser > [!IMPORTANT] > The Gateway still owns active run tasks in process, so production defaults to a single Gateway worker (`GATEWAY_WORKERS=1`). Multi-worker deployments require Postgres, the Redis stream bridge (`stream_bridge.type: redis`), `run_ownership.heartbeat_enabled: true`, and `run_events.backend: db`; process-local memory/JSONL event stores cannot enforce singleton delivery receipts across workers. The bridge shares SSE delivery and bounded `Last-Event-ID` replay across workers. When a valid reconnect cursor has been trimmed, or a subscriber that already established an empty-stream wait falls behind before its first delivery, Memory and Redis emit a machine-readable SSE `gap` event instead of silently returning a partial replay; the Web UI reloads durable thread/event state and resumes from the retained tail. Lease reconciliation marks runs from dead workers as errors, persists their delivery receipts, publishes the terminal stream marker, schedules retained-stream cleanup, and updates the affected thread status. SSE and `/wait` consumers also refresh durable status on heartbeats as a fallback if terminal publication fails. Malformed Redis reconnect IDs live-tail new events instead of replaying the retained buffer, and the rolling retained-buffer TTL (`stream_ttl_seconds`) remains a cleanup safety net rather than a run timeout. IM channel state and other process-local services still need their own multi-worker coordination. > +> With lease heartbeat enabled, a transient RunStore renewal error is retried only until the last confirmed lease expires; the stale worker then cancels local execution and suppresses checkpoint, completion-hook, delivery-receipt, and thread-status finalization. A remote tool side effect already in flight may still be outside local cancellation. +> > Reconciliation uses an atomic takeover claim that re-checks the lease after candidate selection, so a successful owner renewal wins over orphan recovery and only one reconciler can report a run as recovered. See [CONTRIBUTING.md](CONTRIBUTING.md) for detailed Docker development guide. diff --git a/backend/AGENTS.md b/backend/AGENTS.md index 55d6f66da..bb857bb13 100644 --- a/backend/AGENTS.md +++ b/backend/AGENTS.md @@ -500,6 +500,7 @@ JSONL event stores when `GATEWAY_WORKERS > 1`. - `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. +- A local worker's `RunRecord.lease_expires_at` is the last durably confirmed ownership deadline. `_renew_leases()` bounds each renewal attempt by that deadline: transient store exceptions remain retryable while it is valid, but an exception or blocked call that reaches expiry sets the process-local `ownership_lost` fence, raises `abort_event`, and cancels the run task. Fenced workers do not perform subsequent journal/delivery-receipt, progress/completion/status, checkpoint/thread-metadata, or `on_run_completed` writes; the peer recovery path owns the terminal receipt. `RunStore.update_run_completion()` also refuses to replace a different terminal status, closing the peer-takeover/late-finalization race. `grace_seconds` delays peer reclamation for clock skew but is not extra execution time for an owner that can no longer confirm its lease. Already-committed remote tool side effects remain outside this local cancellation boundary. - Startup/orphan reconciliation must claim stale active rows with `RunStore.claim_for_takeover()`, not a plain `update_status()`. The final claim re-checks `status` and lease expiry atomically, so a heartbeat renewal between the candidate scan and the recovery write keeps the run active. - Run admission and independent checkpoint writes are first-class thread operations. `runs.operation_kind` distinguishes user-visible `run` rows from internal `checkpoint_write` reservations, while every active kind shares the existing durable active-thread uniqueness constraint. New operation kinds must go through `RunStore.create_thread_operation_atomic()` and `RunManager.reserve_thread_operation()` rather than adding another lock or metadata marker. Live and lease-less reservations are non-interruptible; an expired leased reservation can be reclaimed immediately by interrupt/rollback admission without waiting for orphan reconciliation. Lease-less rows stay fail-closed because the store cannot distinguish a stale row from a live checkpoint writer in another heartbeat-disabled worker; a rare failed delete therefore requires startup reconciliation, and heartbeat-disabled multi-worker deployment remains unsupported. Reservation bodies are attached to their caller task so loss detected by lease renewal cancels the checkpoint writer before it can continue after takeover; the context manager translates that lease-loss cancellation to `ConflictError` after cleanup so Gateway mutation routes return a retryable 409 instead of dropping the HTTP request. The cleanup scope begins immediately after durable admission, including the await that attaches the caller task, so cancellation cannot strand a locally renewed pending reservation. A failed renewal is revalidated under the manager lock before cancellation; if the reservation completed and unregistered while the store update was in flight, its request task must not be cancelled after the checkpoint write. Reservations are excluded from run history/reporting and from run-only helpers such as `list_by_thread()` and `has_inflight()`, release uses the captured owner rather than ambient user context, and local cleanup still runs when the best-effort store delete fails. `RunStore.create_run_atomic()` remains a deprecated compatibility shim for external stores that only admit normal runs; new stores must implement `create_thread_operation_atomic()` to support internal operation kinds. - Gateway checkpoint mutations outside run execution must use `services.reserve_checkpoint_write()`, which composes the process-local thread lock with the durable `checkpoint_write` reservation. Manual compaction, `POST /threads/{id}/state`, and both goal mutation routes (`PUT` / `DELETE /threads/{id}/goal`, including creation of a missing goal checkpoint) use this boundary, so an existing run blocks the write and the reservation blocks new reject/interrupt/rollback runs across workers. diff --git a/backend/packages/harness/deerflow/persistence/run/sql.py b/backend/packages/harness/deerflow/persistence/run/sql.py index 94c148dc8..a0ea18afb 100644 --- a/backend/packages/harness/deerflow/persistence/run/sql.py +++ b/backend/packages/harness/deerflow/persistence/run/sql.py @@ -315,7 +315,8 @@ class RunRepository(RunStore): ) -> bool: """Update status + token usage + convenience fields on run completion. - Returns ``False`` when no run row matched the requested ``run_id``. + Returns ``False`` when the row is missing or already has a conflicting + terminal outcome. """ values: dict[str, Any] = { "status": status, @@ -336,8 +337,20 @@ class RunRepository(RunStore): values["first_human_message"] = first_human_message[:2000] if error is not None: values["error"] = error + allowed_sources = ["pending", "running"] + if status not in allowed_sources: + allowed_sources.append(status) + if status == "error" and "interrupted" not in allowed_sources: + allowed_sources.append("interrupted") 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_(tuple(allowed_sources)), + ) + .values(**values) + ) await session.commit() return result.rowcount != 0 diff --git a/backend/packages/harness/deerflow/runtime/runs/manager.py b/backend/packages/harness/deerflow/runtime/runs/manager.py index a8924fcab..eeb965c8b 100644 --- a/backend/packages/harness/deerflow/runtime/runs/manager.py +++ b/backend/packages/harness/deerflow/runtime/runs/manager.py @@ -190,6 +190,10 @@ class RunRecord: finalizing: bool = False owner_worker_id: str | None = None lease_expires_at: str | None = None + # Process-local fencing signal. Once set, this worker must not perform + # further durable run/thread finalization because its lease ownership is + # either known to be lost or could not be confirmed before expiry. + ownership_lost: bool = False stop_reason: str | None = None @@ -365,6 +369,13 @@ class RunManager: async def _persist_status(self, record: RunRecord, status: RunStatus, *, error: str | None = None, stop_reason: str | None = None) -> bool: """Best-effort persist a status transition to the backing store.""" + if record.ownership_lost: + logger.warning( + "Skipped status update to %s for run %s after lease ownership was lost", + status.value, + record.run_id, + ) + return False if self._store is None: return True row_recovery_payload = self._store_put_payload(record, error=error, stop_reason=stop_reason) @@ -384,12 +395,25 @@ class RunManager: existing = await self._store.get(record.run_id) if existing is not None: existing_status = existing.get("status") + if existing_status == status.value: + logger.info( + "Run %s status update to %s was already persisted", + record.run_id, + status.value, + ) + return True 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, ) + if self.heartbeat_enabled and not record.store_only: + await self._mark_ownership_lost( + record, + reason="A peer terminalized the run before this worker could persist its outcome.", + require_active=False, + ) else: logger.info( "Run %s status update to %s skipped: store row already at %s (local cancel/completion race)", @@ -446,8 +470,12 @@ class RunManager: async def update_run_completion(self, run_id: str, **kwargs) -> None: """Persist token usage and completion data to the backing store.""" row_recovery_payload: dict[str, Any] | None = None + record: RunRecord | None = None async with self._lock: record = self._runs.get(run_id) + if record is not None and record.ownership_lost: + logger.warning("Skipped completion persistence for run %s after lease ownership was lost", run_id) + return if record is not None: for key, value in kwargs.items(): if key == "status": @@ -465,6 +493,22 @@ class RunManager: lambda: self._store.update_run_completion(run_id, **kwargs), ) if updated is False: + existing = await self._store.get(run_id) + requested_status = kwargs.get("status") + if existing is not None and existing.get("status") != requested_status: + existing_status = existing.get("status") + logger.warning( + "Run completion update for %s skipped because store row is already at %s", + run_id, + existing_status, + ) + if existing_status == "error" and record is not None and self.heartbeat_enabled: + await self._mark_ownership_lost( + record, + reason="A peer terminalized the run before completion data was persisted.", + require_active=False, + ) + return if row_recovery_payload is None: logger.warning("Failed to recreate missing run %s for completion persistence", run_id) return @@ -486,7 +530,7 @@ class RunManager: async with self._lock: record = self._runs.get(run_id) if record is not None: - should_persist = record.status == RunStatus.running + should_persist = record.status == RunStatus.running and not record.ownership_lost if record is not None and should_persist: for key, value in kwargs.items(): if hasattr(record, key) and value is not None: @@ -775,6 +819,13 @@ class RunManager: if record is None: logger.warning("set_status called for unknown run %s", run_id) return + if record.ownership_lost: + logger.warning( + "Skipped local status transition to %s for run %s after lease ownership was lost", + status.value, + run_id, + ) + return record.status = status record.updated_at = _now_iso() if error is not None: @@ -782,7 +833,15 @@ class RunManager: if stop_reason is not None: record.stop_reason = stop_reason if persist: - await self._persist_status(record, status, error=error, stop_reason=stop_reason) + persisted = await self._persist_status(record, status, error=error, stop_reason=stop_reason) + if not persisted and self.heartbeat_enabled and status == RunStatus.success and not record.ownership_lost: + await self._mark_ownership_lost( + record, + reason="Successful completion could not be confirmed in the durable run store.", + require_active=False, + ) + if record.ownership_lost: + return logger.info("Run %s -> %s", run_id, status.value) async def persist_current_status(self, run_id: str) -> bool: @@ -795,7 +854,14 @@ class RunManager: status = record.status error = record.error stop_reason = record.stop_reason - return await self._persist_status(record, status, error=error, stop_reason=stop_reason) + persisted = await self._persist_status(record, status, error=error, stop_reason=stop_reason) + if not persisted and self.heartbeat_enabled and status == RunStatus.success and not record.ownership_lost: + await self._mark_ownership_lost( + record, + reason="Successful completion could not be confirmed in the durable run store.", + require_active=False, + ) + return persisted async def _ensure_delivery_receipt(self, record: RunRecord) -> bool: """Idempotently persist a zero-delivery receipt during recovery.""" @@ -1519,6 +1585,61 @@ class RunManager: """ return self._run_ownership_config.grace_seconds if self._run_ownership_config else 10 + @staticmethod + def _parse_lease_deadline(lease_expires_at: str | None) -> datetime | None: + """Parse the last durably confirmed lease expiry. + + Missing or malformed deadlines are unsafe in heartbeat mode: the local + worker has no bounded interval during which it can prove ownership. + """ + if lease_expires_at is None: + return None + try: + deadline = datetime.fromisoformat(lease_expires_at) + except (TypeError, ValueError): + return None + if deadline.tzinfo is None: + deadline = deadline.replace(tzinfo=UTC) + return deadline + + async def _mark_ownership_lost( + self, + record: RunRecord, + *, + reason: str, + require_active: bool = True, + ) -> bool: + """Fence one local run and cancel its execution task. + + No store write is attempted here: once the last confirmed lease has + expired, this worker is no longer authorized to publish a terminal + outcome. A peer reconciler owns durable terminalization. + """ + task_to_cancel: asyncio.Task | None = None + async with self._lock: + current = self._runs.get(record.run_id) + if current is not record: + return False + if require_active: + if record.status not in (RunStatus.pending, RunStatus.running): + return False + if record.task is not None and record.task.done(): + return False + if record.ownership_lost: + return True + record.ownership_lost = True + record.abort_event.set() + record.status = RunStatus.error + record.error = reason + record.updated_at = _now_iso() + if record.task is not None and not record.task.done() and record.task is not asyncio.current_task(): + task_to_cancel = record.task + + if task_to_cancel is not None: + task_to_cancel.cancel() + logger.error("Run %s lost lease ownership; local execution was fenced: %s", record.run_id, reason) + return True + async def start_heartbeat(self) -> None: """Start the background lease-renewal task. @@ -1595,11 +1716,16 @@ class RunManager: self._schedule_orphan_reconciliation() async def _renew_leases(self) -> None: - """Renew the lease on every locally-owned active run.""" + """Renew locally-owned leases, failing closed at their deadlines. + + ``RunRecord.lease_expires_at`` advances only after a successful durable + renewal, so it is the last confirmed ownership deadline. Transient + exceptions are tolerated before that deadline; a call that blocks or + keeps failing through it fences the local run. + """ if self._store is None or self._run_ownership_config is None: return lease_seconds = self._run_ownership_config.lease_seconds - new_expiry = (datetime.now(UTC) + timedelta(seconds=lease_seconds)).isoformat() async with self._lock: # Renew any pending/running run owned by this worker unless its @@ -1615,17 +1741,34 @@ class RunManager: active_runs = [(rid, record) for rid, record in self._runs.items() if record.status in (RunStatus.pending, RunStatus.running) and record.owner_worker_id == self._worker_id and (record.task is None or not record.task.done())] for run_id, record in active_runs: - try: - updated = await self._call_store_with_retry( - "update_lease", - run_id, - lambda: self._store.update_lease( - run_id, - owner_worker_id=self._worker_id, - lease_expires_at=new_expiry, - ), + confirmed_deadline = self._parse_lease_deadline(record.lease_expires_at) + if confirmed_deadline is None or confirmed_deadline <= datetime.now(UTC): + await self._mark_ownership_lost( + record, + reason="Lease ownership could not be confirmed before the last confirmed lease expired.", ) + continue + + remaining = (confirmed_deadline - datetime.now(UTC)).total_seconds() + new_expiry = (datetime.now(UTC) + timedelta(seconds=lease_seconds)).isoformat() + try: + async with asyncio.timeout(remaining): + updated = await self._call_store_with_retry( + "update_lease", + run_id, + lambda: self._store.update_lease( + run_id, + owner_worker_id=self._worker_id, + lease_expires_at=new_expiry, + ), + ) if updated: + if confirmed_deadline <= datetime.now(UTC): + await self._mark_ownership_lost( + record, + reason="Lease renewal completed after the last confirmed lease had already expired.", + ) + continue # Unsynced write is benign: ``lease_expires_at`` is the # only field on an existing record this path mutates, so # there is no concurrent writer to race against @@ -1641,18 +1784,29 @@ class RunManager: # finalisation. async with self._lock: still_active = self._runs.get(run_id) is record and record.status in (RunStatus.pending, RunStatus.running) and record.owner_worker_id == self._worker_id and (record.task is None or not record.task.done()) - if still_active: - 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() - if record.task is not None: - record.task.cancel() + if still_active: + 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, + ) + await self._mark_ownership_lost( + record, + reason="The durable store rejected lease renewal for this worker.", + ) except Exception: - logger.warning("Failed to renew lease for run %s", run_id, exc_info=True) + if confirmed_deadline <= datetime.now(UTC): + await self._mark_ownership_lost( + record, + reason="Lease ownership could not be confirmed before the last confirmed lease expired.", + ) + else: + logger.warning( + "Failed to renew lease for run %s before its confirmed deadline; will retry", + run_id, + exc_info=True, + ) async def _reconcile_orphans_periodic(self) -> None: """Sweep for expired leases owned by dead peers. diff --git a/backend/packages/harness/deerflow/runtime/runs/store/base.py b/backend/packages/harness/deerflow/runtime/runs/store/base.py index b6796ccec..385d9714e 100644 --- a/backend/packages/harness/deerflow/runtime/runs/store/base.py +++ b/backend/packages/harness/deerflow/runtime/runs/store/base.py @@ -146,7 +146,9 @@ class RunStore(abc.ABC): ) -> bool | None: """Persist final completion fields. - Returns ``False`` when the store can prove no row was updated. + Implementations must not replace a different terminal status. Returns + ``False`` when the row is missing or already has a conflicting terminal + outcome. """ pass diff --git a/backend/packages/harness/deerflow/runtime/runs/store/memory.py b/backend/packages/harness/deerflow/runtime/runs/store/memory.py index eac66fabd..e92356085 100644 --- a/backend/packages/harness/deerflow/runtime/runs/store/memory.py +++ b/backend/packages/harness/deerflow/runtime/runs/store/memory.py @@ -144,14 +144,21 @@ class MemoryRunStore(RunStore): self._unindex_run(run_id, run["thread_id"]) async def update_run_completion(self, run_id, *, status, **kwargs): - if run_id in self._runs: - self._runs[run_id]["status"] = status - for key, value in kwargs.items(): - if value is not None: - self._runs[run_id][key] = value - 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 + current_status = run.get("status") + allowed_sources = {"pending", "running", status} + if status == "error": + allowed_sources.add("interrupted") + if current_status not in allowed_sources: + return False + run["status"] = status + for key, value in kwargs.items(): + if value is not None: + run[key] = value + run["updated_at"] = datetime.now(UTC).isoformat() + return True async def update_run_progress(self, run_id, **kwargs): if run_id in self._runs and self._runs[run_id].get("status") == "running": diff --git a/backend/packages/harness/deerflow/runtime/runs/worker.py b/backend/packages/harness/deerflow/runtime/runs/worker.py index 7af4c76f8..38a5ec9ff 100644 --- a/backend/packages/harness/deerflow/runtime/runs/worker.py +++ b/backend/packages/harness/deerflow/runtime/runs/worker.py @@ -514,7 +514,7 @@ async def run_agent( return started = True - if thread_store is not None: + if not record.ownership_lost and thread_store is not None: try: await thread_store.update_status(thread_id, "running") except Exception: @@ -927,12 +927,18 @@ async def run_agent( ) finally: + if record.ownership_lost: + logger.warning( + "Skipping durable finalization for run %s because this worker no longer owns its lease", + run_id, + ) + # Persist any subagent step events still buffered (#3779) — including on # abort/exception paths, where the stream loop broke before its own flush. - if subagent_events is not None: + if not record.ownership_lost and subagent_events is not None: await subagent_events.flush() - if event_store is not None and pre_run_workspace_snapshot is not None: + if not record.ownership_lost and event_store is not None and pre_run_workspace_snapshot is not None: try: await record_workspace_changes( event_store, @@ -948,7 +954,8 @@ async def run_agent( # receipt uses a run-scoped idempotent write shared with recovery, then # the staged terminal status is persisted. This ordering closes the # crash window where a terminal run could otherwise outlive its receipt. - if journal is not None: + # A fenced worker leaves receipt recovery to the peer that claimed it. + if not record.ownership_lost and journal is not None: try: await journal.flush() except Exception: @@ -961,7 +968,7 @@ async def run_agent( content=journal.get_delivery_content(), ) - if event_store is not None: + if not record.ownership_lost and event_store is not None: try: # Even after bounded receipt retries are exhausted, persist the # real worker outcome. Leaving a successful row inflight would @@ -971,7 +978,7 @@ async def run_agent( except Exception: logger.warning("Failed to persist terminal status for run %s after delivery receipt attempts", run_id, exc_info=True) - if journal is not None and persist_completion: + if not record.ownership_lost and journal is not None and persist_completion: try: # Persist token usage + convenience fields to RunStore completion = journal.get_completion_data() @@ -979,7 +986,7 @@ async def run_agent( except Exception: logger.warning("Failed to persist run completion for %s (non-fatal)", run_id, exc_info=True) - if started and checkpointer is not None and record.status == RunStatus.interrupted: + if started and not record.ownership_lost and checkpointer is not None and record.status == RunStatus.interrupted: try: await run_manager.wait_for_prior_finalizing(thread_id, run_id) if not await run_manager.has_later_started_run(thread_id, run_id): @@ -988,7 +995,7 @@ async def run_agent( logger.debug("Failed to generate interrupted title for thread %s (non-fatal)", thread_id) # Sync title from checkpoint to threads_meta.display_name - if started and checkpointer is not None and thread_store is not None: + if started and not record.ownership_lost and checkpointer is not None and thread_store is not None: try: ckpt_config = {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}} ckpt_tuple = await checkpointer.aget_tuple(ckpt_config) @@ -1002,7 +1009,7 @@ async def run_agent( # Persist run duration to checkpoint metadata so history reads # don't need to correlate runs and events. - if started and checkpointer is not None and record.status == RunStatus.success: + if started and not record.ownership_lost and checkpointer is not None and record.status == RunStatus.success: try: created = datetime.fromisoformat(record.created_at.replace("Z", "+00:00")) updated = datetime.fromisoformat(record.updated_at.replace("Z", "+00:00")) @@ -1020,14 +1027,14 @@ async def run_agent( logger.debug("Failed to persist run duration for thread %s run %s (non-fatal)", thread_id, run_id) # Update threads_meta status based on run outcome - if started and thread_store is not None: + if started and not record.ownership_lost and thread_store is not None: try: final_status = "idle" if record.status == RunStatus.success else record.status.value await thread_store.update_status(thread_id, final_status) except Exception: logger.debug("Failed to update thread_meta status for %s (non-fatal)", thread_id) - if ctx.on_run_completed is not None: + if not record.ownership_lost and ctx.on_run_completed is not None: try: await ctx.on_run_completed(record) except Exception: diff --git a/backend/tests/test_multi_worker_run_ownership.py b/backend/tests/test_multi_worker_run_ownership.py index 57f237748..52e0dd6e0 100644 --- a/backend/tests/test_multi_worker_run_ownership.py +++ b/backend/tests/test_multi_worker_run_ownership.py @@ -859,6 +859,147 @@ async def test_heartbeat_renews_pending_run_before_task_is_spawned(): assert record.lease_expires_at > original_lease +@pytest.mark.anyio +async def test_transient_renewal_exception_before_deadline_keeps_run_alive(): + """A renewal error is retryable while the last confirmed lease is valid.""" + config = _lease_config(lease_seconds=30, heartbeat_enabled=True) + store = MemoryRunStore() + manager = _make_manager(store=store, worker_id="worker-a", run_ownership_config=config) + record = await manager.create("thread-1") + await manager.set_status(record.run_id, RunStatus.running) + record.task = asyncio.create_task(asyncio.sleep(3600)) + original_update_lease = store.update_lease + attempts = 0 + + async def fail_once_then_renew(*args, **kwargs): + nonlocal attempts + attempts += 1 + if attempts == 1: + raise OSError("temporary database outage") + return await original_update_lease(*args, **kwargs) + + store.update_lease = fail_once_then_renew + original_expiry = record.lease_expires_at + + try: + await manager._renew_leases() + + assert record.abort_event.is_set() is False + assert record.task.done() is False + assert record.lease_expires_at == original_expiry + + await manager._renew_leases() + + assert attempts == 2 + assert record.abort_event.is_set() is False + assert record.task.done() is False + assert record.lease_expires_at is not None + assert original_expiry is not None + assert record.lease_expires_at > original_expiry + finally: + record.task.cancel() + with pytest.raises(asyncio.CancelledError): + await record.task + + +@pytest.mark.anyio +async def test_renewal_exception_through_confirmed_expiry_fail_stops_run(): + """The owner must stop once errors outlive its last confirmed lease.""" + config = _lease_config(lease_seconds=30, heartbeat_enabled=True) + store = MemoryRunStore() + manager = _make_manager(store=store, worker_id="worker-a", run_ownership_config=config) + record = await manager.create("thread-1") + await manager.set_status(record.run_id, RunStatus.running) + record.task = asyncio.create_task(asyncio.sleep(3600)) + attempts = 0 + + async def fail_renewal(*_args, **_kwargs): + nonlocal attempts + attempts += 1 + raise OSError("database unreachable") + + store.update_lease = fail_renewal + + await manager._renew_leases() + assert record.abort_event.is_set() is False + assert record.task.done() is False + + expired = (datetime.now(UTC) - timedelta(seconds=1)).isoformat() + record.lease_expires_at = expired + store._runs[record.run_id]["lease_expires_at"] = expired + + await manager._renew_leases() + await asyncio.sleep(0) + + assert attempts == 1 + assert record.ownership_lost is True + assert record.abort_event.is_set() is True + assert record.task.cancelled() + + +@pytest.mark.anyio +async def test_hung_renewal_is_bounded_by_confirmed_lease_deadline(): + """A blocked store call cannot keep execution alive beyond the lease.""" + config = _lease_config(lease_seconds=30, heartbeat_enabled=True) + store = MemoryRunStore() + manager = _make_manager(store=store, worker_id="worker-a", run_ownership_config=config) + record = await manager.create("thread-1") + await manager.set_status(record.run_id, RunStatus.running) + record.task = asyncio.create_task(asyncio.sleep(3600)) + near_expiry = (datetime.now(UTC) + timedelta(milliseconds=50)).isoformat() + record.lease_expires_at = near_expiry + store._runs[record.run_id]["lease_expires_at"] = near_expiry + + async def hang_renewal(*_args, **_kwargs): + await asyncio.Event().wait() + + store.update_lease = hang_renewal + + await asyncio.wait_for(manager._renew_leases(), timeout=1) + await asyncio.sleep(0) + + assert record.ownership_lost is True + assert record.abort_event.is_set() is True + assert record.task.cancelled() + + +@pytest.mark.anyio +async def test_late_successful_renewal_still_fences_local_run(): + """A renewal confirmed after the old deadline cannot revive local work.""" + config = _lease_config(lease_seconds=30, heartbeat_enabled=True) + store = MemoryRunStore() + manager = _make_manager(store=store, worker_id="worker-a", run_ownership_config=config) + record = await manager.create("thread-1") + await manager.set_status(record.run_id, RunStatus.running) + record.task = asyncio.create_task(asyncio.sleep(3600)) + near_expiry = (datetime.now(UTC) + timedelta(milliseconds=50)).isoformat() + record.lease_expires_at = near_expiry + store._runs[record.run_id]["lease_expires_at"] = near_expiry + original_update_lease = store.update_lease + + async def renew_after_timeout_cancellation(*args, **kwargs): + try: + await asyncio.sleep(3600) + except asyncio.CancelledError: + # Simulate a store operation that commits successfully despite the + # caller's deadline cancellation. + pass + return await original_update_lease(*args, **kwargs) + + store.update_lease = renew_after_timeout_cancellation + + await asyncio.wait_for(manager._renew_leases(), timeout=1) + await asyncio.sleep(0) + + row = await store.get(record.run_id) + assert row is not None + assert row["lease_expires_at"] > near_expiry + assert record.lease_expires_at == near_expiry + assert record.ownership_lost is True + assert record.abort_event.is_set() is True + assert record.task.cancelled() + + @pytest.mark.anyio async def test_heartbeat_skips_runs_not_owned_by_this_worker(): """Heartbeat must only renew leases for runs owned by this worker.""" @@ -1779,6 +1920,8 @@ async def test_heartbeat_cancels_task_on_lease_loss(): # Let the event loop process the cancellation (task.cancel() schedules, # doesn't await). await asyncio.sleep(0) + assert record.ownership_lost is True + assert record.abort_event.is_set() is True assert record.task.cancelled() @@ -1850,6 +1993,88 @@ async def test_cancel_action_rollback_finalizes_to_error_in_store(): assert row["error"] == "Rolled back by user" +@pytest.mark.anyio +async def test_peer_reconciliation_fences_late_success_and_completion(): + """A stale owner cannot overwrite a peer's terminal takeover.""" + store = MemoryRunStore() + config = _lease_config(heartbeat_enabled=True, grace_seconds=0) + owner = _make_manager(store=store, worker_id="worker-a", run_ownership_config=config) + peer = _make_manager(store=store, worker_id="worker-b", run_ownership_config=config) + record = await owner.create("thread-1") + await owner.set_status(record.run_id, RunStatus.running) + + expired = (datetime.now(UTC) - timedelta(seconds=1)).isoformat() + record.lease_expires_at = expired + store._runs[record.run_id]["lease_expires_at"] = expired + recovered = await peer.reconcile_orphaned_inflight_runs(error="peer takeover") + + assert [recovered_record.run_id for recovered_record in recovered] == [record.run_id] + assert (await store.get(record.run_id))["status"] == "error" + + await owner.set_status(record.run_id, RunStatus.success) + await owner.update_run_completion(record.run_id, status=record.status.value, total_tokens=1) + + row = await store.get(record.run_id) + assert record.ownership_lost is True + assert record.status == RunStatus.error + assert row["status"] == "error" + assert row["error"] == "peer takeover" + + +@pytest.mark.anyio +async def test_unconfirmed_success_is_fenced_when_heartbeat_is_enabled(): + """A store outage cannot turn an unconfirmed terminal write into local success.""" + store = MemoryRunStore() + manager = _make_manager( + store=store, + worker_id="worker-a", + run_ownership_config=_lease_config(heartbeat_enabled=True), + ) + record = await manager.create("thread-1") + await manager.set_status(record.run_id, RunStatus.running) + + async def fail_status_write(*_args, **_kwargs): + raise OSError("database unreachable") + + store.update_status = fail_status_write + + await manager.set_status(record.run_id, RunStatus.success) + + row = await store.get(record.run_id) + assert record.ownership_lost is True + assert record.status == RunStatus.error + assert record.abort_event.is_set() is True + assert row["status"] == "running" + + +@pytest.mark.anyio +async def test_unconfirmed_staged_success_is_fenced_on_deferred_persistence(): + """Receipt-ordered status persistence retains the success fail-close fence.""" + store = MemoryRunStore() + manager = _make_manager( + store=store, + worker_id="worker-a", + run_ownership_config=_lease_config(heartbeat_enabled=True), + ) + record = await manager.create("thread-1") + await manager.set_status(record.run_id, RunStatus.running) + await manager.set_status(record.run_id, RunStatus.success, persist=False) + + async def fail_status_write(*_args, **_kwargs): + raise OSError("database unreachable") + + store.update_status = fail_status_write + + persisted = await manager.persist_current_status(record.run_id) + + row = await store.get(record.run_id) + assert persisted is False + assert record.ownership_lost is True + assert record.status == RunStatus.error + assert record.abort_event.is_set() is True + assert row["status"] == "running" + + # --------------------------------------------------------------------------- # cancel() claim_for_takeover False → re-read precision # --------------------------------------------------------------------------- diff --git a/backend/tests/test_run_repository.py b/backend/tests/test_run_repository.py index d1d9e0a83..ca797c63a 100644 --- a/backend/tests/test_run_repository.py +++ b/backend/tests/test_run_repository.py @@ -296,6 +296,21 @@ class TestRunRepository: assert updated is False await _cleanup() + @pytest.mark.anyio + async def test_update_run_completion_does_not_replace_terminal_status(self, tmp_path): + """Late completion data cannot rewrite a peer's terminal outcome.""" + repo = await _make_repo(tmp_path) + await repo.put("r1", thread_id="t1", status="running") + await repo.update_status("r1", "error", error="peer takeover") + + updated = await repo.update_run_completion("r1", status="success", total_tokens=1) + + row = await repo.get("r1") + assert updated is False + assert row["status"] == "error" + assert row["error"] == "peer takeover" + await _cleanup() + @pytest.mark.anyio async def test_metadata_preserved(self, tmp_path): repo = await _make_repo(tmp_path) diff --git a/backend/tests/test_run_worker_delivery.py b/backend/tests/test_run_worker_delivery.py index 5cc8123f1..a2ed387ee 100644 --- a/backend/tests/test_run_worker_delivery.py +++ b/backend/tests/test_run_worker_delivery.py @@ -2,7 +2,7 @@ import asyncio from types import SimpleNamespace -from unittest.mock import AsyncMock +from unittest.mock import AsyncMock, MagicMock from uuid import uuid4 import pytest @@ -95,6 +95,44 @@ async def test_delivery_event_presented_zero_without_artifact_production(): assert fetched.status == RunStatus.success +@pytest.mark.anyio +async def test_fenced_worker_leaves_delivery_receipt_to_peer_recovery(): + """A stale worker must not finalize the singleton delivery receipt.""" + run_manager = RunManager() + record = await run_manager.create("thread-lease-lost") + record.ownership_lost = True + record.abort_event.set() + record.status = RunStatus.error + event_store = MemoryRunEventStore() + thread_store = SimpleNamespace( + update_display_name=AsyncMock(), + update_status=AsyncMock(), + ) + on_run_completed = AsyncMock() + agent_factory = MagicMock(side_effect=AssertionError("fenced worker started the agent")) + + await run_agent( + _make_bridge(), + run_manager, + record, + ctx=RunContext( + checkpointer=None, + event_store=event_store, + thread_store=thread_store, + on_run_completed=on_run_completed, + ), + agent_factory=agent_factory, + graph_input={}, + config={}, + ) + + assert await _delivery_events(event_store, record.thread_id, record.run_id) == [] + agent_factory.assert_not_called() + thread_store.update_display_name.assert_not_awaited() + thread_store.update_status.assert_not_awaited() + on_run_completed.assert_not_awaited() + + @pytest.mark.anyio async def test_delivery_event_is_singleton_across_goal_continuations(monkeypatch): run_manager = RunManager() diff --git a/backend/tests/test_run_worker_rollback.py b/backend/tests/test_run_worker_rollback.py index 81813b3dc..abe3f92d1 100644 --- a/backend/tests/test_run_worker_rollback.py +++ b/backend/tests/test_run_worker_rollback.py @@ -3,7 +3,7 @@ import copy from contextlib import suppress from types import SimpleNamespace from typing import Annotated, Any, NotRequired, TypedDict -from unittest.mock import AsyncMock, call, patch +from unittest.mock import AsyncMock, MagicMock, call, patch import pytest from langchain_core.messages import AIMessage, AIMessageChunk, AnyMessage, HumanMessage @@ -2850,3 +2850,45 @@ async def test_worker_finally_block_swallows_helper_exceptions(monkeypatch): assert helper_called.is_set() assert captured_status.get("status") == ("thread-1", "interrupted") bridge.publish_end.assert_awaited_once_with(record.run_id) + + +@pytest.mark.anyio +async def test_worker_skips_execution_and_finalization_after_ownership_loss(): + """A fenced worker closes its stream without starting or finalizing work.""" + run_manager = RunManager() + record = await run_manager.create("thread-lease-lost") + record.ownership_lost = True + record.abort_event.set() + record.status = RunStatus.error + + bridge = SimpleNamespace( + publish=AsyncMock(), + publish_end=AsyncMock(), + cleanup=AsyncMock(), + ) + thread_store = SimpleNamespace( + update_display_name=AsyncMock(), + update_status=AsyncMock(), + ) + on_run_completed = AsyncMock() + agent_factory = MagicMock(side_effect=AssertionError("fenced worker started the agent")) + + await run_agent( + bridge, + run_manager, + record, + ctx=RunContext( + checkpointer=None, + thread_store=thread_store, + on_run_completed=on_run_completed, + ), + agent_factory=agent_factory, + graph_input={"messages": []}, + config={}, + ) + + agent_factory.assert_not_called() + thread_store.update_display_name.assert_not_awaited() + thread_store.update_status.assert_not_awaited() + on_run_completed.assert_not_awaited() + bridge.publish_end.assert_awaited_once_with(record.run_id) diff --git a/config.example.yaml b/config.example.yaml index 1282788e1..38e0bda11 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -1894,6 +1894,11 @@ scheduler: # grace_seconds if your environment cannot keep clocks within a few seconds; # the trade-off is longer recovery latency for genuinely dead workers # (lease_seconds + grace_seconds from last heartbeat to reclaim). +# +# The owning worker itself does NOT use grace_seconds as extra execution time. +# If renewal cannot be confirmed before lease_seconds expires, it cancels the +# local run and suppresses durable finalization; grace_seconds only delays when +# a peer may reclaim the expired row. run_ownership: lease_seconds: 30 # Seconds before a run lease expires if not renewed.