diff --git a/backend/AGENTS.md b/backend/AGENTS.md index 5aa0f108df5..6854612c3db 100644 --- a/backend/AGENTS.md +++ b/backend/AGENTS.md @@ -19,8 +19,26 @@ DeerFlow is a LangGraph-based AI super agent system with a full-stack architectu - Background subagent identity is deliberately split: the provider `tool_call_id` remains the correlation key for `ToolMessage`, `task_*` SSE events, persisted lifecycle events, frontend cards, and the public `ExtensionData.scope_id` contract (stored as `SubagentResult.external_task_id`), while `SubagentExecutor.execute_async()` generates a full server-side `execution_id` for `SubagentResult.task_id`, the process-wide registry, polling, cancellation, timeout handling, and cleanup. Provider IDs are not globally unique across parent runs, so they must never become registry ownership keys; scheduler closures retain their own `SubagentResult` rather than resolving ownership again through the mutable registry. Terminal subagent token usage travels in the current run's `ToolMessage.additional_kwargs` and is attributed from message state, never through a process-global provider-ID cache. - Scheduled-task executions must reuse that same Gateway run lifecycle. The scheduler may decide *when* work runs, but it must dispatch through the existing run path rather than introducing a parallel execution stack. Scheduled launches pass `scheduler.recursion_limit` (default 1000, matching the web UI's `recursion_limit: 1000`, clamped by `max_recursion_limit`) via `launch_scheduled_thread_run`; the value is read from `get_app_config()` at dispatch. - The background scheduler is single-instance by default. `scheduler.multi_instance=true` opts into lease-aware recovery across Gateway instances and requires shared Postgres, `run_ownership.heartbeat_enabled=true`, and `run_events.backend=db`; otherwise startup rejects the configuration. Live scheduled runs are preserved when a peer starts; expired launch claims return to the durable queue, expired run leases are atomically taken over, stale launch writes are fenced by lease ownership, and the Postgres advisory-locked budget makes `max_concurrent_runs` a shared global cap for `launching`/`running` rows. -- Long-running MCP work uses a separate durable task runtime rather than keeping remote task IDs or status polling inside the Agent loop. Explicit `task_toolsets` bind raw submit/status/cancel names; only submit remains Agent-visible, and its wrapper persists the remote handle before returning a local ID. `McpTaskService` claims due rows with leases, resolves a protocol-specific `McpTaskDriver`, and writes normalized snapshots back to `mcp_tasks`; expired leases are the restart-recovery mechanism, and a result returned after expiry or after a cancel request must be discarded even when the owner token still matches. The first cancel request fences an in-flight poll lease, while repeats preserve an active cancellation lease so they cannot issue concurrent remote cancels; cancellation backoff starts when the remote attempt finishes, so a slow timeout cannot consume the retry delay. Cancellation, polling, and notification batches isolate per-task exceptions; an unexpected cancellation/poll failure leaves that record's lease to expire, while notification failures release only the affected lease for retry. Input-required and terminal event snapshots are delivered by idempotent Agent runs and marked delivered only after run success; the trusted notification instruction stays outside the input boundary while the serialized remote event is framed as untrusted data. A busy-thread conflict is normalized back to the service boundary so the queued snapshot coalesces to the latest task event. A missing dispatched run becomes a failed delivery attempt, while transient run-store hydration errors stay distinguishable and retry the same lookup. The database is the source of truth; `ThreadState` receives only a bounded current-thread projection, and display names are neutralized at that model-state boundary. The installed process-local submitter is the source of truth for management-tool exposure; hot `mcp_tasks` edits take effect only after restart, and active skills must explicitly declare the list/cancel business tools. -- MCP notification failures use a consecutive counter separate from the idempotency-key `dispatch_attempt`, capped exponential backoff, latest-event rebuilding before a run launches, and a five-attempt budget before `dead_letter`. A permanently missing/mismatched target thread is dead-lettered immediately instead of being recreated or reclaimed. HTTP and Agent cancellation requests return after the durable cancel fence; the background loop alone owns the potentially slow remote call and retry schedule. The HTTP cancel endpoint rejects requests with 503 when the loop is not running (`mcp_tasks_available` false, e.g. `mcp_tasks.enabled=false` with SQL persistence), so a cancellation is never acknowledged without a worker to perform it. The bounded notification error/count/status join poll and cancellation diagnostics in the task detail API and expanded card. +- Long-running MCP work uses the durable task runtime in `app/mcp_tasks/service.py`: + the database is authoritative, leases provide recovery, and results after lease + expiry or a cancel fence are discarded. HTTP and Agent cancellation return after + the durable fence; the worker owns remote cancellation and retry. Notification + delivery is idempotent and bounded by retry/dead-letter policy, with delivered + state written only after the Agent run succeeds. Detailed claim/release, batch, + delivery, and tool-exposure rules live in the [MCP guide](packages/harness/deerflow/mcp/AGENTS.md); + cancellation, batch, journal, and run-ownership rules live in the [runtime guide](packages/harness/deerflow/runtime/AGENTS.md). +- Cancellation ownership crosses MCP, RunManager, and journal boundaries. Each + cancellation drain and `McpTaskService.stop()` currently use the module-local + `_CANCELLATION_DRAIN_TIMEOUT_SECONDS = 5.0`. When + caller-cancellation draining times out: + the exact asyncio operation task remains retained by the subsystem-owned registry after timeout. + Caller cancellation is re-raised, while normal `McpTaskService.stop()` and + `RunManager.shutdown()` record/log the deadline and return as retained work + continues. `RunManager.shutdown(timeout=5.0)` keeps a caller-provided hard + total budget (the default is overrideable), establishes one absolute deadline, + and gives nested waits only the remaining time. RunManager fencing is local; + durable terminalization may be reconciled by peers/orphans. See the [runtime guide](packages/harness/deerflow/runtime/AGENTS.md) + for the boundary-specific rules and source map. - Scheduled-task dispatch enforces at most one non-terminal occurrence per task through `uq_scheduled_task_run_active` (`task_id WHERE status IN ('queued','launching','running')`). `queued` is durable and survives restart; `launching` carries a short owner/expiry lease and is the only state that may call the normal Gateway launch path; `running` is associated with the durable run. Each occurrence also supplies a stable run-admission idempotency key, so a recovered launch retry reuses the same durable run. A reused-thread `ConflictError` moves `launching` back to `queued`, while non-conflict launch errors become terminal `failed`. Waiting rows do not consume `max_concurrent_runs`; the atomic queue claim enforces the budget. Repeated triggers coalesce on the one active row, and same-thread FIFO treats older `queued`, `launching`, and `running` rows as blockers. The task definition stays immutable for all three active states because queue admission, PATCH/resume, pause, and delete serialize on the parent task row before touching the occurrence row. Pause/delete atomically interrupt existing `queued` rows and reject `launching`/`running` rows; PATCH/resume reject every active state, and mutation errors advertise pause cancellation only for `queued` work. A manual trigger may queue and run while the parent schedule remains paused. Recovery and multi-instance reconciliation lock task/run pairs in deterministic task-id/run-id order and must reconstruct `run_id`, `started_at`, and the live error state before releasing the short launch claim. Launch/failure/timeout bookkeeping changes the occurrence and its parent task in one parent-first transaction so a peer cannot claim the released task between those writes. Queue timeout marks the occurrence failed and advances a scheduled occurrence so it cannot immediately requeue forever; repository write boundaries coerce serialized task timestamps before binding SQL `DateTime` fields. - `extensions_config.json` is written at runtime by the Gateway (`PUT`/`PATCH /api/mcp/config`, the MCP enable switch, skill updates), so the production compose mounts it read-write while `config.yaml` stays `:ro`; Helm copies its ConfigMap seed into a writable home-volume directory before Gateway starts. Every read-modify-write holds both `extensions_config_write_lock` and the sidecar advisory `extensions_config_file_lock`, because the process-local lock alone loses updates across workers. Docker mounts the compose file as its own mount point, and Linux refuses `rename()` over a mount point with `EBUSY` even when the mount is writable — so `atomic_write_extensions_config` keeps the temp-file-plus-rename path and falls back to an in-place overwrite only on `EBUSY`. That fallback is deliberately non-atomic (a crash mid-write truncates the file); it exists because the alternative is a write that can never succeed, and only its first occurrence per target is logged at warning level. Any other `errno` still propagates. Pinned by `tests/test_compose_extensions_config_writable.py`, `tests/test_extensions_config_atomic_write.py`, and `tests/test_helm_extensions_config_writable.py`. diff --git a/backend/app/mcp_tasks/service.py b/backend/app/mcp_tasks/service.py index 6a318cb9238..644f3e88ae6 100644 --- a/backend/app/mcp_tasks/service.py +++ b/backend/app/mcp_tasks/service.py @@ -6,7 +6,8 @@ import socket import uuid from collections.abc import Awaitable, Callable -from dataclasses import replace +from contextvars import ContextVar +from dataclasses import dataclass, replace from datetime import UTC, datetime, timedelta from typing import Any @@ -25,6 +26,7 @@ TaskSubmitRequest, ) from deerflow.persistence.mcp_tasks import DuplicateMcpRemoteTaskError +from deerflow.runtime.cancellation import wait_for_task_until from deerflow.runtime.runs.manager import ConflictError from deerflow.runtime.runs.schemas import RunStatus @@ -33,7 +35,27 @@ _MAX_PERSISTED_ERROR_CHARS = 4_000 _MAX_INPUT_REQUIRED_BYTES = 65_536 _MAX_NOTIFICATION_ATTEMPTS = 5 -_UNTRACKED_TASK_COMPENSATION_WAIT_SECONDS = 5.0 +_CANCELLATION_DRAIN_TIMEOUT_SECONDS = 5.0 + + +@dataclass(slots=True) +class _BatchRecordState: + record: dict[str, Any] + ordinary_release_task: asyncio.Future[Any] | None = None + ordinary_release_terminal: bool = False + cancellation_release_task: asyncio.Future[Any] | None = None + cancellation_release_terminal: bool = False + + +@dataclass(slots=True) +class _BatchState: + cancellation_requested: bool = False + + +_current_batch_record: ContextVar[_BatchRecordState | None] = ContextVar( + "mcp_task_current_batch_record", + default=None, +) def _bound_error(error: str | None) -> str | None: @@ -42,6 +64,21 @@ def _bound_error(error: str | None) -> str | None: return error[:_MAX_PERSISTED_ERROR_CHARS] +def _consume_task_error(task: asyncio.Future[Any]) -> BaseException | None: + try: + return task.exception() + except asyncio.CancelledError as exc: + return exc + + +def _task_has_cancelled_terminal_state(task: asyncio.Future[Any]) -> bool: + if not task.done(): + return False + if task.cancelled(): + return True + return isinstance(_consume_task_error(task), asyncio.CancelledError) + + class McpTaskService: """Persist and poll long-running MCP tasks outside the Agent loop.""" @@ -75,7 +112,10 @@ def __init__( self._get_run = get_run self._lease_owner = f"{socket.gethostname()}:{uuid.uuid4().hex}" self._task: asyncio.Task[None] | None = None - self._compensation_tasks: set[asyncio.Task[Any]] = set() + self._stopping_task: asyncio.Task[None] | None = None + self._stop_deadline: float | None = None + self._stop_timeout_logged = False + self._compensation_tasks: set[asyncio.Future[Any]] = set() self._stop = asyncio.Event() @property @@ -86,6 +126,112 @@ def drivers(self) -> McpTaskDriverRegistry: def tracking_degraded_after_errors(self) -> int: return self._tracking_degraded_after_errors + def _observe_batch_release_task( + self, + state: _BatchRecordState, + task: asyncio.Future[Any], + *, + ordinary: bool, + action: str, + ) -> None: + terminal_field = "ordinary_release_terminal" if ordinary else "cancellation_release_terminal" + if getattr(state, terminal_field) or not task.done(): + return + setattr(state, terminal_field, True) + error = _consume_task_error(task) + if error is None: + return + self._log_batch_release_error( + error, + action=action, + task_id=state.record.get("id"), + ) + + @staticmethod + def _log_batch_release_error(error: BaseException, *, action: str, task_id: Any) -> None: + logger.error( + "MCP task batch release failed (%s, task_id=%s): %s", + action, + task_id, + error, + exc_info=(type(error), error, error.__traceback__), + ) + + @staticmethod + def _log_claim_error(error: BaseException, *, action: str) -> None: + logger.error( + "MCP task claim operation failed (%s, task_id=batch): %s", + action, + error, + exc_info=(type(error), error, error.__traceback__), + ) + + def _track_batch_release_task( + self, + state: _BatchRecordState, + task: asyncio.Future[Any], + *, + ordinary: bool, + action: str, + ) -> None: + task.add_done_callback( + lambda completed: self._observe_batch_release_task( + state, + completed, + ordinary=ordinary, + action=action, + ) + ) + + async def _release_ordinary_batch_record( + self, + record: dict[str, Any], + *, + release: Callable[[], Awaitable[Any]], + action: str, + ) -> None: + state = _current_batch_record.get() + if state is None: + await release() + return + + if state.cancellation_release_task is not None: + self._observe_batch_release_task( + state, + state.cancellation_release_task, + ordinary=False, + action="cancellation release", + ) + return + + task = state.ordinary_release_task + if task is None: + task = asyncio.create_task( + release(), + name=f"mcp-{action.replace(' ', '-')}-ordinary-release-{record.get('id', 'unknown')}", + ) + state.ordinary_release_task = task + self._track_batch_release_task( + state, + task, + ordinary=True, + action=action, + ) + try: + await asyncio.shield(task) + except asyncio.CancelledError: + caller_cancelling = asyncio.current_task().cancelling() + self._observe_batch_release_task(state, task, ordinary=True, action=action) + release_cancelled = state.ordinary_release_terminal or _task_has_cancelled_terminal_state(task) + if release_cancelled and not caller_cancelling: + return + raise + except Exception: + self._observe_batch_release_task(state, task, ordinary=True, action=action) + raise + else: + self._observe_batch_release_task(state, task, ordinary=True, action=action) + async def submit( self, *, @@ -174,12 +320,9 @@ async def _cancel_untracked_task( ) self._compensation_tasks.add(compensation) - def finalize(task: asyncio.Task[Any]) -> None: + def finalize(task: asyncio.Future[Any]) -> None: self._compensation_tasks.discard(task) - try: - error = task.exception() - except asyncio.CancelledError as exc: - error = exc + error = _consume_task_error(task) if error is None: return logger.error( @@ -193,46 +336,295 @@ def finalize(task: asyncio.Task[Any]) -> None: compensation.add_done_callback(finalize) loop = asyncio.get_running_loop() - deadline = loop.time() + _UNTRACKED_TASK_COMPENSATION_WAIT_SECONDS - while not compensation.done(): - remaining = deadline - loop.time() - if remaining <= 0: - logger.warning( - "Timed out after %.1f seconds waiting for untracked MCP task compensation after %s; cancellation continues in the background (task_id=%s, driver=%s, remote_task_id=%s)", - _UNTRACKED_TASK_COMPENSATION_WAIT_SECONDS, - reason, - task_reference.local_task_id, - driver_name, - task_reference.remote_task_id, + deadline = loop.time() + _CANCELLATION_DRAIN_TIMEOUT_SECONDS + if not await wait_for_task_until(compensation, deadline=deadline): + logger.warning( + "Timed out after %.1f seconds waiting for untracked MCP task compensation after %s; cancellation continues in the background (task_id=%s, driver=%s, remote_task_id=%s)", + _CANCELLATION_DRAIN_TIMEOUT_SECONDS, + reason, + task_reference.local_task_id, + driver_name, + task_reference.remote_task_id, + ) + + def _track_compensation_task(self, task: asyncio.Future[Any], *, action: str, task_id: str) -> None: + if task in self._compensation_tasks: + return + self._compensation_tasks.add(task) + + def finalize(completed: asyncio.Future[Any]) -> None: + self._compensation_tasks.discard(completed) + error = _consume_task_error(completed) + if error is None: + return + logger.error( + "MCP task cancellation operation failed (%s, task_id=%s): %s", + action, + task_id, + error, + exc_info=(type(error), error, error.__traceback__), + ) + + task.add_done_callback(finalize) + + async def _drain_cancellation_task( + self, + task: asyncio.Future[Any], + *, + action: str, + task_id: str, + deadline: float, + ) -> tuple[bool, Any]: + if not await wait_for_task_until(task, deadline=deadline): + self._track_compensation_task(task, action=action, task_id=task_id) + logger.warning( + "Timed out after %.1f seconds waiting for MCP task cancellation operation; it continues in the background (%s, task_id=%s)", + _CANCELLATION_DRAIN_TIMEOUT_SECONDS, + action, + task_id, + ) + return False, None + + error = _consume_task_error(task) + if error is not None: + logger.error( + "MCP task cancellation operation failed (%s, task_id=%s): %s", + action, + task_id, + error, + exc_info=(type(error), error, error.__traceback__), + ) + return False, None + return True, task.result() + + async def _drain_cancellation_compensation( + self, + compensation: Awaitable[Any], + *, + action: str, + task_id: str, + ) -> tuple[bool, Any]: + task = asyncio.ensure_future(compensation) + deadline = asyncio.get_running_loop().time() + _CANCELLATION_DRAIN_TIMEOUT_SECONDS + return await self._drain_cancellation_task( + task, + action=action, + task_id=task_id, + deadline=deadline, + ) + + async def _release_owned_batch_record( + self, + state: _BatchRecordState, + *, + release: Callable[[dict[str, Any]], Awaitable[None]], + ) -> None: + ordinary_task = state.ordinary_release_task + if ordinary_task is not None: + self._observe_batch_release_task( + state, + ordinary_task, + ordinary=True, + action="ordinary retry release", + ) + return + + task = state.cancellation_release_task + if task is None: + task = asyncio.create_task( + release(state.record), + name=f"mcp-cancellation-release-{state.record.get('id', 'unknown')}", + ) + state.cancellation_release_task = task + self._track_batch_release_task( + state, + task, + ordinary=False, + action="cancellation release", + ) + try: + await asyncio.shield(task) + except asyncio.CancelledError: + self._observe_batch_release_task( + state, + task, + ordinary=False, + action="cancellation release", + ) + except Exception: + self._observe_batch_release_task( + state, + task, + ordinary=False, + action="cancellation release", + ) + else: + self._observe_batch_release_task( + state, + task, + ordinary=False, + action="cancellation release", + ) + + async def _finish_cancelled_batch( + self, + supervisor: asyncio.Task[list[Any]], + children: list[asyncio.Task[Any]], + states: list[_BatchRecordState], + *, + release: Callable[[dict[str, Any]], Awaitable[None]], + action: str, + ) -> None: + # The handoff owns both the supervisor and every release task. Keeping + # all of them in this frame lets a timed-out handoff finish safely in + # the background without starting a second release. + async def release_uncompleted(state: _BatchRecordState) -> None: + if state.ordinary_release_task is not None: + try: + await state.ordinary_release_task + except asyncio.CancelledError: + pass + except Exception: + pass + self._observe_batch_release_task( + state, + state.ordinary_release_task, + ordinary=True, + action="ordinary retry release", ) return + await self._release_owned_batch_record(state, release=release) + + release_tasks = [ + asyncio.create_task( + release_uncompleted(state), + name=f"mcp-{action.replace(' ', '-')}-release-{index}-{state.record.get('id', 'unknown')}", + ) + for index, state in enumerate(states) + ] + results = await asyncio.gather(supervisor, *release_tasks, return_exceptions=True) + supervisor_result = results[0] + if isinstance(supervisor_result, BaseException): + logger.error( + "MCP task batch supervisor failed during cancellation handoff (action=%s): %s", + action, + supervisor_result, + exc_info=(type(supervisor_result), supervisor_result, supervisor_result.__traceback__), + ) + for state, child in zip(states, children, strict=True): + if child.done(): + error = _consume_task_error(child) + if error is None or isinstance(error, asyncio.CancelledError): + continue + failure_action = "cancellation" if action == "cancel" else action + lease_suffix = "; the lease will expire for recovery" if action in {"poll", "cancel"} else "" + logger.error( + "Unexpected MCP task %s failure (task_id=%s)%s", + failure_action, + state.record.get("id"), + lease_suffix, + exc_info=(type(error), error, error.__traceback__), + ) + + async def _run_claimed_batch( + self, + records: list[dict[str, Any]], + *, + operation: Callable[[dict[str, Any]], Awaitable[Any]], + release: Callable[[dict[str, Any]], Awaitable[None]], + action: str, + ) -> tuple[list[_BatchRecordState], list[Any]]: + states = [_BatchRecordState(record) for record in records] + batch_state = _BatchState() + parent_task = asyncio.current_task() + + async def run_one(state: _BatchRecordState) -> Any: + context_token = _current_batch_record.set(state) try: - await asyncio.wait({compensation}, timeout=remaining) + if parent_task is not None and parent_task.cancelling(): + batch_state.cancellation_requested = True + if batch_state.cancellation_requested: + return None + return await operation(state.record) except asyncio.CancelledError: - # Repeated caller cancellation does not propagate through - # asyncio.wait() to the compensation task. Keep waiting only - # until the original deadline. - continue + await self._release_owned_batch_record(state, release=release) + raise + finally: + _current_batch_record.reset(context_token) + if batch_state.cancellation_requested: + await self._release_owned_batch_record(state, release=release) + + task_prefix = action.replace(" ", "-") + tasks = [ + asyncio.create_task( + run_one(state), + name=f"mcp-{task_prefix}-{index}-{state.record.get('id', 'unknown')}", + ) + for index, state in enumerate(states) + ] + + async def supervise() -> list[Any]: + return await asyncio.gather(*tasks, return_exceptions=True) + + supervisor = asyncio.create_task( + supervise(), + name=f"mcp-{task_prefix}-supervisor", + ) + try: + results = await asyncio.shield(supervisor) + except asyncio.CancelledError as original_cancel: + batch_state.cancellation_requested = True + for task in tasks: + task.cancel() + handoff = asyncio.create_task( + self._finish_cancelled_batch( + supervisor, + tasks, + states, + release=release, + action=action, + ), + name=f"mcp-{task_prefix}-cancellation-handoff", + ) + await self._drain_cancellation_task( + handoff, + action=f"finish {action} batch handoff", + task_id="batch", + deadline=asyncio.get_running_loop().time() + _CANCELLATION_DRAIN_TIMEOUT_SECONDS, + ) + # A task that was cancelled before entering this handler stores a + # special cancelled state; raising that same exception object can + # make asyncio discard its message when the task is awaited. + # Recreate it with the first cancellation's args instead. + raise asyncio.CancelledError(*original_cancel.args) + + return states, results async def run_once(self, *, now: datetime) -> None: await self._run_cancellations(now=now) - claimed = await self._repository.claim_due_tasks( - now=now, - lease_owner=self._lease_owner, - lease_seconds=self._lease_seconds, - limit=self._max_concurrent_polls, + claimed = await self._claim_with_cancellation_release( + self._repository.claim_due_tasks( + now=now, + lease_owner=self._lease_owner, + lease_seconds=self._lease_seconds, + limit=self._max_concurrent_polls, + ), + action="poll claim", + release=self._release_poll_after_cancellation, ) if claimed: - results = await asyncio.gather( - *(self._poll_one(task, now=now) for task in claimed), - return_exceptions=True, + states, results = await self._run_claimed_batch( + claimed, + operation=lambda record: self._poll_one_claimed(record, now=now), + release=self._release_poll_after_cancellation, + action="poll", ) - for record, result in zip(claimed, results, strict=True): + for state, result in zip(states, results, strict=True): if isinstance(result, BaseException): logger.error( "Unexpected MCP task poll failure (task_id=%s); the lease will expire for recovery", - record.get("id"), + state.record.get("id"), exc_info=(type(result), result, result.__traceback__), ) @@ -299,26 +691,39 @@ async def _run_cancellations(self, *, now: datetime) -> None: claim = getattr(self._repository, "claim_cancel_requests", None) if claim is None: return - records = await claim( - now=now, - lease_owner=self._lease_owner, - lease_seconds=self._lease_seconds, - limit=self._max_concurrent_polls, + records = await self._claim_with_cancellation_release( + claim( + now=now, + lease_owner=self._lease_owner, + lease_seconds=self._lease_seconds, + limit=self._max_concurrent_polls, + ), + action="cancel claim", + release=self._release_cancel_after_cancellation, ) if records: - results = await asyncio.gather( - *(self._cancel_one(record) for record in records), - return_exceptions=True, + states, results = await self._run_claimed_batch( + records, + operation=self._cancel_one_claimed, + release=self._release_cancel_after_cancellation, + action="cancel", ) - for record, result in zip(records, results, strict=True): + for state, result in zip(states, results, strict=True): if isinstance(result, BaseException): logger.error( "Unexpected MCP task cancellation failure (task_id=%s); the lease will expire for recovery", - record.get("id"), + state.record.get("id"), exc_info=(type(result), result, result.__traceback__), ) async def _cancel_one(self, record: dict[str, Any]) -> None: + try: + await self._cancel_one_claimed(record) + except asyncio.CancelledError: + await self._release_cancel_after_cancellation(record) + raise + + async def _cancel_one_claimed(self, record: dict[str, Any]) -> None: driver_name = str(record.get("driver_name") or "") driver = self._drivers.get(driver_name) try: @@ -343,52 +748,59 @@ async def _cancel_one(self, record: dict[str, Any]) -> None: attempts = max(0, int(record.get("cancel_attempt_count") or 1) - 1) retry_seconds = min(self._poll_interval_seconds * (2 ** min(attempts, 16)), self._max_poll_backoff_seconds) failed_at = datetime.now(UTC) - await self._repository.release_cancel_claim( - record["id"], - lease_owner=self._lease_owner, - next_cancel_at=failed_at + timedelta(seconds=retry_seconds), - error=_bound_error(str(exc) or type(exc).__name__), + retry_error = _bound_error(str(exc) or type(exc).__name__) + await self._release_ordinary_batch_record( + record, + release=lambda: self._repository.release_cancel_claim( + record["id"], + lease_owner=self._lease_owner, + next_cancel_at=failed_at + timedelta(seconds=retry_seconds), + error=retry_error, + ), + action="release cancel retry", ) async def _run_notifications(self, *, now: datetime) -> None: if self._launch_notification is None or self._get_run is None: return - records = await self._repository.claim_notification_work( - now=now, - lease_owner=self._lease_owner, - lease_seconds=self._lease_seconds, - limit=self._max_concurrent_polls, - tracking_degraded_after_errors=self._tracking_degraded_after_errors, + records = await self._claim_with_cancellation_release( + self._repository.claim_notification_work( + now=now, + lease_owner=self._lease_owner, + lease_seconds=self._lease_seconds, + limit=self._max_concurrent_polls, + tracking_degraded_after_errors=self._tracking_degraded_after_errors, + ), + action="notification claim", + release=self._release_notification_after_cancellation, ) if records: - results = await asyncio.gather( - *(self._notify_one(record, now=now) for record in records), - return_exceptions=True, + states, results = await self._run_claimed_batch( + records, + operation=lambda record: self._notify_one_claimed(record, now=now), + release=self._release_notification_after_cancellation, + action="notification", ) - for record, result in zip(records, results, strict=True): - if not isinstance(result, BaseException): + for state, result in zip(states, results, strict=True): + if not isinstance(result, BaseException) or isinstance(result, asyncio.CancelledError): continue + record = state.record error = _bound_error(str(result) or type(result).__name__) or type(result).__name__ logger.error( "Unexpected MCP task notification failure (task_id=%s)", record.get("id"), exc_info=(type(result), result, result.__traceback__), ) - try: - await self._repository.release_notification_lease( - record["id"], - lease_owner=self._lease_owner, - next_notification_at=now + timedelta(seconds=self._notification_retry_seconds(record)), - error=error, - count_failure=True, - ) - except Exception: # noqa: BLE001 - retain the original task-scoped failure - logger.exception( - "Failed to release MCP task notification lease (task_id=%s)", - record.get("id"), - ) + await self._release_notification_failure(record, now=now, error=error) async def _notify_one(self, record: dict[str, Any], *, now: datetime) -> None: + try: + await self._notify_one_claimed(record, now=now) + except asyncio.CancelledError: + await self._release_notification_after_cancellation(record) + raise + + async def _notify_one_claimed(self, record: dict[str, Any], *, now: datetime) -> None: task_id = record["id"] dispatch_version = int(record.get("dispatch_version") or 0) notification_attempts = max(0, int(record.get("notification_attempt_count") or 0)) @@ -470,22 +882,32 @@ async def _notify_one(self, record: dict[str, Any], *, now: datetime) -> None: ) return except ConflictError as exc: - await self._repository.release_notification_claim( - task_id, - lease_owner=self._lease_owner, - next_notification_at=now + timedelta(seconds=self._poll_interval_seconds), - error=_bound_error(str(exc)), - replace_with_latest=True, + retry_error = _bound_error(str(exc)) + await self._release_ordinary_batch_record( + record, + release=lambda: self._repository.release_notification_claim( + task_id, + lease_owner=self._lease_owner, + next_notification_at=now + timedelta(seconds=self._poll_interval_seconds), + error=retry_error, + replace_with_latest=True, + ), + action="release notification conflict retry", ) return except Exception as exc: # noqa: BLE001 - retry the same idempotency key - await self._repository.release_notification_claim( - task_id, - lease_owner=self._lease_owner, - next_notification_at=now + timedelta(seconds=self._notification_retry_seconds(record)), - error=_bound_error(str(exc) or type(exc).__name__), - replace_with_latest=True, - count_failure=True, + retry_error = _bound_error(str(exc) or type(exc).__name__) + await self._release_ordinary_batch_record( + record, + release=lambda: self._repository.release_notification_claim( + task_id, + lease_owner=self._lease_owner, + next_notification_at=now + timedelta(seconds=self._notification_retry_seconds(record)), + error=retry_error, + replace_with_latest=True, + count_failure=True, + ), + action="release notification retry", ) return await self._repository.mark_notification_dispatched( @@ -504,13 +926,219 @@ def _notification_retry_seconds(self, record: dict[str, Any]) -> int: ) async def _poll_one(self, record: dict, *, now: datetime) -> None: + try: + await self._poll_one_claimed(record, now=now) + except asyncio.CancelledError: + await self._release_poll_after_cancellation(record) + raise + + async def _claim_with_cancellation_release( + self, + claim: Awaitable[list[dict[str, Any]]], + *, + action: str, + release: Callable[[dict[str, Any]], Awaitable[None]], + ) -> list[dict[str, Any]]: + claim_task = asyncio.ensure_future(claim) + try: + return await asyncio.shield(claim_task) + except asyncio.CancelledError: + caller_cancelling = asyncio.current_task().cancelling() + claim_cancelled = _task_has_cancelled_terminal_state(claim_task) + if claim_cancelled and not caller_cancelling: + error = _consume_task_error(claim_task) + if error is not None: + self._log_claim_error(error, action=action) + return [] + handoff = asyncio.create_task( + self._finish_cancelled_claim_handoff( + claim_task, + action=action, + release=release, + ), + name=f"mcp-cancelled-{action.replace(' ', '-')}-handoff", + ) + loop = asyncio.get_running_loop() + await self._drain_cancellation_task( + handoff, + action=f"finish {action} handoff", + task_id="batch", + deadline=loop.time() + _CANCELLATION_DRAIN_TIMEOUT_SECONDS, + ) + raise + + async def _finish_cancelled_claim_handoff( + self, + claim_task: asyncio.Future[list[dict[str, Any]]], + *, + action: str, + release: Callable[[dict[str, Any]], Awaitable[None]], + ) -> None: + try: + records = await claim_task + except asyncio.CancelledError: + logger.error("MCP task claim operation was cancelled (%s, task_id=batch)", action) + return + except Exception as exc: # noqa: BLE001 - claim recovery is best-effort + logger.error( + "MCP task claim operation failed (%s, task_id=batch): %s", + action, + exc, + exc_info=(type(exc), exc, exc.__traceback__), + ) + return + if records: + await self._release_claimed_records(records, release=release) + + async def _release_claimed_records( + self, + records: list[dict[str, Any]], + *, + release: Callable[[dict[str, Any]], Awaitable[None]], + ) -> None: + async def release_one(record: dict[str, Any]) -> None: + try: + await release(record) + except asyncio.CancelledError: + raise + except Exception: # noqa: BLE001 - release every record in the claimed batch + logger.exception( + "Unexpected MCP task claim release failure (task_id=%s)", + record.get("id"), + ) + + release_tasks = [ + asyncio.create_task( + release_one(record), + name=f"mcp-release-claimed-{record.get('id', 'unknown')}", + ) + for record in records + ] + completion = asyncio.gather(*release_tasks, return_exceptions=True) + deadline = asyncio.get_running_loop().time() + _CANCELLATION_DRAIN_TIMEOUT_SECONDS + if not await wait_for_task_until(completion, deadline=deadline): + self._track_compensation_task( + completion, + action="release claimed MCP task batch", + task_id="batch", + ) + logger.warning( + "Timed out after %.1f seconds waiting for MCP task claim releases; they continue in the background", + _CANCELLATION_DRAIN_TIMEOUT_SECONDS, + ) + return + results = completion.result() + for record, result in zip(records, results, strict=True): + if isinstance(result, asyncio.CancelledError): + logger.error( + "MCP task claim release was cancelled (task_id=%s)", + record.get("id"), + ) + + async def _release_notification_failure( + self, + record: dict[str, Any], + *, + now: datetime, + error: str, + ) -> None: + task = asyncio.create_task( + self._repository.release_notification_lease( + record["id"], + lease_owner=self._lease_owner, + next_notification_at=now + timedelta(seconds=self._notification_retry_seconds(record)), + error=error, + count_failure=True, + ), + name=f"mcp-release-notification-failure-{record.get('id', 'unknown')}", + ) + try: + await asyncio.shield(task) + except asyncio.CancelledError: + caller_cancelling = asyncio.current_task().cancelling() + release_cancelled = _task_has_cancelled_terminal_state(task) + if release_cancelled and not caller_cancelling: + error = _consume_task_error(task) + if error is not None: + self._log_batch_release_error( + error, + action="release notification failure", + task_id=record["id"], + ) + return + await self._drain_cancellation_task( + task, + action="release notification failure", + task_id=record["id"], + deadline=asyncio.get_running_loop().time() + _CANCELLATION_DRAIN_TIMEOUT_SECONDS, + ) + raise + except Exception: # noqa: BLE001 - retain the task-scoped failure + logger.exception( + "Failed to release MCP task notification lease (task_id=%s)", + record.get("id"), + ) + + async def _release_cancel_after_cancellation(self, record: dict[str, Any]) -> None: + await self._drain_cancellation_compensation( + self._repository.release_cancel_claim( + record["id"], + lease_owner=self._lease_owner, + next_cancel_at=datetime.now(UTC), + error="cancelled", + ), + action="release cancel claim", + task_id=record["id"], + ) + + async def _release_notification_after_cancellation(self, record: dict[str, Any]) -> None: + task_id = record["id"] + if record.get("notification_status") == "dispatched": + compensation = self._repository.release_notification_lease( + task_id, + lease_owner=self._lease_owner, + next_notification_at=datetime.now(UTC), + error="cancelled", + count_failure=False, + ) + action = "release dispatched notification lease" + else: + compensation = self._repository.release_notification_claim( + task_id, + lease_owner=self._lease_owner, + next_notification_at=datetime.now(UTC), + error="cancelled", + replace_with_latest=False, + ) + action = "release notification claim" + await self._drain_cancellation_compensation( + compensation, + action=action, + task_id=task_id, + ) + + async def _release_poll_after_cancellation(self, record: dict[str, Any]) -> None: + await self._drain_cancellation_compensation( + self._repository.release_poll_claim_after_cancellation( + record["id"], + lease_owner=self._lease_owner, + ), + action="release poll claim", + task_id=record["id"], + ) + + async def _poll_one_claimed(self, record: dict, *, now: datetime) -> None: driver_name = str(record.get("driver_name") or "") driver = self._drivers.get(driver_name) if driver is None: - await self._release_after_error( + await self._release_ordinary_batch_record( record, - now=now, - error=f"No MCP task driver registered as {driver_name!r}", + release=lambda: self._release_after_error( + record, + now=now, + error=f"No MCP task driver registered as {driver_name!r}", + ), + action="release poll retry", ) return @@ -537,7 +1165,16 @@ async def _poll_one(self, record: dict, *, now: datetime) -> None: driver_name, exc_info=True, ) - await self._release_after_error(record, now=polled_at, error=str(exc) or type(exc).__name__) + retry_error = str(exc) or type(exc).__name__ + await self._release_ordinary_batch_record( + record, + release=lambda: self._release_after_error( + record, + now=polled_at, + error=retry_error, + ), + action="release poll retry", + ) return polled_at = datetime.now(UTC) @@ -644,20 +1281,61 @@ async def start(self) -> None: if self._task is not None: return self._stop.clear() - self._task = asyncio.create_task(self._run_loop(), name="deerflow-mcp-task-poller") + self._stopping_task = None + self._stop_deadline = None + self._stop_timeout_logged = False + task = asyncio.create_task(self._run_loop(), name="deerflow-mcp-task-poller") + self._task = task + task.add_done_callback(self._poller_done) + + def _poller_done(self, task: asyncio.Task[None]) -> None: + if self._task is task: + self._task = None + self._stopping_task = None + self._stop_deadline = None + self._stop_timeout_logged = False + error = _consume_task_error(task) + if error is None or isinstance(error, asyncio.CancelledError): + return + logger.error( + "MCP task poller failed: %s", + error, + exc_info=(type(error), error, error.__traceback__), + ) + + def _log_stop_timeout(self, task: asyncio.Task[None]) -> None: + if self._stopping_task is not task or self._stop_timeout_logged: + return + self._stop_timeout_logged = True + logger.warning( + "Timed out after %.1f seconds waiting for MCP task poller cleanup; cleanup continues in the background", + _CANCELLATION_DRAIN_TIMEOUT_SECONDS, + ) async def stop(self) -> None: task = self._task if task is None: return + loop = asyncio.get_running_loop() + if self._stopping_task is not task: + self._stopping_task = task + self._stop_deadline = loop.time() + _CANCELLATION_DRAIN_TIMEOUT_SECONDS + self._stop_timeout_logged = False + task.cancel() self._stop.set() - task.cancel() + deadline = self._stop_deadline + assert deadline is not None try: - await task + done, _ = await asyncio.wait( + {task}, + timeout=max(0.0, deadline - loop.time()), + ) except asyncio.CancelledError: - pass - finally: - self._task = None + if not await wait_for_task_until(task, deadline=deadline): + self._log_stop_timeout(task) + raise + if task not in done: + self._log_stop_timeout(task) async def _run_loop(self) -> None: while not self._stop.is_set(): diff --git a/backend/packages/harness/deerflow/mcp/AGENTS.md b/backend/packages/harness/deerflow/mcp/AGENTS.md index 275d9416903..f05772263aa 100644 --- a/backend/packages/harness/deerflow/mcp/AGENTS.md +++ b/backend/packages/harness/deerflow/mcp/AGENTS.md @@ -1,7 +1,7 @@ ### MCP System (`packages/harness/deerflow/mcp/`) - Uses `langchain-mcp-adapters` `MultiServerMCPClient` for multi-server management -- **Long-running task foundation**: `mcp/tasks/` defines the protocol-neutral `McpTaskDriver` contract and normalized `TaskSnapshot` states (`submitted`, `working`, `input_required`, `completed`, `failed`, `cancelled`). A driver-supplied `poll_after_seconds` must be a finite positive number, validated at the `TaskSnapshot` boundary so every driver is held to the same invariant rather than each one guarding the consumer that turns the interval into a `timedelta`. `persistence/mcp_tasks/` owns the durable remote-handle mapping, poll schedule, notification state, lease owner, and separate consecutive poll/delivery error counters; `app/mcp_tasks/McpTaskService` performs status, cancellation, and notification work outside the Agent/LLM loop. Notification retries keep their idempotency attempt separate from the delivery-failure count, use capped exponential backoff, and stop after five failures; strict existing-thread admission dead-letters a deleted/mismatched target immediately. A status result is applied only when the worker still owns an unexpired lease, so a stale result cannot be written after expiry even before another worker reclaims the row. Poll timestamps and retry schedules are based on the remote call's completion time rather than the scan start. If submission succeeds but persistence fails or the caller is cancelled while persistence is in flight, the service best-effort cancels the remote task and preserves the original error or cancellation if that compensation also fails. The exact `uq_mcp_tasks_user_server_remote` conflict is different: an existing durable row already owns the remote handle, so the conflict surfaces without cancelling that tracked task. Unexpected per-task poll failures are isolated from sibling claims and remain recoverable through lease expiry; Gateway shutdown cancels the poller so a hung external status call cannot block process exit. `input_required` and terminal states stop polling and become `notification_status=pending` for later Agent/UI delivery. Durable recovery requires a SQL database backend (`sqlite` or `postgres`); the in-memory backend leaves the repository/service unavailable. The runtime is startup-configured by `mcp_tasks` and disabled by default until a concrete driver is registered; this foundation does not alter ordinary MCP tool behavior on its own. +- **Long-running task foundation**: `mcp/tasks/` defines the protocol-neutral `McpTaskDriver` contract and normalized `TaskSnapshot` states (`submitted`, `working`, `input_required`, `completed`, `failed`, `cancelled`). A driver-supplied `poll_after_seconds` must be a finite positive number, validated at the `TaskSnapshot` boundary so every driver is held to the same invariant rather than each one guarding the consumer that turns the interval into a `timedelta`. `persistence/mcp_tasks/` owns the durable remote-handle mapping, poll schedule, notification state, lease owner, and separate consecutive poll/delivery error counters; `app/mcp_tasks/McpTaskService` performs status, cancellation, and notification work outside the Agent/LLM loop. Notification retries keep their idempotency attempt separate from the delivery-failure count, use capped exponential backoff, and stop after five failures; strict existing-thread admission dead-letters a deleted/mismatched target immediately. A status result is applied only when the worker still owns an unexpired lease, so a stale result cannot be written after expiry even before another worker reclaims the row. Poll timestamps and retry schedules are based on the remote call's completion time rather than the scan start. If submission succeeds but persistence fails or the caller is cancelled while persistence is in flight, the service best-effort cancels the remote task and preserves the original error or cancellation if that compensation also fails. The exact `uq_mcp_tasks_user_server_remote` conflict is different: an existing durable row already owns the remote handle, so the conflict surfaces without cancelling that tracked task. Unexpected per-task poll failures are isolated from sibling claims and remain recoverable through lease expiry; Gateway shutdown cancels the poller so a hung external status call cannot block process exit. Cancelling an in-flight poll releases only the owner-fenced lease and preserves its preclaim schedule and poll-failure state; real poll failures retain exponential backoff and tracking-degradation behavior. `input_required` and terminal states stop polling and become `notification_status=pending` for later Agent/UI delivery. Durable recovery requires a SQL database backend (`sqlite` or `postgres`); the in-memory backend leaves the repository/service unavailable. The runtime is startup-configured by `mcp_tasks` and disabled by default until a concrete driver is registered; this foundation does not alter ordinary MCP tool behavior on its own. - **Runtime availability boundary**: the installed process-local submitter is the source of truth for durable task-management tool exposure. `mcp_tasks` is startup-only; changing it on disk does not alter the live toolset until the Gateway restarts. - **Long-running ordinary task driver**: `extensions_config.json -> mcpServers..task_toolsets` binds exact raw submit/status/cancel names; one raw tool may occupy only one role across that server's groups. `mcp/tools.py` hides status/cancel and replaces submit with a wrapper that returns only the local task ID after persistence. `ordinary.py` reads only MCP `structuredContent`, maps remote `running` to `working`, and treats `error_code=task_not_found` or malformed structured output as permanent failure. A status call with `isError=true` is a retryable call failure: the first text content block is retained as a bounded diagnostic, while a permanent remote-task outcome must arrive in a normal result with structured `status=failed`. `task_tool_caller.py` restores the same `(server_name, user_id:thread_id)` stdio session scope; HTTP/SSE calls remain ephemeral, apply `session_init_timeout` to initialization and `tool_call_timeout` to task calls, and support server-level OAuth refresh outside an Agent run. `McpTaskService` exponentially backs off transient status/cancel errors without a maximum attempt count, derives API `tracking_degraded` from the consecutive-error threshold, keeps `input_required` on a slower poll, and caps finite positive remote poll hints at 24 hours. Task-enabled server runtime/binding configuration and `mcpInterceptors` are frozen to the Gateway startup snapshot; hot drift fails clearly before tool discovery can diverge from background calls, while presentation-only fields and non-task servers remain reloadable. Configured task toolsets fail startup when the runtime is disabled or persistence is memory. Users still cannot submit an answer back to an `input_required` remote task. - **Durable task payload bounds**: persisted task errors are capped at 4,000 characters. `input_required` and `result_artifact` must each serialize as valid JSON within 64 KiB; an invalid or oversized payload becomes a permanent protocol failure rather than being truncated and changing its semantics. Remote task IDs/task names are limited to 255 characters and task-enabled server names to 128, matching the SQL schema; an oversized submitted remote ID is rejected only after the Service has the handle so compensation cancellation still runs. Oversized results retain the existing bounded preview/truncation/artifact behavior. diff --git a/backend/packages/harness/deerflow/persistence/mcp_tasks/sql.py b/backend/packages/harness/deerflow/persistence/mcp_tasks/sql.py index db9a9b2bf9a..7ee4c8681d7 100644 --- a/backend/packages/harness/deerflow/persistence/mcp_tasks/sql.py +++ b/backend/packages/harness/deerflow/persistence/mcp_tasks/sql.py @@ -304,6 +304,27 @@ async def release_claim( await session.commit() return True + async def release_poll_claim_after_cancellation( + self, + task_id: str, + *, + lease_owner: str, + ) -> bool: + """Release a cancelled poll's lease without recording a poll failure.""" + stmt = ( + update(McpTaskRow) + .where(McpTaskRow.id == task_id, McpTaskRow.lease_owner == lease_owner) + .values( + lease_owner=None, + lease_expires_at=None, + updated_at=datetime.now(UTC), + ) + ) + async with self._sf() as session: + result = await session.execute(stmt) + await session.commit() + return bool(result.rowcount) + async def request_cancel( self, task_id: str, diff --git a/backend/packages/harness/deerflow/runtime/AGENTS.md b/backend/packages/harness/deerflow/runtime/AGENTS.md index 2c27ba96740..27340f1bc30 100644 --- a/backend/packages/harness/deerflow/runtime/AGENTS.md +++ b/backend/packages/harness/deerflow/runtime/AGENTS.md @@ -109,3 +109,80 @@ PYTHONPATH=. uv run python scripts/benchmark/checkpoint/bench_production.py \ PYTHONPATH=. uv run python scripts/benchmark/checkpoint/summarize_production.py \ /tmp/production-bench.jsonl ``` + +### Bounded Cancellation Drains + +Source map (paths relative to `backend/`): +`packages/harness/deerflow/runtime/cancellation.py` provides the shared absolute- +deadline wait primitive; `app/mcp_tasks/service.py` owns MCP claim, batch, +release, and stop handoffs; `packages/harness/deerflow/runtime/journal.py` owns +event-store writes; and `packages/harness/deerflow/runtime/runs/manager.py` owns +run cancellation, finalization, and shutdown. Inner-vs-outer cancellation +predicates remain in the MCP service. + +Cancellation drains and `McpTaskService.stop()` currently use the module-local +`_CANCELLATION_DRAIN_TIMEOUT_SECONDS = 5.0` monotonic deadline (`loop.time()`); +this constant does not define `RunManager.shutdown`. Repeated caller cancellation +is absorbed while the same absolute deadline remains in force; cancellation never +renews the wait. When caller-cancellation draining times out: +the exact asyncio operation task remains retained by the subsystem-owned registry after timeout. +The original `CancelledError` is re-raised, while +normal `McpTaskService.stop()` and `RunManager.shutdown()` record/log the +deadline and return as retained work continues. `RunManager.shutdown(timeout=5.0)` +keeps a caller-provided hard total budget (the default is overrideable), computes +one absolute deadline, and gives nested waits only the remaining time. The +retained owner consumes the eventual success, failure, or cancellation exactly +once; it must not blindly retry or requeue an operation whose durable outcome is +unknown. + +RunManager applies its caller-provided shutdown hard total budget across +cancellation-cleanup producers, manager-lock waiters, in-flight run +cancellation, heartbeat stop, orphan recovery, and trailing interrupted-status +persistence. Lock waiters are +cancelled once, tracked until a late result arrives, and release the manager +lock at most once if they acquire it after the foreground deadline. Heartbeat +and orphan tasks keep one background owner after a timed-out stop, while late +failures are consumed once. Shutdown preserves a run's real final outcome when +it settles during the drain and only marks/persists `interrupted` for work that +did not settle. + +MCP claimed batches run under a shielded supervisor. Each record has one +release state and sibling release tasks start concurrently, so started and +never-started records are released exactly once. Ordinary retry release and +cancellation release are mutually exclusive. Inner self-cancellation is +consumed and logged at the task boundary without killing the poller; an outer +cancellation keeps its original signal and releases the whole batch without +duplicating child cleanup. Claim cancellation uses the same inner-vs-outer +predicate: a self-cancelled claim is consumed only when no caller cancellation +is pending, otherwise the caller's cancellation wins. Notification cancellation +does not enter ordinary retry handling. The supervisor task, every child task, +and every ordinary or cancellation release task are retrieved and their +terminal success, failure, or cancellation is consumed exactly once. Poll, +cancellation, and notification batch results are observed and recorded at the +per-record boundary; unexpected child failures are logged with that record's +task ID, and notification failures enter that record's release handling. +Supervisor, release, and background-ownership failures receive one contextual +log. Successful completion clears task ownership without emitting a second log. + +`McpTaskService.stop()` establishes one absolute `loop.time()` deadline for the +poller cleanup. Repeated or concurrent `stop()` calls reuse that deadline and +do not recancel the same poller. If the deadline expires, the poller remains +supervised in the background; `start()` will not create an overlapping poller +while that owner is still live. + +### Cancellation-Safe Run Journal Writes + +`RunJournal` transfers each detached batch to a dedicated `put_batch` task and +shields that task from caller cancellation. Threshold flushes and the worker's +final explicit `flush()` serialize all pending and detached predecessors before +starting a later write; a normal, uncancelled flush may wait for that serialized +predecessor. A later flush that is itself cancelled only observes predecessors +through its bounded drain deadline. A successful write is discarded from the +detached registry; a late explicit failure or self-cancellation prepends its +batch exactly once. No automatic retry is launched. If the write remains +ambiguous after its drain deadline, ownership stays with its supervised +background task until the final result, so a JSONL `to_thread` append or +database commit cannot be duplicated by a blind retry. If a scheduled flush is +cancelled before its first coroutine step, its completion callback restores the +still-unowned detached batch instead. Do not reintroduce a direct +`except CancelledError: requeue` around event-store writes. diff --git a/backend/packages/harness/deerflow/runtime/cancellation.py b/backend/packages/harness/deerflow/runtime/cancellation.py new file mode 100644 index 00000000000..6e2cd99e89b --- /dev/null +++ b/backend/packages/harness/deerflow/runtime/cancellation.py @@ -0,0 +1,24 @@ +from __future__ import annotations + +import asyncio +from typing import TypeVar + +T = TypeVar("T") + + +async def wait_for_task_until( # noqa: UP047 + task: asyncio.Future[T], *, deadline: float +) -> bool: + """Wait through repeated caller cancellation without cancelling task.""" + loop = asyncio.get_running_loop() + while not task.done(): + remaining = deadline - loop.time() + if remaining <= 0: + return False + try: + done, _ = await asyncio.wait({task}, timeout=remaining) + except asyncio.CancelledError: + continue + if task in done: + return True + return True diff --git a/backend/packages/harness/deerflow/runtime/journal.py b/backend/packages/harness/deerflow/runtime/journal.py index 6f3007fa1f2..c81713863fc 100644 --- a/backend/packages/harness/deerflow/runtime/journal.py +++ b/backend/packages/harness/deerflow/runtime/journal.py @@ -21,6 +21,7 @@ import logging import time from collections.abc import Awaitable, Callable, Mapping, Sequence +from dataclasses import dataclass from datetime import UTC, datetime from typing import TYPE_CHECKING, Any, cast from uuid import UUID @@ -30,6 +31,7 @@ from langgraph.types import Command from deerflow.agents.human_input import read_human_input_response +from deerflow.runtime.cancellation import wait_for_task_until from deerflow.runtime.events.catalog import ( LLM_AI_RESPONSE_EVENT, LLM_ERROR_EVENT, @@ -51,6 +53,13 @@ _LEGACY_SUMMARY_MESSAGE_NAME = "summary" _RECONCILED_TOOL_MESSAGE_NAMES = frozenset({"ask_clarification"}) _PERSISTED_HIDDEN_HUMAN_INPUT_RESPONSE_SOURCES = frozenset({"ask_clarification"}) +_CANCELLATION_DRAIN_TIMEOUT_SECONDS = 5.0 + + +@dataclass +class _DetachedFlush: + batch: list[dict] + started: bool = False def _should_persist_human_input_message(message: BaseMessage) -> bool: @@ -243,6 +252,7 @@ def __init__( # Write buffer self._buffer: list[dict] = [] self._pending_flush_tasks: set[asyncio.Task[None]] = set() + self._detached_write_tasks: dict[asyncio.Future[Any], list[dict]] = {} self._pending_progress_task: asyncio.Task[None] | None = None self._pending_progress_delayed = False self._progress_dirty = False @@ -661,22 +671,133 @@ def _flush_sync(self) -> None: return # Skip if a flush is already in flight — avoids concurrent writes # to the same SQLite file from multiple fire-and-forget tasks. - if self._pending_flush_tasks: + if self._pending_flush_tasks or self._detached_write_tasks: return try: loop = asyncio.get_running_loop() except RuntimeError: # No event loop — keep events in buffer for later async flush. return - batch = self._buffer.copy() + detached = _DetachedFlush(self._buffer.copy()) self._buffer.clear() - task = loop.create_task(self._flush_async(batch)) + task = loop.create_task(self._flush_async(detached.batch, detached=detached)) self._pending_flush_tasks.add(task) - task.add_done_callback(self._on_flush_done) + task.add_done_callback(lambda completed: self._on_flush_done(completed, detached=detached)) + + async def _put_batch_cancellation_safe(self, batch: list[dict]) -> None: + """Drain an in-flight write before deciding whether its batch can be retried.""" + write_task = asyncio.create_task(self._store.put_batch(batch)) + write_task.set_name(f"deerflow-journal-put-batch-{self.run_id}") + try: + await asyncio.shield(write_task) + except asyncio.CancelledError: + completed = await wait_for_task_until( + write_task, + deadline=asyncio.get_running_loop().time() + _CANCELLATION_DRAIN_TIMEOUT_SECONDS, + ) + if not completed: + self._track_detached_write(write_task, batch) + logger.warning( + "Timed out after %.1f seconds draining journal write for run %s; %d events remain owned by the background write", + _CANCELLATION_DRAIN_TIMEOUT_SECONDS, + self.run_id, + len(batch), + ) + raise + try: + write_task.result() + except asyncio.CancelledError: + logger.warning( + "Journal write was cancelled while draining cancellation for run %s; returning %d events to buffer", + self.run_id, + len(batch), + ) + self._buffer = batch + self._buffer + except BaseException as error: + logger.warning( + "Journal write failed while draining cancellation for run %s — returning %d events to buffer", + self.run_id, + len(batch), + exc_info=(type(error), error, error.__traceback__), + ) + self._buffer = batch + self._buffer + raise + + def _track_detached_write(self, task: asyncio.Future[Any], batch: list[dict]) -> None: + """Keep ownership of an ambiguous write until its terminal result is known.""" + if task in self._detached_write_tasks: + return + self._detached_write_tasks[task] = batch + task.add_done_callback(self._resolve_detached_write) + + def _resolve_detached_write(self, task: asyncio.Future[Any]) -> None: + """Resolve a detached write exactly once when its task reaches a terminal state.""" + batch = self._detached_write_tasks.pop(task, None) + if batch is None: + return + try: + error = task.exception() + except asyncio.CancelledError as exc: + error = exc + if error is None: + return + logger.warning( + "Detached journal write failed for run %s; returning %d events to buffer", + self.run_id, + len(batch), + exc_info=(type(error), error, error.__traceback__), + ) + self._buffer = batch + self._buffer - async def _flush_async(self, batch: list[dict]) -> None: + async def _await_detached_writes(self) -> None: + """Observe detached predecessors without cancelling them on caller cancellation.""" + tasks = tuple(self._detached_write_tasks) + if not tasks: + return + try: + await asyncio.gather(*(asyncio.shield(task) for task in tasks), return_exceptions=True) + except asyncio.CancelledError: + deadline = asyncio.get_running_loop().time() + _CANCELLATION_DRAIN_TIMEOUT_SECONDS + for task in tasks: + if await wait_for_task_until(task, deadline=deadline): + self._resolve_detached_write(task) + raise + for task in tasks: + self._resolve_detached_write(task) + + async def _await_pending_flush_tasks(self) -> None: + """Observe threshold flushes without letting caller cancellation cancel them.""" + tasks = tuple(self._pending_flush_tasks) + if not tasks: + return + try: + await asyncio.gather(*(asyncio.shield(task) for task in tasks), return_exceptions=True) + except asyncio.CancelledError: + deadline = asyncio.get_running_loop().time() + _CANCELLATION_DRAIN_TIMEOUT_SECONDS + for task in tasks: + await wait_for_task_until(task, deadline=deadline) + raise + + async def _await_write_predecessors(self) -> None: + """Observe pending and detached writes until no predecessor can appear.""" + while True: + has_detached_writes = bool(self._detached_write_tasks) + has_pending_flushes = bool(self._pending_flush_tasks) + if not has_detached_writes and not has_pending_flushes: + return + if has_detached_writes: + await self._await_detached_writes() + if has_pending_flushes: + await self._await_pending_flush_tasks() + # Let completion callbacks move a cancelled threshold write into + # the detached registry before checking the fixed point again. + await asyncio.sleep(0) + + async def _flush_async(self, batch: list[dict], *, detached: _DetachedFlush | None = None) -> None: + if detached is not None: + detached.started = True try: - await self._store.put_batch(batch) + await self._put_batch_cancellation_safe(batch) except Exception: logger.warning( "Failed to flush %d events for run %s — returning to buffer", @@ -687,9 +808,11 @@ async def _flush_async(self, batch: list[dict]) -> None: # Return failed events to buffer for retry on next flush self._buffer = batch + self._buffer - def _on_flush_done(self, task: asyncio.Task) -> None: + def _on_flush_done(self, task: asyncio.Task, *, detached: _DetachedFlush) -> None: self._pending_flush_tasks.discard(task) if task.cancelled(): + if not detached.started: + self._buffer = detached.batch + self._buffer return exc = task.exception() if exc: @@ -886,8 +1009,7 @@ def record_delivery(self) -> None: async def flush(self) -> None: """Force flush remaining buffer. Called in worker's finally block.""" - if self._pending_flush_tasks: - await asyncio.gather(*tuple(self._pending_flush_tasks), return_exceptions=True) + await self._await_write_predecessors() while self._pending_progress_task is not None and not self._pending_progress_task.done(): if self._pending_progress_delayed: self._pending_progress_task.cancel() @@ -901,7 +1023,7 @@ async def flush(self) -> None: batch = self._buffer[: self._flush_threshold] del self._buffer[: self._flush_threshold] try: - await self._store.put_batch(batch) + await self._put_batch_cancellation_safe(batch) except Exception: self._buffer = batch + self._buffer raise diff --git a/backend/packages/harness/deerflow/runtime/runs/manager.py b/backend/packages/harness/deerflow/runtime/runs/manager.py index 9827afd118e..a2688ee2d3e 100644 --- a/backend/packages/harness/deerflow/runtime/runs/manager.py +++ b/backend/packages/harness/deerflow/runtime/runs/manager.py @@ -7,6 +7,7 @@ import socket import sqlite3 import uuid +import weakref from collections.abc import AsyncIterator, Awaitable, Callable from contextlib import asynccontextmanager from dataclasses import dataclass, field @@ -16,6 +17,7 @@ from sqlalchemy.exc import IntegrityError as SAIntegrityError +from deerflow.runtime.cancellation import wait_for_task_until from deerflow.runtime.user_context import AUTO, _AutoSentinel, resolve_user_id from deerflow.utils.time import is_lease_expired from deerflow.utils.time import now_iso as _now_iso @@ -33,6 +35,7 @@ ORPHAN_RECOVERY_STOP_REASON = "orphan_recovered" STARTUP_ORPHAN_RECOVERY_ERROR = "Gateway restarted before this run reached a durable final state." LEASE_ORPHAN_RECOVERY_ERROR = "Run lease expired — owning worker is unreachable." +_CANCELLATION_DRAIN_TIMEOUT_SECONDS = 5.0 _RETRYABLE_SQLITE_MESSAGES = ( "database is locked", @@ -250,6 +253,300 @@ def __init__( self._heartbeat_task: asyncio.Task | None = None self._heartbeat_stop: asyncio.Event | None = None self._orphan_recovery_task: asyncio.Task[None] | None = None + self._heartbeat_cancel_requested = False + self._orphan_recovery_cancel_requested = False + self._lock_waiters: set[asyncio.Task[bool]] = set() + self._lock_waiter_cancel_requested: set[asyncio.Task[bool]] = set() + self._cancellation_cleanup_tasks: set[asyncio.Future[None]] = set() + self._cancellation_cleanup_producers: set[object] = set() + self._cancellation_cleanup_state_changed = asyncio.Event() + self._shutdown_persistence_tasks: set[asyncio.Task[bool | BaseException]] = set() + self._shutdown_persistence_records: dict[asyncio.Task[bool | BaseException], RunRecord] = {} + self._shutdown_persistence_observed: weakref.WeakSet[asyncio.Task[bool | BaseException]] = weakref.WeakSet() + + def _begin_cancellation_cleanup_producer(self) -> Callable[[], None]: + """Register a producer that may later track cancellation cleanup. + + Callers must acquire this token before scheduling any callback that can + call ``_track_cancellation_cleanup`` and release it only after that + callback has either tracked its cleanup or established that none is + needed. Direct, unsupervised calls to ``_track_cancellation_cleanup`` + are supported only when no registration gap can exist. + """ + token = object() + self._cancellation_cleanup_producers.add(token) + self._cancellation_cleanup_state_changed.set() + released = False + + def release() -> None: + nonlocal released + if released: + return + released = True + self._cancellation_cleanup_producers.discard(token) + self._cancellation_cleanup_state_changed.set() + + return release + + def _track_cancellation_cleanup( + self, + task: asyncio.Future[None], + *, + action: str, + run_id: str, + ) -> None: + """Keep an ambiguous cancellation cleanup alive until it settles.""" + if task in self._cancellation_cleanup_tasks: + return + self._cancellation_cleanup_tasks.add(task) + self._cancellation_cleanup_state_changed.set() + + def finalize(completed: asyncio.Future[None]) -> None: + if completed in self._cancellation_cleanup_tasks: + self._cancellation_cleanup_tasks.discard(completed) + self._cancellation_cleanup_state_changed.set() + try: + error = completed.exception() + except asyncio.CancelledError as exc: + error = exc + if error is None: + return + logger.error( + "Run cancellation cleanup failed (%s, run_id=%s): %s", + action, + run_id, + error, + exc_info=(type(error), error, error.__traceback__), + ) + + task.add_done_callback(finalize) + + async def _wait_for_cancellation_cleanup_state_change(self, *, deadline: float) -> bool: + """Wait for producer or cleanup ownership to change until *deadline*.""" + loop = asyncio.get_running_loop() + remaining = deadline - loop.time() + if remaining <= 0: + return False + self._cancellation_cleanup_state_changed.clear() + try: + async with asyncio.timeout(remaining): + await self._cancellation_cleanup_state_changed.wait() + except TimeoutError: + return False + self._cancellation_cleanup_state_changed.clear() + return True + + async def _drain_cancellation_cleanup( + self, + cleanup: Awaitable[None], + *, + action: str, + run_id: str, + ) -> None: + """Drain one cleanup until its fixed cancellation deadline.""" + task = asyncio.ensure_future(cleanup) + deadline = asyncio.get_running_loop().time() + _CANCELLATION_DRAIN_TIMEOUT_SECONDS + if not await wait_for_task_until(task, deadline=deadline): + self._track_cancellation_cleanup(task, action=action, run_id=run_id) + logger.warning( + "Timed out after %.1f seconds waiting for run cancellation cleanup; continuing in the background (%s, run_id=%s)", + _CANCELLATION_DRAIN_TIMEOUT_SECONDS, + action, + run_id, + ) + return + try: + task.result() + except asyncio.CancelledError as exc: + logger.error( + "Run cancellation cleanup was cancelled (%s, run_id=%s): %s", + action, + run_id, + exc, + exc_info=(type(exc), exc, exc.__traceback__), + ) + except Exception as exc: # noqa: BLE001 - preserve the caller's cancellation + logger.error( + "Run cancellation cleanup failed (%s, run_id=%s): %s", + action, + run_id, + exc, + exc_info=(type(exc), exc, exc.__traceback__), + ) + + def _observe_lock_waiter_result(self, waiter: asyncio.Task[bool], *, late: bool) -> bool: + self._lock_waiter_cancel_requested.discard(waiter) + try: + return bool(waiter.result()) + except asyncio.CancelledError as exc: + if late: + logger.warning( + "Manager lock waiter was cancelled while settling in the background (task=%s): %s", + waiter.get_name(), + exc, + exc_info=(type(exc), exc, exc.__traceback__), + ) + return False + except BaseException as exc: # noqa: BLE001 - consume late waiter failures + logger.warning( + "Manager lock waiter failed while settling%s (task=%s): %s", + " in the background" if late else "", + waiter.get_name(), + exc, + exc_info=(type(exc), exc, exc.__traceback__), + ) + return False + + def _lock_waiter_done(self, waiter: asyncio.Task[bool]) -> None: + if waiter not in self._lock_waiters: + return + self._lock_waiters.discard(waiter) + acquired = self._observe_lock_waiter_result(waiter, late=True) + if not acquired: + return + try: + self._lock.release() + except RuntimeError as exc: + logger.warning( + "Late manager lock waiter acquired an already-unlocked manager lock (task=%s): %s", + waiter.get_name(), + exc, + exc_info=(type(exc), exc, exc.__traceback__), + ) + + def _track_lock_waiter(self, waiter: asyncio.Task[bool]) -> None: + if waiter in self._lock_waiters: + return + self._lock_waiters.add(waiter) + waiter.add_done_callback(self._lock_waiter_done) + logger.warning( + "Manager lock waiter did not settle before the deadline; continuing in the background (task=%s)", + waiter.get_name(), + ) + + def _cancel_lock_waiter_once(self, waiter: asyncio.Task[bool]) -> None: + if waiter.done() or waiter in self._lock_waiter_cancel_requested: + return + self._lock_waiter_cancel_requested.add(waiter) + waiter.cancel() + + async def _cancel_and_settle_lock_waiter(self, waiter: asyncio.Task[bool], *, deadline: float) -> bool: + """Cancel one lock waiter and settle it within the caller's deadline.""" + self._cancel_lock_waiter_once(waiter) + if not await wait_for_task_until(waiter, deadline=deadline): + self._track_lock_waiter(waiter) + return False + return self._observe_lock_waiter_result(waiter, late=False) + + async def _acquire_lock_until(self, deadline: float) -> bool: + """Acquire the manager lock without exceeding an absolute deadline.""" + remaining = deadline - asyncio.get_running_loop().time() + if remaining <= 0: + return False + + waiter = asyncio.create_task(self._lock.acquire()) + try: + done, _ = await asyncio.wait((waiter,), timeout=remaining) + except asyncio.CancelledError: + acquired = await self._cancel_and_settle_lock_waiter(waiter, deadline=deadline) + if acquired: + self._lock.release() + raise + + if waiter in done: + self._lock_waiter_cancel_requested.discard(waiter) + try: + acquired = waiter.result() + except asyncio.CancelledError: + return False + if acquired and asyncio.get_running_loop().time() >= deadline: + self._lock.release() + return False + return acquired + + acquired = await self._cancel_and_settle_lock_waiter(waiter, deadline=deadline) + if acquired and asyncio.get_running_loop().time() >= deadline: + self._lock.release() + return False + return acquired + + def _observe_shutdown_persistence( + self, + task: asyncio.Task[bool | BaseException], + record: RunRecord, + ) -> None: + """Consume one shutdown persistence result, at most once.""" + if task in self._shutdown_persistence_observed: + return + self._shutdown_persistence_observed.add(task) + try: + result = task.result() + except asyncio.CancelledError: + logger.warning( + "Shutdown status persistence was cancelled for run %s", + record.run_id, + ) + except BaseException as exc: # noqa: BLE001 - callback must consume all failures + logger.warning( + "Shutdown status persistence failed for run %s: %s", + record.run_id, + exc, + ) + else: + if isinstance(result, asyncio.CancelledError): + logger.warning( + "Shutdown status persistence was cancelled for run %s", + record.run_id, + ) + elif isinstance(result, BaseException): + logger.warning( + "Shutdown status persistence failed for run %s: %s", + record.run_id, + result, + ) + elif result is False: + logger.warning( + "Could not persist interrupted status for run %s during shutdown", + record.run_id, + ) + + def _shutdown_persistence_done(self, task: asyncio.Task[bool | BaseException]) -> None: + """Release task ownership and consume a late persistence result.""" + record = self._shutdown_persistence_records.pop(task, None) + self._shutdown_persistence_tasks.discard(task) + if record is not None: + self._observe_shutdown_persistence(task, record) + + async def _run_shutdown_persistence(self, record: RunRecord) -> bool | BaseException: + """Return persistence failures for the supervisor to observe.""" + try: + return await self._persist_status(record, RunStatus.interrupted) + except BaseException as exc: # noqa: BLE001 - supervisor owns the result + return exc + + def _start_shutdown_persistence(self, record: RunRecord) -> asyncio.Task[bool | BaseException]: + """Start and immediately supervise one shutdown status persistence.""" + task = asyncio.create_task(self._run_shutdown_persistence(record)) + task.set_name(f"deerflow-shutdown-persist-status-{record.run_id}") + self._shutdown_persistence_tasks.add(task) + self._shutdown_persistence_records[task] = record + task.add_done_callback(self._shutdown_persistence_done) + return task + + @staticmethod + def _observe_completed_run_task_exceptions( + run_tasks: list[asyncio.Task | None], + ) -> None: + """Consume terminal exceptions without mutating run state.""" + for task in run_tasks: + if task is None or not task.done() or task.cancelled(): + continue + try: + task.exception() + except BaseException: + # Retrieving an exception is best effort and must never replace + # shutdown's own cancellation or deadline outcome. + pass def _index_run_locked(self, record: RunRecord) -> None: """Register *record* in the thread index. Caller must hold ``self._lock``.""" @@ -974,26 +1271,31 @@ async def set_status_if_not_cancelled( return None try: - result = await self._call_store_with_retry( - "finalize_if_not_cancelled", - run_id, - lambda: self._store.finalize_if_not_cancelled( + try: + result = await self._call_store_with_retry( + "finalize_if_not_cancelled", run_id, - status=status.value, - error=error, - stop_reason=stop_reason, - ), - ) - except Exception: - async with self._lock: - record = self._runs.get(run_id) - if record is not None: - await self._mark_ownership_lost( - record, - reason=("The durable store could not confirm whether cancellation or completion won."), - require_active=False, + lambda: self._store.finalize_if_not_cancelled( + run_id, + status=status.value, + error=error, + stop_reason=stop_reason, + ), ) - return None + except Exception: + await self._fence_unconfirmed_finalization(run_id) + return None + except asyncio.CancelledError: + release_producer = self._begin_cancellation_cleanup_producer() + try: + await self._drain_cancellation_cleanup( + self._fence_unconfirmed_finalization(run_id), + action="fence cancelled finalization", + run_id=run_id, + ) + finally: + release_producer() + raise if result.cancel_action is not None: async with self._lock: @@ -1012,6 +1314,16 @@ async def set_status_if_not_cancelled( ) return None + async def _fence_unconfirmed_finalization(self, run_id: str) -> None: + async with self._lock: + record = self._runs.get(run_id) + if record is not None: + await self._mark_ownership_lost( + record, + reason="The durable store could not confirm whether cancellation or completion won.", + require_active=False, + ) + async def _ensure_delivery_receipt(self, record: RunRecord) -> bool: """Idempotently persist a zero-delivery receipt during recovery.""" if self._event_store is None: @@ -1681,21 +1993,15 @@ def reuse_idempotent_run(conflict: RunIdempotencyConflict) -> RunRecord: 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 + release_producer = self._begin_cancellation_cleanup_producer() 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) + await self._drain_cancellation_cleanup( + self._close_cancelled_admission(record), + action="close cancelled admission", + run_id=record.run_id, + ) + finally: + release_producer() raise logger.info("Run created: run_id=%s thread_id=%s", run_id, thread_id) @@ -1950,31 +2256,71 @@ async def start_heartbeat(self) -> None: """ if not self.heartbeat_enabled: return - if self._heartbeat_task is not None and not self._heartbeat_task.done(): - return + if self._heartbeat_task is not None: + if not self._heartbeat_task.done(): + return + self._clear_heartbeat_owner(self._heartbeat_task) self._heartbeat_stop = asyncio.Event() + self._heartbeat_cancel_requested = False task = asyncio.create_task(self._heartbeat_loop()) task.set_name("deerflow-run-lease-heartbeat") self._heartbeat_task = task + task.add_done_callback(self._heartbeat_done) logger.info("Run lease heartbeat started for worker %s", self._worker_id) + def _clear_heartbeat_owner(self, task: asyncio.Task[None]) -> None: + if self._heartbeat_task is task: + self._heartbeat_task = None + self._heartbeat_stop = None + self._heartbeat_cancel_requested = False + + def _heartbeat_done(self, task: asyncio.Task[None]) -> None: + """Release the heartbeat owner and consume its terminal result.""" + self._clear_heartbeat_owner(task) + if task.cancelled(): + return + try: + task.result() + except Exception: + logger.warning("Run lease heartbeat failed; its task has stopped", exc_info=True) + async def stop_heartbeat(self, *, timeout: float = 5.0) -> None: - """Stop the background heartbeat task within ``timeout`` seconds.""" + """Stop the heartbeat within ``timeout`` and supervise late completion.""" if self._heartbeat_stop is not None: self._heartbeat_stop.set() - if self._heartbeat_task is not None and not self._heartbeat_task.done(): - _, pending = await asyncio.wait( - (self._heartbeat_task,), - timeout=max(0.0, timeout), + task = self._heartbeat_task + if task is None: + self._heartbeat_stop = None + logger.info("Run lease heartbeat stopped for worker %s", self._worker_id) + return + if task.done(): + self._clear_heartbeat_owner(task) + logger.info("Run lease heartbeat stopped for worker %s", self._worker_id) + return + if self._heartbeat_cancel_requested or task.cancelling(): + logger.warning( + "Run lease heartbeat cancellation is already pending; it continues in the background (task=%s)", + task.get_name(), ) - if pending: - self._heartbeat_task.cancel() - try: - await self._heartbeat_task - except asyncio.CancelledError: - pass - self._heartbeat_task = None - self._heartbeat_stop = None + return + + _, pending = await asyncio.wait( + (task,), + timeout=max(0.0, timeout), + ) + if pending: + task.cancel() + self._heartbeat_cancel_requested = True + logger.warning( + "Run lease heartbeat did not stop within %.1fs; cancellation continues in the background (task=%s)", + timeout, + task.get_name(), + ) + await asyncio.sleep(0) + if task.done(): + await asyncio.sleep(0) + else: + self._clear_heartbeat_owner(task) logger.info("Run lease heartbeat stopped for worker %s", self._worker_id) async def _heartbeat_loop(self) -> None: @@ -2009,6 +2355,9 @@ async def _heartbeat_loop(self) -> None: except Exception: logger.warning("Heartbeat renewal cycle failed", exc_info=True) + if stop.is_set(): + break + # Reconcile every 3rd cycle (= every lease_seconds). Startup # reconciliation (in langgraph_runtime) covers the initial # sweep; this periodic pass catches orphans whose lease @@ -2169,12 +2518,14 @@ def _schedule_orphan_reconciliation(self) -> None: task = asyncio.create_task(self._reconcile_orphans_periodic()) task.set_name("deerflow-periodic-orphan-recovery") self._orphan_recovery_task = task + self._orphan_recovery_cancel_requested = False task.add_done_callback(self._orphan_reconciliation_done) def _orphan_reconciliation_done(self, task: asyncio.Task[None]) -> None: """Clear and inspect the supervised single-flight recovery task.""" if self._orphan_recovery_task is task: self._orphan_recovery_task = None + self._orphan_recovery_cancel_requested = False if task.cancelled(): return try: @@ -2183,18 +2534,29 @@ def _orphan_reconciliation_done(self, task: asyncio.Task[None]) -> None: logger.warning("Periodic orphan reconciliation failed", exc_info=True) async def _drain_orphan_recovery_task(self, *, timeout: float) -> None: - """Boundedly await the supervised recovery pass during shutdown.""" + """Observe orphan recovery without extending the caller's deadline.""" task = self._orphan_recovery_task if task is None or task.done(): return - _, pending = await asyncio.wait((task,), timeout=max(0.0, timeout)) - if pending: - task.cancel() - await asyncio.gather(task, return_exceptions=True) + if self._orphan_recovery_cancel_requested or task.cancelling(): logger.warning( - "Orphan recovery drain exceeded %.1fs on shutdown; cancelled the active pass", - timeout, + "Orphan recovery cancellation is already pending; it continues in the background (task=%s)", + task.get_name(), ) + return + _, pending = await asyncio.wait((task,), timeout=max(0.0, timeout)) + if not pending: + return + task.cancel() + self._orphan_recovery_cancel_requested = True + logger.warning( + "Orphan recovery drain exceeded %.1fs; cancellation continues in the background (task=%s)", + timeout, + task.get_name(), + ) + await asyncio.sleep(0) + if task.done(): + await asyncio.sleep(0) async def shutdown(self, *, timeout: float = 5.0) -> None: """Cancel and bounded-await all in-flight runs on process shutdown. @@ -2228,72 +2590,112 @@ async def shutdown(self, *, timeout: float = 5.0) -> None: after ``timeout`` are logged and may still race teardown. """ loop = asyncio.get_running_loop() - deadline = loop.time() + timeout + deadline = loop.time() + max(0.0, timeout) - async with self._lock: - inflight = [record for record in self._runs.values() if record.status in (RunStatus.pending, RunStatus.running) and record.task is not None and not record.task.done()] - for record in inflight: - record.abort_action = "interrupt" - record.abort_event.set() - record.task.cancel() # type: ignore[union-attr] # filtered above - # Status is decided AFTER the drain (below), not here: a run that - # completes on its own during the drain must keep its real status. + inflight: list[RunRecord] = [] + initial_lock_acquired = await self._acquire_lock_until(deadline) + if initial_lock_acquired: + try: + inflight = [record for record in self._runs.values() if record.status in (RunStatus.pending, RunStatus.running) and record.task is not None and not record.task.done()] + for record in inflight: + record.abort_action = "interrupt" + record.abort_event.set() + record.task.cancel() # type: ignore[union-attr] # filtered above + # Status is decided AFTER the drain (below), not here: a run that + # completes on its own during the drain must keep its real status. + finally: + self._lock.release() + else: + logger.warning( + "Run shutdown could not acquire manager lock before deadline; unable to snapshot or cancel in-flight runs", + ) await self.stop_heartbeat(timeout=max(0.0, deadline - loop.time())) - if not inflight: - await self._drain_orphan_recovery_task(timeout=max(0.0, deadline - loop.time())) - return - - tasks = [record.task for record in inflight] - _, pending = await asyncio.wait(tasks, timeout=max(0.0, deadline - loop.time())) + # Run cancellation handlers may register cleanup after their run task + # receives the shutdown cancellation. Re-snapshot ownership after each + # wait so those tasks share the same top-level deadline. + run_tasks = [record.task for record in inflight] + pending_run_tasks: set[asyncio.Future[None]] = set() + while True: + active_run_tasks = {task for task in run_tasks if task is not None and not task.done()} + active_cleanup_tasks = {task for task in self._cancellation_cleanup_tasks if not task.done()} + owned_tasks = active_run_tasks | active_cleanup_tasks + if owned_tasks: + remaining = deadline - loop.time() + if remaining <= 0: + pending_run_tasks = active_run_tasks + break + _, pending = await asyncio.wait(owned_tasks, timeout=remaining) + pending_run_tasks = {task for task in pending if task in active_run_tasks} + if loop.time() >= deadline: + break + continue + if not self._cancellation_cleanup_producers: + break + if not await self._wait_for_cancellation_cleanup_state_change(deadline=deadline): + break # Only mark/persist ``interrupted`` for runs that did not settle on their # own (still pending after the timeout, or ended cancelled). A run that # finished normally during the drain keeps the status it set for itself. to_persist: list[RunRecord] = [] - async with self._lock: - for record in inflight: - task = record.task - if task not in pending and not task.cancelled(): - # Completed on its own — retrieve any surfaced exception so it - # is not reported as "never retrieved", and keep its status. - task.exception() # type: ignore[union-attr] # done & not cancelled - continue - if record.status in (RunStatus.pending, RunStatus.running): - record.status = RunStatus.interrupted - record.updated_at = _now_iso() - to_persist.append(record) + post_wait_lock_acquired = False + try: + post_wait_lock_acquired = await self._acquire_lock_until(deadline) + finally: + self._observe_completed_run_task_exceptions(run_tasks) + if post_wait_lock_acquired: + try: + for record in inflight: + task = record.task + if task not in pending_run_tasks and not task.cancelled(): + continue + if record.status in (RunStatus.pending, RunStatus.running): + record.status = RunStatus.interrupted + record.updated_at = _now_iso() + to_persist.append(record) + finally: + self._lock.release() + elif initial_lock_acquired: + logger.warning( + "Run shutdown could not acquire manager lock before deadline after draining runs; skipping status mutation and persistence", + ) # Bound the trailing status persistence within the remaining budget so a # slow store (``_call_store_with_retry`` can back off under DB pressure) # cannot push shutdown past ``timeout``. if to_persist: + persistence_records = [(self._start_shutdown_persistence(record), record) for record in to_persist] remaining = deadline - loop.time() if remaining <= 0: - logger.warning("Run drain budget exhausted before persisting %d interrupted run(s) on shutdown", len(to_persist)) + logger.warning("Run drain budget exhausted before persisting %d interrupted run(s) on shutdown; tasks remain supervised in the background", len(to_persist)) + done: set[asyncio.Task[bool | BaseException]] = set() else: - try: - results = await asyncio.wait_for( - asyncio.gather(*(self._persist_status(record, RunStatus.interrupted) for record in to_persist), return_exceptions=True), - timeout=remaining, - ) - except TimeoutError: - logger.warning("Run drain status persistence exceeded the %.1fs budget; %d record(s) may not be persisted", timeout, len(to_persist)) - else: - # ``_persist_status`` is best-effort: it catches and logs its - # own failures, returning ``False``. Inspect the aggregate so a - # partial failure is surfaced at shutdown level (with the - # run_id) instead of being silently swallowed by the gather. - for record, result in zip(to_persist, results): - if isinstance(result, Exception): - logger.warning("Unexpected error persisting interrupted status for run %s during shutdown: %r", record.run_id, result) - elif result is False: - logger.warning("Could not persist interrupted status for run %s during shutdown", record.run_id) + done, _ = await asyncio.wait( + [task for task, _record in persistence_records], + timeout=remaining, + ) + for task, record in persistence_records: + if task in done: + self._observe_shutdown_persistence(task, record) + pending_persistence_count = sum(1 for task, _record in persistence_records if not task.done()) + if pending_persistence_count: + logger.warning( + "Run drain status persistence exceeded the %.1fs budget; %d record(s) remain supervised in the background", + timeout, + pending_persistence_count, + ) - if pending: - logger.warning("Run drain exceeded %.1fs on shutdown; %d run task(s) still active and may race checkpointer teardown", timeout, len(pending)) - logger.info("Drained %d in-flight run(s) on shutdown (%d settled within %.1fs)", len(inflight), len(inflight) - len(pending), timeout) + pending_cleanup_count = sum(1 for task in self._cancellation_cleanup_tasks if not task.done()) + if pending_cleanup_count: + logger.warning( + "Run shutdown deadline expired with %d cancellation cleanup task(s) still supervised", + pending_cleanup_count, + ) + if pending_run_tasks: + logger.warning("Run drain exceeded %.1fs on shutdown; %d run task(s) still active and may race checkpointer teardown", timeout, len(pending_run_tasks)) + logger.info("Drained %d in-flight run(s) on shutdown (%d settled within %.1fs)", len(inflight), len(inflight) - len(pending_run_tasks), timeout) await self._drain_orphan_recovery_task(timeout=max(0.0, deadline - loop.time())) diff --git a/backend/tests/test_mcp_task_repository.py b/backend/tests/test_mcp_task_repository.py index 79e5af32f6b..7bcf3f10f35 100644 --- a/backend/tests/test_mcp_task_repository.py +++ b/backend/tests/test_mcp_task_repository.py @@ -288,6 +288,58 @@ async def test_release_claim_retries_transient_poll_failure(tmp_path): assert stored["lease_owner"] is None +@pytest.mark.asyncio +async def test_release_poll_claim_after_cancellation_preserves_poll_failure_state(tmp_path): + repo = await _make_repo(tmp_path) + now = datetime.now(UTC) + await _create_working_task(repo, task_id="task-cancelled-poll", now=now) + await repo.claim_due_tasks(now=now, lease_owner="worker-1", lease_seconds=60, limit=10) + retry_at = now + timedelta(seconds=30) + await repo.release_claim( + "task-cancelled-poll", + lease_owner="worker-1", + next_poll_at=retry_at, + error="temporary network failure", + ) + before = await repo.get("task-cancelled-poll", user_id="user-1") + assert before is not None + + await repo.claim_due_tasks(now=retry_at, lease_owner="worker-2", lease_seconds=60, limit=10) + released = await repo.release_poll_claim_after_cancellation( + "task-cancelled-poll", + lease_owner="worker-2", + ) + + assert released is True + stored = await repo.get("task-cancelled-poll", user_id="user-1") + assert stored is not None + assert stored["next_poll_at"] == before["next_poll_at"] + assert stored["last_poll_error"] == before["last_poll_error"] + assert stored["consecutive_poll_error_count"] == before["consecutive_poll_error_count"] + assert stored["poll_attempt_count"] == before["poll_attempt_count"] + 1 + assert stored["lease_owner"] is None + assert stored["lease_expires_at"] is None + + +@pytest.mark.asyncio +async def test_release_poll_claim_after_cancellation_requires_current_owner(tmp_path): + repo = await _make_repo(tmp_path) + now = datetime.now(UTC) + await _create_working_task(repo, task_id="task-stale-cancel", now=now) + await repo.claim_due_tasks(now=now, lease_owner="worker-current", lease_seconds=60, limit=10) + + released = await repo.release_poll_claim_after_cancellation( + "task-stale-cancel", + lease_owner="worker-stale", + ) + + assert released is False + stored = await repo.get("task-stale-cancel", user_id="user-1") + assert stored is not None + assert stored["lease_owner"] == "worker-current" + assert stored["lease_expires_at"] is not None + + @pytest.mark.asyncio async def test_consecutive_poll_error_count_increments_and_resets_on_success(tmp_path): repo = await _make_repo(tmp_path) diff --git a/backend/tests/test_mcp_task_service.py b/backend/tests/test_mcp_task_service.py index 54a4c94f949..f27b31f0b5f 100644 --- a/backend/tests/test_mcp_task_service.py +++ b/backend/tests/test_mcp_task_service.py @@ -48,6 +48,10 @@ async def release_claim(self, task_id, **kwargs): self.released.append((task_id, kwargs)) return True + async def release_poll_claim_after_cancellation(self, task_id, **kwargs): + self.released.append((task_id, kwargs)) + return True + class FailingApplyRepository(FakeRepository): async def apply_snapshot(self, task_id, **kwargs): @@ -347,7 +351,7 @@ async def test_submit_repeated_cancellation_does_not_interrupt_compensation(): @pytest.mark.asyncio async def test_submit_stops_waiting_for_hung_compensation_without_cancelling_it(monkeypatch, caplog): - monkeypatch.setattr(service_module, "_UNTRACKED_TASK_COMPENSATION_WAIT_SECONDS", 0) + monkeypatch.setattr(service_module, "_CANCELLATION_DRAIN_TIMEOUT_SECONDS", 0) repo = BlockingCreateRepository() driver = BlockingCancelDriver( submission=TaskSubmission( @@ -1417,3 +1421,2176 @@ async def test_stop_cancels_a_hung_driver_poll(): await asyncio.wait_for(service.stop(), timeout=1) assert driver.cancelled is True + + +@pytest.mark.asyncio +async def test_stop_callers_share_deadline_and_log_one_timeout(monkeypatch, caplog): + monkeypatch.setattr(service_module, "_CANCELLATION_DRAIN_TIMEOUT_SECONDS", 0.05) + clock = [0.0] + wait_timeouts = [] + wait_started = [asyncio.Event(), asyncio.Event()] + release_wait = asyncio.Event() + + async def fake_wait(_tasks, *, timeout): + wait_timeouts.append(timeout) + wait_started[len(wait_timeouts) - 1].set() + if len(wait_timeouts) == 1: + # The second caller arrives 40ms into the first caller's budget. + clock[0] = 0.04 + await release_wait.wait() + return set(), set() + + monkeypatch.setattr(service_module.asyncio, "wait", fake_wait) + monkeypatch.setattr( + service_module.asyncio, + "get_running_loop", + lambda: SimpleNamespace(time=lambda: clock[0]), + ) + + poller_started = asyncio.Event() + cleanup_started = asyncio.Event() + finish = asyncio.Event() + cancel_count = 0 + + async def stubborn_poller(): + nonlocal cancel_count + poller_started.set() + try: + await asyncio.Future() + except asyncio.CancelledError: + cancel_count += 1 + cleanup_started.set() + await finish.wait() + + service = McpTaskService( + repository=FakeRepository(), + drivers=McpTaskDriverRegistry(), + poll_interval_seconds=60, + lease_seconds=120, + max_concurrent_polls=3, + ) + monkeypatch.setattr(service, "_run_loop", stubborn_poller) + + await service.start() + await poller_started.wait() + poller = service._task + assert poller is not None + first_stop = second_stop = None + + try: + with caplog.at_level(logging.WARNING): + first_stop = asyncio.create_task(service.stop()) + await wait_started[0].wait() + await cleanup_started.wait() + + second_stop = asyncio.create_task(service.stop()) + await wait_started[1].wait() + + assert wait_timeouts == pytest.approx([0.05, 0.01]) + release_wait.set() + await asyncio.gather(first_stop, second_stop) + + assert cancel_count == 1 + assert sum("Timed out after" in record.getMessage() for record in caplog.records) == 1 + finally: + release_wait.set() + if first_stop is not None and not first_stop.done(): + await first_stop + if second_stop is not None and not second_stop.done(): + await second_stop + finish.set() + await poller + + +@pytest.mark.asyncio +async def test_poller_done_clears_stop_state_and_ignores_stale_callback(monkeypatch, caplog): + monkeypatch.setattr(service_module, "_CANCELLATION_DRAIN_TIMEOUT_SECONDS", 0.05) + clock = [0.0] + wait_timeouts = [] + wait_started = [asyncio.Event(), asyncio.Event()] + release_wait = [asyncio.Event(), asyncio.Event()] + + async def fake_wait(_tasks, *, timeout): + index = len(wait_timeouts) + wait_timeouts.append(timeout) + wait_started[index].set() + await release_wait[index].wait() + return set(), set() + + monkeypatch.setattr(service_module.asyncio, "wait", fake_wait) + monkeypatch.setattr( + service_module.asyncio, + "get_running_loop", + lambda: SimpleNamespace(time=lambda: clock[0]), + ) + + first_started = asyncio.Event() + first_finish = asyncio.Event() + second_started = asyncio.Event() + second_finish = asyncio.Event() + + async def first_poller(): + first_started.set() + try: + await asyncio.Future() + except asyncio.CancelledError: + await first_finish.wait() + + async def second_poller(): + second_started.set() + try: + await asyncio.Future() + except asyncio.CancelledError: + await second_finish.wait() + + pollers = iter((first_poller, second_poller)) + + async def run_loop(): + await next(pollers)() + + service = McpTaskService( + repository=FakeRepository(), + drivers=McpTaskDriverRegistry(), + poll_interval_seconds=60, + lease_seconds=120, + max_concurrent_polls=3, + ) + monkeypatch.setattr(service, "_run_loop", run_loop) + + await service.start() + await first_started.wait() + first_task = service._task + assert first_task is not None + + with caplog.at_level(logging.WARNING): + first_stop = asyncio.create_task(service.stop()) + await wait_started[0].wait() + release_wait[0].set() + await first_stop + + assert service._stop_deadline == pytest.approx(0.05) + assert service._stop_timeout_logged is True + + first_finish.set() + await first_task + await asyncio.sleep(0) + assert service._task is None + assert service._stopping_task is None + assert service._stop_deadline is None + assert service._stop_timeout_logged is False + + clock[0] = 10.0 + await service.start() + await second_started.wait() + second_task = service._task + assert second_task is not None + + second_stop = asyncio.create_task(service.stop()) + await wait_started[1].wait() + assert wait_timeouts == pytest.approx([0.05, 0.05]) + + # A callback from the completed poller must not clear the new episode. + service._poller_done(first_task) + assert service._task is second_task + assert service._stopping_task is second_task + assert service._stop_deadline == pytest.approx(10.05) + assert service._stop_timeout_logged is False + + release_wait[1].set() + await second_stop + assert service._stop_timeout_logged is True + + second_finish.set() + await second_task + await asyncio.sleep(0) + + assert service._task is None + assert service._stopping_task is None + assert service._stop_deadline is None + assert service._stop_timeout_logged is False + assert sum("Timed out after" in record.getMessage() for record in caplog.records) == 2 + + +@pytest.mark.asyncio +async def test_stop_returns_with_timed_out_poller_and_start_does_not_overlap(monkeypatch): + monkeypatch.setattr(service_module, "_CANCELLATION_DRAIN_TIMEOUT_SECONDS", 0.01) + poller_started = asyncio.Event() + cleanup_started = asyncio.Event() + finish = asyncio.Event() + + async def stubborn_poller(): + poller_started.set() + try: + await asyncio.Future() + except asyncio.CancelledError: + cleanup_started.set() + try: + await finish.wait() + except asyncio.CancelledError: + # Make a second poller cancellation observable while keeping + # the test cleanup deterministic. + return + + service = McpTaskService( + repository=FakeRepository(), + drivers=McpTaskDriverRegistry(), + poll_interval_seconds=60, + lease_seconds=120, + max_concurrent_polls=3, + ) + monkeypatch.setattr(service, "_run_loop", stubborn_poller) + + await service.start() + await poller_started.wait() + poller = service._task + assert poller is not None + + try: + await asyncio.wait_for(service.stop(), timeout=0.2) + await cleanup_started.wait() + + assert service._task is poller + assert not poller.done() + + await service.start() + assert service._task is poller + + finish.set() + await asyncio.wait_for(poller, timeout=0.2) + await asyncio.sleep(0) + assert service._task is None + finally: + finish.set() + if not poller.done(): + await asyncio.wait_for(poller, timeout=0.2) + + +@pytest.mark.asyncio +async def test_stop_caller_cancellation_is_bounded_and_does_not_recancel_poller(monkeypatch, caplog): + monkeypatch.setattr(service_module, "_CANCELLATION_DRAIN_TIMEOUT_SECONDS", 0.05) + poller_started = asyncio.Event() + cleanup_started = asyncio.Event() + finish = asyncio.Event() + cancellation_count = 0 + + async def stubborn_poller(): + nonlocal cancellation_count + poller_started.set() + try: + await asyncio.Future() + except asyncio.CancelledError: + cancellation_count += 1 + cleanup_started.set() + while not finish.is_set(): + try: + await finish.wait() + except asyncio.CancelledError: + cancellation_count += 1 + + service = McpTaskService( + repository=FakeRepository(), + drivers=McpTaskDriverRegistry(), + poll_interval_seconds=60, + lease_seconds=120, + max_concurrent_polls=3, + ) + monkeypatch.setattr(service, "_run_loop", stubborn_poller) + + await service.start() + await poller_started.wait() + poller = service._task + assert poller is not None + caller = asyncio.create_task(service.stop()) + await cleanup_started.wait() + + try: + caller.cancel() + await asyncio.sleep(0) + caller.cancel() + + with caplog.at_level(logging.WARNING): + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(asyncio.shield(caller), timeout=0.2) + + assert cancellation_count == 1 + assert service._task is poller + assert not poller.done() + assert "cleanup continues in the background" in caplog.text + + await service.start() + assert service._task is poller + + await service.stop() + assert cancellation_count == 1 + assert service._task is poller + + finish.set() + await asyncio.wait_for(poller, timeout=0.2) + await asyncio.sleep(0) + assert service._task is None + finally: + finish.set() + if not poller.done(): + await asyncio.wait_for(poller, timeout=0.2) + if not caller.done(): + caller.cancel() + with pytest.raises(asyncio.CancelledError): + await caller + + +@pytest.mark.asyncio +async def test_finished_poller_failure_is_logged_and_clears_task(monkeypatch, caplog): + poller_started = asyncio.Event() + fail = asyncio.Event() + + async def failing_poller(): + poller_started.set() + await fail.wait() + raise RuntimeError("poller cleanup failed") + + service = McpTaskService( + repository=FakeRepository(), + drivers=McpTaskDriverRegistry(), + poll_interval_seconds=60, + lease_seconds=120, + max_concurrent_polls=3, + ) + monkeypatch.setattr(service, "_run_loop", failing_poller) + + with caplog.at_level(logging.ERROR): + await service.start() + await poller_started.wait() + poller = service._task + assert poller is not None + fail.set() + await asyncio.wait({poller}) + await asyncio.sleep(0) + + assert service._task is None + assert "MCP task poller failed" in caplog.text + assert "poller cleanup failed" in caplog.text + + +class CancellationBlockingApplyRepo(FakeRepository): + """``apply_cancel_snapshot`` blocks so the caller can be cancelled mid-flight.""" + + def __init__(self, *, release_error: Exception | None = None, block_release: bool = False): + super().__init__() + self.apply_started = asyncio.Event() + self.release_cancel_calls = [] + self.release_error = release_error + self.block_release = block_release + self.release_started = asyncio.Event() + self.finish_release = asyncio.Event() + self.release_completed = False + self.release_interrupted = False + + async def apply_cancel_snapshot(self, task_id, **kwargs): + self.applied.append((task_id, kwargs)) + self.apply_started.set() + await asyncio.Event().wait() + + async def release_cancel_claim(self, task_id, **kwargs): + self.release_cancel_calls.append((task_id, kwargs)) + self.release_started.set() + if self.block_release: + try: + await self.finish_release.wait() + except asyncio.CancelledError: + self.release_interrupted = True + raise + if self.release_error is not None: + raise self.release_error + self.release_completed = True + return True + + +def test_cancel_one_releases_claim_when_cancelled(): + """A cancel that lands mid-``_cancel_one`` must still release the claim.""" + repo = CancellationBlockingApplyRepo() + driver = FakeDriver() # cancel() returns a CANCELLED snapshot + drivers = McpTaskDriverRegistry() + drivers.register("fake", driver) + service = McpTaskService( + repository=repo, + drivers=drivers, + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_polls=3, + ) + record = _claimed_row() + + async def main(): + task = asyncio.create_task(service._cancel_one(record)) + await repo.apply_started.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert repo.release_cancel_calls + assert repo.release_cancel_calls[0][0] == record["id"] + + asyncio.run(main()) + + +@pytest.mark.asyncio +async def test_cancel_one_preserves_cancellation_when_release_fails(caplog): + repo = CancellationBlockingApplyRepo(release_error=RuntimeError("release unavailable")) + drivers = McpTaskDriverRegistry() + drivers.register("fake", FakeDriver()) + service = McpTaskService( + repository=repo, + drivers=drivers, + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_polls=3, + ) + + task = asyncio.create_task(service._cancel_one(_claimed_row())) + await repo.apply_started.wait() + task.cancel() + + with caplog.at_level(logging.ERROR), pytest.raises(asyncio.CancelledError): + await task + + assert repo.release_cancel_calls + assert "release unavailable" in caplog.text + + +@pytest.mark.asyncio +async def test_cancel_one_repeated_cancellation_does_not_interrupt_release(): + repo = CancellationBlockingApplyRepo(block_release=True) + drivers = McpTaskDriverRegistry() + drivers.register("fake", FakeDriver()) + service = McpTaskService( + repository=repo, + drivers=drivers, + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_polls=3, + ) + + task = asyncio.create_task(service._cancel_one(_claimed_row())) + await repo.apply_started.wait() + task.cancel() + await repo.release_started.wait() + task.cancel() + repo.finish_release.set() + + with pytest.raises(asyncio.CancelledError): + await task + + assert repo.release_completed is True + assert repo.release_interrupted is False + + +@pytest.mark.asyncio +async def test_cancel_one_failure_release_retries_after_caller_cancellation(): + repo = CancellationBlockingApplyRepo(block_release=True) + drivers = McpTaskDriverRegistry() + drivers.register("fake", FakeDriver(cancel_error=RuntimeError("remote unavailable"))) + service = McpTaskService( + repository=repo, + drivers=drivers, + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_polls=3, + ) + + task = asyncio.create_task(service._cancel_one(_claimed_row())) + await repo.release_started.wait() + task.cancel() + repo.finish_release.set() + + with pytest.raises(asyncio.CancelledError): + await task + + assert repo.release_completed is True + assert repo.release_interrupted is True + assert len(repo.release_cancel_calls) == 2 + + +class CancelAfterClaimRepository(FakeRepository): + def __init__(self, *, phase: str): + super().__init__() + self.phase = phase + self.caller_task = None + self.cancel_releases = [] + self.notification_claim_releases = [] + self.notification_lease_releases = [] + + def _cancel_caller(self): + task = self.caller_task + assert task is not None + task.cancel() + + async def claim_cancel_requests(self, **_kwargs): + if self.phase != "cancel": + return [] + self._cancel_caller() + return [_claimed_row()] + + async def claim_due_tasks(self, **_kwargs): + if self.phase != "poll": + return [] + self._cancel_caller() + return [_claimed_row()] + + async def claim_notification_work(self, **_kwargs): + if not self.phase.startswith("notification_"): + return [] + status = self.phase.removeprefix("notification_") + self._cancel_caller() + return [ + { + **_claimed_row(), + "notification_status": status, + "notification_run_id": "notify-run-1" if status == "dispatched" else None, + "dispatch_version": 2, + "dispatch_attempt": 0, + "dispatch_event": {"status": "completed"}, + } + ] + + async def release_cancel_claim(self, task_id, **kwargs): + self.cancel_releases.append((task_id, kwargs)) + return True + + async def release_notification_claim(self, task_id, **kwargs): + self.notification_claim_releases.append((task_id, kwargs)) + return True + + async def release_notification_lease(self, task_id, **kwargs): + self.notification_lease_releases.append((task_id, kwargs)) + return True + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("phase", "released_attr"), + [ + ("poll", "released"), + ("cancel", "cancel_releases"), + ("notification_claimed", "notification_claim_releases"), + ("notification_dispatched", "notification_lease_releases"), + ], +) +async def test_cancellation_immediately_after_claim_releases_every_record(phase, released_attr): + repo = CancelAfterClaimRepository(phase=phase) + repo.caller_task = asyncio.current_task() + drivers = McpTaskDriverRegistry() + drivers.register("fake", FakeDriver()) + service = McpTaskService( + repository=repo, + drivers=drivers, + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_polls=3, + launch_notification=AsyncMock(return_value={"run_id": "notify-run-1"}), + get_run=AsyncMock(return_value=SimpleNamespace(assistant_id="lead_agent")), + ) + + with pytest.raises(asyncio.CancelledError): + await service.run_once(now=datetime.now(UTC)) + + assert [task_id for task_id, _kwargs in getattr(repo, released_attr)] == ["task-1"] + + +class DurableClaimHandoffRepository(CancelAfterClaimRepository): + def __init__(self, *, phase: str): + super().__init__(phase=phase) + self.claim_committed = asyncio.Event() + self.allow_claim_return = asyncio.Event() + self.claim_cancelled = False + + async def _return_after_commit(self, records): + self.claim_committed.set() + try: + await self.allow_claim_return.wait() + except asyncio.CancelledError: + self.claim_cancelled = True + raise + return records + + async def claim_cancel_requests(self, **_kwargs): + if self.phase != "cancel": + return [] + return await self._return_after_commit([_claimed_row()]) + + async def claim_due_tasks(self, **_kwargs): + if self.phase != "poll": + return [] + return await self._return_after_commit([_claimed_row()]) + + async def claim_notification_work(self, **_kwargs): + if not self.phase.startswith("notification_"): + return [] + status = self.phase.removeprefix("notification_") + return await self._return_after_commit( + [ + { + **_claimed_row(), + "notification_status": status, + "notification_run_id": "notify-run-1" if status == "dispatched" else None, + "dispatch_version": 2, + "dispatch_attempt": 0, + "dispatch_event": {"status": "completed"}, + } + ] + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("phase", "released_attr"), + [ + ("poll", "released"), + ("cancel", "cancel_releases"), + ("notification_claimed", "notification_claim_releases"), + ("notification_dispatched", "notification_lease_releases"), + ], +) +async def test_cancellation_during_durable_claim_handoff_drains_and_releases(phase, released_attr): + repo = DurableClaimHandoffRepository(phase=phase) + service = McpTaskService( + repository=repo, + drivers=McpTaskDriverRegistry(), + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_polls=3, + launch_notification=AsyncMock(), + get_run=AsyncMock(), + ) + + task = asyncio.create_task(service.run_once(now=datetime.now(UTC))) + await repo.claim_committed.wait() + task.cancel() + await asyncio.sleep(0) + task.cancel() + await asyncio.sleep(0) + repo.allow_claim_return.set() + + with pytest.raises(asyncio.CancelledError): + await task + + assert repo.claim_cancelled is False + assert [task_id for task_id, _kwargs in getattr(repo, released_attr)] == ["task-1"] + + +class NotificationFallbackCancellationRepo(FakeRepository): + def __init__(self): + super().__init__() + self.release_started = asyncio.Event() + self.finish_release = asyncio.Event() + self.release_calls = [] + self.release_interrupted = False + self.release_completed = False + + async def claim_cancel_requests(self, **_kwargs): + return [] + + async def claim_due_tasks(self, **_kwargs): + return [] + + async def claim_notification_work(self, **_kwargs): + if self.claimed: + return [] + self.claimed = True + return [ + { + **_claimed_row(), + "notification_status": "dispatched", + "notification_run_id": "notify-run-1", + "dispatch_version": 2, + } + ] + + async def release_notification_lease(self, task_id, **kwargs): + self.release_calls.append((task_id, kwargs)) + self.release_started.set() + try: + await self.finish_release.wait() + except asyncio.CancelledError: + self.release_interrupted = True + raise + self.release_completed = True + return True + + +@pytest.mark.asyncio +async def test_notification_batch_fallback_release_survives_caller_cancellation(): + repo = NotificationFallbackCancellationRepo() + service = McpTaskService( + repository=repo, + drivers=McpTaskDriverRegistry(), + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_polls=3, + launch_notification=AsyncMock(), + get_run=AsyncMock(side_effect=RuntimeError("run store unavailable")), + ) + + task = asyncio.create_task(service._run_notifications(now=datetime.now(UTC))) + await repo.release_started.wait() + task.cancel() + repo.finish_release.set() + + with pytest.raises(asyncio.CancelledError): + await task + + assert repo.release_interrupted is False + assert repo.release_completed is True + assert len(repo.release_calls) == 1 + + +class NotificationPersistenceRepo: + def __init__(self, *, release_error: BaseException | None = None): + self.claimed = False + self.mark_started = asyncio.Event() + self.release_finished = asyncio.Event() + self.release_calls = [] + self.release_error = release_error + + async def claim_due_tasks(self, **_kwargs): + return [] + + async def claim_notification_work(self, **_kwargs): + if self.claimed: + return [] + self.claimed = True + return [ + { + **_claimed_row(), + "notification_status": "claimed", + "dispatch_version": 2, + "dispatch_attempt": 0, + "dispatch_event": {"status": "completed"}, + } + ] + + async def mark_notification_dispatched(self, *_args, **_kwargs): + self.mark_started.set() + await asyncio.Event().wait() + + async def release_notification_claim(self, task_id, **kwargs): + self.release_calls.append((task_id, kwargs)) + if self.release_error is not None: + self.release_finished.set() + raise self.release_error + return True + + async def release_notification_lease(self, task_id, **kwargs): + self.release_calls.append((task_id, kwargs)) + if self.release_error is not None: + self.release_finished.set() + raise self.release_error + return True + + +@pytest.mark.asyncio +async def test_stop_releases_notification_claim_during_dispatched_persistence(): + repo = NotificationPersistenceRepo() + service = McpTaskService( + repository=repo, + drivers=McpTaskDriverRegistry(), + poll_interval_seconds=60, + lease_seconds=120, + max_concurrent_polls=3, + launch_notification=AsyncMock(return_value={"run_id": "notify-run-1"}), + get_run=AsyncMock(return_value=SimpleNamespace(assistant_id="lead_agent")), + ) + + await service.start() + await asyncio.wait_for(repo.mark_started.wait(), timeout=1) + await asyncio.wait_for(service.stop(), timeout=1) + + assert repo.release_calls + assert {task_id for task_id, _kwargs in repo.release_calls} == {"task-1"} + + +@pytest.mark.asyncio +async def test_notification_cancellation_preserves_cancelled_error_when_release_fails(caplog): + repo = NotificationPersistenceRepo(release_error=RuntimeError("notification release unavailable")) + service = McpTaskService( + repository=repo, + drivers=McpTaskDriverRegistry(), + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_polls=3, + launch_notification=AsyncMock(return_value={"run_id": "notify-run-1"}), + get_run=AsyncMock(return_value=SimpleNamespace(assistant_id="lead_agent")), + ) + record = (await repo.claim_notification_work())[0] + + task = asyncio.create_task(service._notify_one(record, now=datetime.now(UTC))) + await repo.mark_started.wait() + task.cancel() + + with caplog.at_level(logging.ERROR), pytest.raises(asyncio.CancelledError): + await task + + assert repo.release_calls + assert "notification release unavailable" in caplog.text + + +class NotificationFailureReleaseRepo(NotificationPersistenceRepo): + async def claim_notification_work(self, **_kwargs): + if self.claimed: + return [] + self.claimed = True + return [ + { + **_claimed_row(), + "notification_status": "dispatched", + "notification_run_id": "notify-run-1", + "dispatch_version": 2, + } + ] + + +@pytest.mark.asyncio +async def test_notification_failure_release_self_cancellation_does_not_kill_poller(caplog): + repo = NotificationFailureReleaseRepo(release_error=asyncio.CancelledError("notification release cancelled itself")) + service = McpTaskService( + repository=repo, + drivers=McpTaskDriverRegistry(), + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_polls=3, + launch_notification=AsyncMock(return_value={"run_id": "notify-run-1"}), + get_run=AsyncMock(side_effect=RuntimeError("run store unavailable")), + ) + + try: + with caplog.at_level(logging.ERROR): + await service.start() + await repo.release_finished.wait() + async with asyncio.timeout(1): + while not any("MCP task batch release failed" in record.message for record in caplog.records): + await asyncio.sleep(0) + + assert service._task is not None + assert not service._task.done() + assert [task_id for task_id, _kwargs in repo.release_calls] == ["task-1"] + release_logs = [record for record in caplog.records if "MCP task batch release failed" in record.message] + assert len(release_logs) == 1 + assert "release notification failure" in release_logs[0].message + assert "task_id=task-1" in release_logs[0].message + finally: + await service.stop() + + +class SameTickNotificationFailureReleaseRepo(NotificationFailureReleaseRepo): + def __init__(self): + super().__init__(release_error=asyncio.CancelledError("notification release cancelled itself")) + self.caller_task = None + + async def release_notification_lease(self, task_id, **kwargs): + self.release_calls.append((task_id, kwargs)) + assert self.caller_task is not None + self.caller_task.cancel("same tick notification cancellation") + self.release_finished.set() + raise asyncio.CancelledError("notification release cancelled itself") + + +@pytest.mark.asyncio +async def test_notification_failure_release_same_tick_outer_cancellation_wins(caplog): + repo = SameTickNotificationFailureReleaseRepo() + service = McpTaskService( + repository=repo, + drivers=McpTaskDriverRegistry(), + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_polls=3, + launch_notification=AsyncMock(return_value={"run_id": "notify-run-1"}), + get_run=AsyncMock(side_effect=RuntimeError("run store unavailable")), + ) + caller = asyncio.create_task(service._run_notifications(now=datetime.now(UTC))) + repo.caller_task = caller + + with caplog.at_level(logging.ERROR), pytest.raises(asyncio.CancelledError) as caught: + await caller + + assert caught.value.args == ("same tick notification cancellation",) + assert repo.release_finished.is_set() + assert [task_id for task_id, _kwargs in repo.release_calls] == ["task-1"] + + +@pytest.mark.asyncio +async def test_notification_cancellation_during_source_run_lookup_releases_claim(): + lookup_started = asyncio.Event() + + async def get_run(*_args, **_kwargs): + lookup_started.set() + await asyncio.Event().wait() + + repo = SimpleNamespace(release_notification_claim=AsyncMock(return_value=True)) + service = McpTaskService( + repository=repo, + drivers=McpTaskDriverRegistry(), + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_polls=3, + launch_notification=AsyncMock(return_value={"run_id": "notify-run-1"}), + get_run=get_run, + ) + record = { + **_claimed_row(), + "notification_status": "claimed", + "dispatch_version": 2, + "dispatch_attempt": 0, + "dispatch_event": {"status": "completed"}, + } + + task = asyncio.create_task(service._notify_one(record, now=datetime.now(UTC))) + await lookup_started.wait() + task.cancel() + + with pytest.raises(asyncio.CancelledError): + await task + + repo.release_notification_claim.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_dispatched_notification_cancellation_preserves_phase_when_releasing_lease(): + lookup_started = asyncio.Event() + + async def get_run(*_args, **_kwargs): + lookup_started.set() + await asyncio.Event().wait() + + repo = SimpleNamespace( + release_notification_claim=AsyncMock(return_value=True), + release_notification_lease=AsyncMock(return_value=True), + ) + service = McpTaskService( + repository=repo, + drivers=McpTaskDriverRegistry(), + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_polls=3, + launch_notification=AsyncMock(), + get_run=get_run, + ) + record = { + **_claimed_row(), + "notification_status": "dispatched", + "notification_run_id": "notify-run-1", + "dispatch_version": 2, + } + + task = asyncio.create_task(service._notify_one(record, now=datetime.now(UTC))) + await lookup_started.wait() + task.cancel() + + with pytest.raises(asyncio.CancelledError): + await task + + repo.release_notification_lease.assert_awaited_once() + repo.release_notification_claim.assert_not_awaited() + + +class PollPersistenceRepo(FakeRepository): + def __init__(self, *, release_error: Exception | None = None): + super().__init__([_claimed_row()]) + self.apply_started = asyncio.Event() + self.release_error = release_error + self.cancelled_releases = [] + + async def apply_snapshot(self, task_id, **kwargs): + self.applied.append((task_id, kwargs)) + self.apply_started.set() + await asyncio.Event().wait() + + async def release_claim(self, task_id, **kwargs): + self.released.append((task_id, kwargs)) + if self.release_error is not None: + raise self.release_error + return True + + async def release_poll_claim_after_cancellation(self, task_id, **kwargs): + self.cancelled_releases.append((task_id, kwargs)) + if self.release_error is not None: + raise self.release_error + return True + + +@pytest.mark.asyncio +async def test_stop_releases_poll_claim_during_snapshot_persistence(): + repo = PollPersistenceRepo() + drivers = McpTaskDriverRegistry() + drivers.register("fake", FakeDriver(snapshots=[TaskSnapshot(status=TaskStatus.WORKING)])) + service = McpTaskService( + repository=repo, + drivers=drivers, + poll_interval_seconds=60, + lease_seconds=120, + max_concurrent_polls=3, + ) + + await service.start() + await asyncio.wait_for(repo.apply_started.wait(), timeout=1) + await asyncio.wait_for(service.stop(), timeout=1) + + assert repo.cancelled_releases + assert {task_id for task_id, _kwargs in repo.cancelled_releases} == {"task-1"} + assert repo.released == [] + + +@pytest.mark.asyncio +async def test_poll_cancellation_releases_only_the_current_poll_lease(): + repo = PollPersistenceRepo() + driver = HangingDriver() + drivers = McpTaskDriverRegistry() + drivers.register("fake", driver) + service = McpTaskService( + repository=repo, + drivers=drivers, + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_polls=3, + ) + + task = asyncio.create_task(service._poll_one(_claimed_row(), now=datetime.now(UTC))) + await driver.started.wait() + task.cancel() + + with pytest.raises(asyncio.CancelledError): + await task + + assert repo.cancelled_releases == [("task-1", {"lease_owner": service._lease_owner})] + assert repo.released == [] + + +@pytest.mark.asyncio +async def test_poll_cancellation_preserves_cancelled_error_when_release_fails(caplog): + repo = PollPersistenceRepo(release_error=RuntimeError("poll release unavailable")) + driver = HangingDriver() + drivers = McpTaskDriverRegistry() + drivers.register("fake", driver) + service = McpTaskService( + repository=repo, + drivers=drivers, + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_polls=3, + ) + + task = asyncio.create_task(service._poll_one(_claimed_row(), now=datetime.now(UTC))) + await driver.started.wait() + task.cancel() + + with caplog.at_level(logging.ERROR), pytest.raises(asyncio.CancelledError): + await task + + assert repo.cancelled_releases + assert repo.released == [] + assert "poll release unavailable" in caplog.text + + +async def _wait_for_compensation_tasks_to_clear(service: McpTaskService) -> None: + async with asyncio.timeout(1): + while service._compensation_tasks: + await asyncio.sleep(0) + + +@pytest.mark.asyncio +async def test_cancelled_hung_claim_returns_then_releases_delayed_result(monkeypatch): + monkeypatch.setattr(service_module, "_CANCELLATION_DRAIN_TIMEOUT_SECONDS", 0.01) + claim_started = asyncio.Event() + claim_gate = asyncio.Event() + release_calls = [] + + async def claim(): + claim_started.set() + await claim_gate.wait() + return [_claimed_row()] + + async def release(record): + release_calls.append(record["id"]) + + service = McpTaskService( + repository=SimpleNamespace(), + drivers=McpTaskDriverRegistry(), + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_polls=3, + ) + caller = asyncio.create_task( + service._claim_with_cancellation_release( + claim(), + action="probe claim", + release=release, + ) + ) + await claim_started.wait() + caller.cancel() + + try: + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(asyncio.shield(caller), timeout=0.2) + + assert release_calls == [] + assert len(service._compensation_tasks) == 1 + + claim_gate.set() + await _wait_for_compensation_tasks_to_clear(service) + + assert release_calls == ["task-1"] + assert not service._compensation_tasks + finally: + claim_gate.set() + if not caller.done(): + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(caller, timeout=0.2) + await _wait_for_compensation_tasks_to_clear(service) + + +@pytest.mark.asyncio +async def test_batch_release_starts_sibling_when_first_release_hangs(): + first_release_gate = asyncio.Event() + second_release_completed = asyncio.Event() + release_calls = [] + + async def release(record): + release_calls.append(record["id"]) + if record["id"] == "first": + await first_release_gate.wait() + else: + second_release_completed.set() + + service = McpTaskService( + repository=SimpleNamespace(), + drivers=McpTaskDriverRegistry(), + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_polls=3, + ) + batch = asyncio.create_task( + service._release_claimed_records( + [ + {**_claimed_row(), "id": "first"}, + {**_claimed_row(), "id": "second"}, + ], + release=release, + ) + ) + + try: + await asyncio.wait_for(second_release_completed.wait(), timeout=0.2) + assert release_calls == ["first", "second"] + finally: + first_release_gate.set() + await asyncio.wait_for(batch, timeout=0.2) + + +@pytest.mark.asyncio +async def test_batch_release_logs_cancelled_record_and_finishes_sibling(caplog): + sibling_completed = asyncio.Event() + release_calls = [] + + async def release(record): + release_calls.append(record["id"]) + if record["id"] == "cancelled": + raise asyncio.CancelledError("release cancelled") + sibling_completed.set() + + service = McpTaskService( + repository=SimpleNamespace(), + drivers=McpTaskDriverRegistry(), + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_polls=3, + ) + + with caplog.at_level(logging.ERROR): + await service._release_claimed_records( + [ + {**_claimed_row(), "id": "cancelled"}, + {**_claimed_row(), "id": "sibling"}, + ], + release=release, + ) + + assert sibling_completed.is_set() + assert release_calls.count("cancelled") == 1 + assert release_calls.count("sibling") == 1 + cancellations = [record for record in caplog.records if "MCP task claim release was cancelled" in record.getMessage()] + assert len(cancellations) == 1 + assert "task_id=cancelled" in cancellations[0].getMessage() + + +@pytest.mark.asyncio +async def test_duplicate_compensation_registration_logs_failure_once(caplog): + service = McpTaskService( + repository=SimpleNamespace(), + drivers=McpTaskDriverRegistry(), + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_polls=3, + ) + compensation = asyncio.get_running_loop().create_future() + + with caplog.at_level(logging.ERROR): + service._track_compensation_task(compensation, action="release poll claim", task_id="task-1") + service._track_compensation_task(compensation, action="release poll claim", task_id="task-1") + assert service._compensation_tasks == {compensation} + + compensation.set_exception(RuntimeError("release remained unavailable")) + await _wait_for_compensation_tasks_to_clear(service) + + failures = [record for record in caplog.records if "MCP task cancellation operation failed" in record.getMessage()] + assert len(failures) == 1 + assert "release remained unavailable" in failures[0].getMessage() + + +@pytest.mark.asyncio +async def test_hung_cancellation_compensation_transfers_to_background(monkeypatch): + monkeypatch.setattr(service_module, "_CANCELLATION_DRAIN_TIMEOUT_SECONDS", 0.01) + release_started = asyncio.Event() + release_gate = asyncio.Event() + release_calls = [] + + async def release_poll_claim_after_cancellation(task_id, **_kwargs): + release_calls.append(task_id) + release_started.set() + await release_gate.wait() + return True + + driver = HangingDriver() + drivers = McpTaskDriverRegistry() + drivers.register("fake", driver) + service = McpTaskService( + repository=SimpleNamespace(release_poll_claim_after_cancellation=release_poll_claim_after_cancellation), + drivers=drivers, + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_polls=3, + ) + + task = asyncio.create_task(service._poll_one(_claimed_row(), now=datetime.now(UTC))) + await driver.started.wait() + task.cancel() + await release_started.wait() + + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, timeout=0.2) + + assert release_calls == ["task-1"] + assert len(service._compensation_tasks) == 1 + + release_gate.set() + await _wait_for_compensation_tasks_to_clear(service) + + assert not service._compensation_tasks + + +@pytest.mark.asyncio +async def test_background_compensation_failure_is_consumed(monkeypatch, caplog): + monkeypatch.setattr(service_module, "_CANCELLATION_DRAIN_TIMEOUT_SECONDS", 0.01) + release_started = asyncio.Event() + release_gate = asyncio.Event() + release_calls = [] + + async def release_poll_claim_after_cancellation(task_id, **_kwargs): + release_calls.append(task_id) + release_started.set() + await release_gate.wait() + raise RuntimeError("release remained unavailable") + + driver = HangingDriver() + drivers = McpTaskDriverRegistry() + drivers.register("fake", driver) + service = McpTaskService( + repository=SimpleNamespace(release_poll_claim_after_cancellation=release_poll_claim_after_cancellation), + drivers=drivers, + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_polls=3, + ) + + task = asyncio.create_task(service._poll_one(_claimed_row(), now=datetime.now(UTC))) + await driver.started.wait() + task.cancel() + await release_started.wait() + + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, timeout=0.2) + + assert release_calls == ["task-1"] + assert len(service._compensation_tasks) == 1 + + with caplog.at_level(logging.ERROR): + release_gate.set() + await _wait_for_compensation_tasks_to_clear(service) + + failures = [record for record in caplog.records if "MCP task cancellation operation failed" in record.getMessage()] + assert len(failures) == 1 + assert "release remained unavailable" in failures[0].getMessage() + assert not service._compensation_tasks + + +class BatchCancellationRepository(FakeRepository): + def __init__(self, rows, *, phase="cancel"): + super().__init__(rows) + self.phase = phase + self.cancel_releases = [] + self.caller_task = None + + async def claim_due_tasks(self, **_kwargs): + if self.phase != "poll": + return [] + return [dict(row) for row in self.rows] + + async def claim_cancel_requests(self, **_kwargs): + if self.phase != "cancel": + return [] + return [dict(row) for row in self.rows] + + async def release_cancel_claim(self, task_id, **kwargs): + self.cancel_releases.append((task_id, kwargs)) + return True + + +class OuterCancellingDriver(FakeDriver): + def __init__(self, *, caller_task, phase): + super().__init__() + self.caller_task = caller_task + self.phase = phase + self.started = [] + + async def _run(self, task): + self.started.append(task.local_task_id) + if task.local_task_id != "task-1": + raise AssertionError("task-2 should be released by the batch fallback") + self.caller_task.cancel() + await asyncio.Event().wait() + + async def get_status(self, task): + return await self._run(task) + + async def cancel(self, task): + return await self._run(task) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("phase", ["poll", "cancel"]) +async def test_batch_outer_cancellation_releases_started_and_never_started_once(phase): + rows = [ + _claimed_row(), + {**_claimed_row(), "id": "task-2", "remote_task_id": "remote-2"}, + ] + repo = BatchCancellationRepository(rows, phase=phase) + drivers = McpTaskDriverRegistry() + service = McpTaskService( + repository=repo, + drivers=drivers, + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_polls=3, + ) + caller = asyncio.create_task(service.run_once(now=datetime.now(UTC)) if phase == "poll" else service._run_cancellations(now=datetime.now(UTC))) + repo.caller_task = caller + driver = OuterCancellingDriver(caller_task=caller, phase=phase) + drivers.register("fake", driver) + + with pytest.raises(asyncio.CancelledError): + await caller + + released = repo.released if phase == "poll" else repo.cancel_releases + assert sorted(task_id for task_id, _kwargs in released) == ["task-1", "task-2"] + assert driver.started == ["task-1"] + + +class SelfCancellingDriver(FakeDriver): + async def get_status(self, task): + raise asyncio.CancelledError("child poll cancelled itself") + + async def cancel(self, task): + raise asyncio.CancelledError("child cancel cancelled itself") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("phase", ["poll", "cancel"]) +async def test_batch_child_self_cancellation_releases_once(phase): + repo = BatchCancellationRepository([_claimed_row()], phase=phase) + drivers = McpTaskDriverRegistry() + drivers.register("fake", SelfCancellingDriver()) + service = McpTaskService( + repository=repo, + drivers=drivers, + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_polls=3, + ) + + if phase == "poll": + await service.run_once(now=datetime.now(UTC)) + assert [task_id for task_id, _kwargs in repo.released] == ["task-1"] + else: + await service._run_cancellations(now=datetime.now(UTC)) + assert [task_id for task_id, _kwargs in repo.cancel_releases] == ["task-1"] + + +@pytest.mark.asyncio +async def test_batch_outer_cancellation_logs_unexpected_child_failure_once(caplog): + service = McpTaskService( + repository=SimpleNamespace(), + drivers=McpTaskDriverRegistry(), + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_polls=3, + ) + rows = [ + _claimed_row(), + {**_claimed_row(), "id": "task-2", "remote_task_id": "remote-2"}, + ] + child_started = {row["id"]: asyncio.Event() for row in rows} + release_finished = {row["id"]: asyncio.Event() for row in rows} + release_calls = [] + + async def operation(record): + child_started[record["id"]].set() + try: + await asyncio.Future() + except asyncio.CancelledError: + raise RuntimeError(f"child failed during cancellation handoff ({record['id']})") + + async def release(record): + release_calls.append(record["id"]) + release_finished[record["id"]].set() + + caller = asyncio.create_task( + service._run_claimed_batch( + rows, + operation=operation, + release=release, + action="poll", + ) + ) + await asyncio.gather(*(event.wait() for event in child_started.values())) + caller.cancel("outer cancellation") + + with caplog.at_level(logging.ERROR), pytest.raises(asyncio.CancelledError): + await caller + + failures = [record for record in caplog.records if "Unexpected MCP task poll failure" in record.getMessage()] + assert len(failures) == len(rows) + for row in rows: + task_id = row["id"] + assert release_finished[task_id].is_set() + assert release_calls.count(task_id) == 1 + matching_failures = [failure for failure in failures if f"task_id={task_id}" in failure.getMessage()] + assert len(matching_failures) == 1 + assert f"child failed during cancellation handoff ({task_id})" in caplog.text + + +class SelfCancellingNotificationRepository(FakeRepository): + def __init__(self): + super().__init__() + self.notification_releases = [] + + async def claim_notification_work(self, **_kwargs): + if self.claimed: + return [] + self.claimed = True + return [ + { + **_claimed_row(), + "notification_status": "claimed", + "dispatch_version": 1, + "dispatch_attempt": 0, + "dispatch_event": {"status": "completed"}, + } + ] + + async def release_notification_claim(self, task_id, **kwargs): + self.notification_releases.append((task_id, kwargs)) + return True + + +@pytest.mark.asyncio +async def test_notification_child_self_cancellation_releases_once(): + repo = SelfCancellingNotificationRepository() + service = McpTaskService( + repository=repo, + drivers=McpTaskDriverRegistry(), + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_polls=3, + launch_notification=AsyncMock(side_effect=asyncio.CancelledError("child notification cancelled itself")), + get_run=AsyncMock(return_value=None), + ) + + await service._run_notifications(now=datetime.now(UTC)) + + assert [task_id for task_id, _kwargs in repo.notification_releases] == ["task-1"] + + +class BatchNotificationRepository(FakeRepository): + def __init__(self, rows): + super().__init__(rows) + self.notification_releases = [] + + async def claim_notification_work(self, **_kwargs): + if self.claimed: + return [] + self.claimed = True + return [dict(row) for row in self.rows] + + async def release_notification_claim(self, task_id, **kwargs): + self.notification_releases.append((task_id, kwargs)) + return True + + +@pytest.mark.asyncio +async def test_notification_outer_cancellation_releases_started_and_never_started_once(): + rows = [ + { + **_claimed_row(), + "notification_status": "claimed", + "dispatch_version": 1, + "dispatch_attempt": 0, + "dispatch_event": {"status": "completed"}, + }, + { + **_claimed_row(), + "id": "task-2", + "remote_task_id": "remote-2", + "notification_status": "claimed", + "dispatch_version": 1, + "dispatch_attempt": 0, + "dispatch_event": {"status": "completed"}, + }, + ] + repo = BatchNotificationRepository(rows) + caller = None + release_gate = asyncio.Event() + launch_calls = [] + + async def launch_notification(**kwargs): + launch_calls.append(kwargs["task_id"]) + if kwargs["task_id"] != "task-1": + raise AssertionError("task-2 should be released by the batch fallback") + caller.cancel() + await release_gate.wait() + return {"run_id": "notify-run-1"} + + service = McpTaskService( + repository=repo, + drivers=McpTaskDriverRegistry(), + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_polls=3, + launch_notification=launch_notification, + get_run=AsyncMock(return_value=None), + ) + caller = asyncio.create_task(service._run_notifications(now=datetime.now(UTC))) + + with pytest.raises(asyncio.CancelledError): + await caller + + assert sorted(task_id for task_id, _kwargs in repo.notification_releases) == ["task-1", "task-2"] + assert launch_calls == ["task-1"] + release_gate.set() + + +class SuppressingBatchDriver(FakeDriver): + def __init__(self, *, release_gate): + super().__init__() + self.release_gate = release_gate + self.started = asyncio.Event() + self.swallowed = asyncio.Event() + + async def _run(self, task): + if task.local_task_id == "task-1": + self.started.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + self.swallowed.set() + await self.release_gate.wait() + if task.local_task_id == "task-1": + return TaskSnapshot(status=TaskStatus.WORKING) + return TaskSnapshot(status=TaskStatus.WORKING) + + async def get_status(self, task): + return await self._run(task) + + async def cancel(self, task): + return await self._run(task) + + +def _batch_probe_rows(*, notification=False, count=2): + rows = [] + for index in range(count): + row = { + **_claimed_row(), + "id": f"task-{index + 1}", + "remote_task_id": f"remote-{index + 1}", + } + if notification: + row.update( + notification_status="claimed", + dispatch_version=1, + dispatch_attempt=0, + dispatch_event={"status": "completed"}, + ) + rows.append(row) + return rows + + +async def _run_batch_probe(service, phase): + now = datetime.now(UTC) + if phase == "poll": + await service.run_once(now=now) + elif phase == "cancel": + await service._run_cancellations(now=now) + else: + await service._run_notifications(now=now) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("phase", ["poll", "cancel", "notification"]) +async def test_outer_cancel_returns_before_suppressing_child_and_releases_all_once(phase, monkeypatch): + monkeypatch.setattr(service_module, "_CANCELLATION_DRAIN_TIMEOUT_SECONDS", 0.05) + release_gate = asyncio.Event() + driver = SuppressingBatchDriver(release_gate=release_gate) + + if phase == "notification": + repo = BatchNotificationRepository(_batch_probe_rows(notification=True)) + + async def launch_notification(**kwargs): + if kwargs["task_id"] == "task-1": + driver.started.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + driver.swallowed.set() + await release_gate.wait() + return {"run_id": f"notify-{kwargs['task_id']}"} + + service = McpTaskService( + repository=repo, + drivers=McpTaskDriverRegistry(), + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_polls=3, + launch_notification=launch_notification, + get_run=AsyncMock(return_value=None), + ) + else: + repo = BatchCancellationRepository(_batch_probe_rows(), phase=phase) + drivers = McpTaskDriverRegistry() + drivers.register("fake", driver) + service = McpTaskService( + repository=repo, + drivers=drivers, + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_polls=3, + ) + + caller = asyncio.create_task(_run_batch_probe(service, phase)) + await driver.started.wait() + caller.cancel("first cancellation") + timed_out = False + try: + with pytest.raises(asyncio.CancelledError) as caught: + await asyncio.wait_for(caller, timeout=0.2) + assert caught.value.args == ("first cancellation",) + except TimeoutError: + timed_out = True + + assert timed_out is False + assert driver.swallowed.is_set() + if phase == "poll": + released = repo.released + elif phase == "cancel": + released = repo.cancel_releases + else: + released = repo.notification_releases + assert sorted(task_id for task_id, _kwargs in released) == ["task-1", "task-2"] + assert len(service._compensation_tasks) == 1 + + release_gate.set() + await _wait_for_compensation_tasks_to_clear(service) + assert len(released) == 2 + + +@pytest.mark.asyncio +async def test_started_child_swallowing_cancel_then_returning_still_releases_once(monkeypatch): + monkeypatch.setattr(service_module, "_CANCELLATION_DRAIN_TIMEOUT_SECONDS", 0.05) + release_gate = asyncio.Event() + driver = SuppressingBatchDriver(release_gate=release_gate) + repo = BatchCancellationRepository([_claimed_row()], phase="poll") + drivers = McpTaskDriverRegistry() + drivers.register("fake", driver) + service = McpTaskService( + repository=repo, + drivers=drivers, + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_polls=3, + ) + + caller = asyncio.create_task(service.run_once(now=datetime.now(UTC))) + await driver.started.wait() + caller.cancel("first cancellation") + await driver.swallowed.wait() + assert repo.released == [] + release_gate.set() + + with pytest.raises(asyncio.CancelledError) as caught: + await caller + assert caught.value.args == ("first cancellation",) + assert [task_id for task_id, _kwargs in repo.released] == ["task-1"] + await _wait_for_compensation_tasks_to_clear(service) + + +@pytest.mark.asyncio +async def test_repeated_batch_cancel_preserves_first_cancel_args(monkeypatch): + monkeypatch.setattr(service_module, "_CANCELLATION_DRAIN_TIMEOUT_SECONDS", 0.05) + release_gate = asyncio.Event() + driver = SuppressingBatchDriver(release_gate=release_gate) + repo = BatchCancellationRepository([_claimed_row()], phase="poll") + drivers = McpTaskDriverRegistry() + drivers.register("fake", driver) + service = McpTaskService( + repository=repo, + drivers=drivers, + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_polls=3, + ) + + caller = asyncio.create_task(service.run_once(now=datetime.now(UTC))) + await driver.started.wait() + caller.cancel("first cancellation") + await driver.swallowed.wait() + caller.cancel("second cancellation") + + try: + with pytest.raises(asyncio.CancelledError) as caught: + await asyncio.wait_for(caller, timeout=0.2) + assert caught.value.args == ("first cancellation",) + finally: + release_gate.set() + if not caller.done(): + with pytest.raises(asyncio.CancelledError): + await caller + + +@pytest.mark.asyncio +async def test_batch_outer_cancel_releases_duplicate_ids_by_position(monkeypatch): + monkeypatch.setattr(service_module, "_CANCELLATION_DRAIN_TIMEOUT_SECONDS", 0.05) + release_gate = asyncio.Event() + driver = SuppressingBatchDriver(release_gate=release_gate) + rows = [_claimed_row(), _claimed_row()] + repo = BatchCancellationRepository(rows, phase="poll") + drivers = McpTaskDriverRegistry() + drivers.register("fake", driver) + service = McpTaskService( + repository=repo, + drivers=drivers, + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_polls=3, + ) + + caller = asyncio.create_task(service.run_once(now=datetime.now(UTC))) + await driver.started.wait() + caller.cancel("first cancellation") + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(caller, timeout=0.2) + + assert [task_id for task_id, _kwargs in repo.released] == ["task-1", "task-1"] + release_gate.set() + await _wait_for_compensation_tasks_to_clear(service) + + +@pytest.mark.asyncio +async def test_batch_completion_cancel_race_releases_once_for_100_rounds(): + service = McpTaskService( + repository=SimpleNamespace(), + drivers=McpTaskDriverRegistry(), + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_polls=3, + ) + release_calls = [] + + async def operation(_record): + await asyncio.sleep(0) + return None + + async def release(record): + release_calls.append(record["position"]) + + for position in range(100): + caller = asyncio.create_task( + service._run_claimed_batch( + [{"id": "duplicate", "position": position}], + operation=operation, + release=release, + action="race", + ) + ) + await asyncio.sleep(0) + caller.cancel("race cancellation") + with pytest.raises(asyncio.CancelledError): + await caller + + assert sorted(release_calls) == list(range(100)) + + +class OrdinaryReleaseBatchRepository: + def __init__(self, *, phase, outcome): + self.phase = phase + self.outcome = outcome + self.release_started = asyncio.Event() + self.release_gate = asyncio.Event() + self.release_calls = [] + self.release_interrupted = False + self.release_completed = False + self.release_finished = asyncio.Event() + self.caller_task = None + + def _records(self, phase): + if self.phase != phase: + return [] + record = _claimed_row() + if phase == "notification": + record.update( + notification_status="claimed", + dispatch_version=1, + dispatch_attempt=0, + dispatch_event={"status": "completed"}, + ) + return [record] + + async def claim_due_tasks(self, **_kwargs): + return self._records("poll") + + async def claim_cancel_requests(self, **_kwargs): + return self._records("cancel") + + async def claim_notification_work(self, **_kwargs): + return self._records("notification") + + async def _release(self, task_id, **_kwargs): + self.release_calls.append(task_id) + self.release_started.set() + if self.outcome in {"same_tick", "same_tick_self_cancel"}: + assert self.caller_task is not None + self.caller_task.cancel("same tick cancellation") + if self.outcome == "same_tick_self_cancel": + self.release_finished.set() + raise asyncio.CancelledError("ordinary release cancelled itself") + self.release_completed = True + return True + try: + await self.release_gate.wait() + except asyncio.CancelledError: + self.release_interrupted = True + raise + if self.outcome == "failure": + raise RuntimeError("ordinary release unavailable") + if self.outcome == "self_cancel": + self.release_finished.set() + raise asyncio.CancelledError("ordinary release cancelled itself") + self.release_completed = True + return True + + async def release_claim(self, task_id, **kwargs): + return await self._release(task_id, **kwargs) + + async def release_cancel_claim(self, task_id, **kwargs): + return await self._release(task_id, **kwargs) + + async def release_notification_claim(self, task_id, **kwargs): + return await self._release(task_id, **kwargs) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("phase", ["poll", "cancel", "notification"]) +@pytest.mark.parametrize("outcome", ["success", "failure", "self_cancel"]) +async def test_ordinary_batch_release_is_handed_off_without_duplication(phase, outcome, monkeypatch, caplog): + monkeypatch.setattr(service_module, "_CANCELLATION_DRAIN_TIMEOUT_SECONDS", 0.01) + repo = OrdinaryReleaseBatchRepository(phase=phase, outcome=outcome) + drivers = McpTaskDriverRegistry() + if phase == "poll": + drivers.register("fake", FakeDriver(error=RuntimeError("poll failed"))) + elif phase == "cancel": + drivers.register("fake", FakeDriver(cancel_error=RuntimeError("cancel failed"))) + + service = McpTaskService( + repository=repo, + drivers=drivers, + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_polls=3, + launch_notification=AsyncMock(side_effect=ConflictError("thread busy")), + get_run=AsyncMock(return_value=None), + ) + caller = asyncio.create_task(_run_batch_probe(service, phase)) + await repo.release_started.wait() + caller.cancel("first cancellation") + + try: + with caplog.at_level(logging.ERROR), pytest.raises(asyncio.CancelledError) as caught: + await asyncio.wait_for(caller, timeout=0.2) + assert caught.value.args == ("first cancellation",) + assert repo.release_interrupted is False + assert repo.release_calls == ["task-1"] + assert len(service._compensation_tasks) == 1 + + repo.release_gate.set() + await _wait_for_compensation_tasks_to_clear(service) + assert repo.release_calls == ["task-1"] + if outcome == "success": + assert repo.release_completed is True + else: + assert repo.release_completed is False + if outcome == "failure": + assert "ordinary release unavailable" in caplog.text + else: + assert "MCP task batch release failed" in caplog.text + assert not service._compensation_tasks + finally: + repo.release_gate.set() + if not caller.done(): + with pytest.raises(asyncio.CancelledError): + await caller + await _wait_for_compensation_tasks_to_clear(service) + + +@pytest.mark.asyncio +async def test_ordinary_batch_release_completion_same_tick_is_terminal(monkeypatch): + repo = OrdinaryReleaseBatchRepository(phase="poll", outcome="same_tick") + drivers = McpTaskDriverRegistry() + drivers.register("fake", FakeDriver(error=RuntimeError("poll failed"))) + service = McpTaskService( + repository=repo, + drivers=drivers, + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_polls=3, + ) + caller = asyncio.create_task(_run_batch_probe(service, "poll")) + repo.caller_task = caller + + with pytest.raises(asyncio.CancelledError) as caught: + await caller + + assert caught.value.args == ("same tick cancellation",) + assert repo.release_calls == ["task-1"] + assert repo.release_completed is True + assert not service._compensation_tasks + + +@pytest.mark.asyncio +@pytest.mark.parametrize("phase", ["poll", "cancel", "notification"]) +async def test_repeated_outer_cancellation_keeps_one_ordinary_release(phase, monkeypatch): + monkeypatch.setattr(service_module, "_CANCELLATION_DRAIN_TIMEOUT_SECONDS", 0.01) + repo = OrdinaryReleaseBatchRepository(phase=phase, outcome="success") + drivers = McpTaskDriverRegistry() + if phase == "poll": + drivers.register("fake", FakeDriver(error=RuntimeError("poll failed"))) + elif phase == "cancel": + drivers.register("fake", FakeDriver(cancel_error=RuntimeError("cancel failed"))) + service = McpTaskService( + repository=repo, + drivers=drivers, + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_polls=3, + launch_notification=AsyncMock(side_effect=ConflictError("thread busy")), + get_run=AsyncMock(return_value=None), + ) + caller = asyncio.create_task(_run_batch_probe(service, phase)) + await repo.release_started.wait() + caller.cancel("first cancellation") + + async with asyncio.timeout(0.2): + while len(service._compensation_tasks) != 1: + await asyncio.sleep(0) + caller.cancel("second cancellation") + + with pytest.raises(asyncio.CancelledError) as caught: + await caller + assert caught.value.args == ("first cancellation",) + assert repo.release_calls == ["task-1"] + assert repo.release_interrupted is False + + repo.release_gate.set() + await _wait_for_compensation_tasks_to_clear(service) + assert repo.release_completed is True + + +@pytest.mark.asyncio +async def test_inner_ordinary_release_cancellation_does_not_kill_poller(caplog): + repo = OrdinaryReleaseBatchRepository(phase="poll", outcome="self_cancel") + drivers = McpTaskDriverRegistry() + drivers.register("fake", FakeDriver(error=RuntimeError("poll failed"))) + service = McpTaskService( + repository=repo, + drivers=drivers, + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_polls=3, + ) + + try: + with caplog.at_level(logging.ERROR): + await service.start() + await repo.release_started.wait() + repo.release_gate.set() + await repo.release_finished.wait() + await asyncio.sleep(0) + + assert service._task is not None + assert not service._task.done() + assert repo.release_calls == ["task-1"] + release_logs = [record for record in caplog.records if "MCP task batch release failed" in record.message] + assert len(release_logs) == 1 + assert "release poll retry" in release_logs[0].message + assert "task_id=task-1" in release_logs[0].message + finally: + await service.stop() + + +@pytest.mark.asyncio +async def test_terminal_ordinary_release_cancellation_is_consumed_once(caplog): + service = McpTaskService( + repository=SimpleNamespace(), + drivers=McpTaskDriverRegistry(), + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_polls=3, + ) + state = service_module._BatchRecordState(_claimed_row()) + task = asyncio.get_running_loop().create_future() + task.set_exception(asyncio.CancelledError("ordinary release cancelled itself")) + state.ordinary_release_task = task + token = service_module._current_batch_record.set(state) + try: + with caplog.at_level(logging.ERROR): + await service._release_ordinary_batch_record( + state.record, + release=AsyncMock(), + action="release poll retry", + ) + finally: + service_module._current_batch_record.reset(token) + + assert state.ordinary_release_terminal is True + release_logs = [record for record in caplog.records if "MCP task batch release failed" in record.message] + assert len(release_logs) == 1 + assert "release poll retry" in release_logs[0].message + assert "task_id=task-1" in release_logs[0].message + + +@pytest.mark.asyncio +async def test_same_tick_outer_cancellation_wins_over_inner_ordinary_release(caplog): + repo = OrdinaryReleaseBatchRepository(phase="poll", outcome="same_tick_self_cancel") + drivers = McpTaskDriverRegistry() + drivers.register("fake", FakeDriver(error=RuntimeError("poll failed"))) + service = McpTaskService( + repository=repo, + drivers=drivers, + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_polls=3, + ) + caller = asyncio.create_task(_run_batch_probe(service, "poll")) + repo.caller_task = caller + + with caplog.at_level(logging.ERROR), pytest.raises(asyncio.CancelledError) as caught: + await caller + + assert caught.value.args == ("same tick cancellation",) + assert repo.release_calls == ["task-1"] + assert repo.release_finished.is_set() + assert not service._compensation_tasks + + +class SelfCancellingClaimRepository: + def __init__(self): + self.claim_started = asyncio.Event() + self.claim_calls = 0 + self.release_calls = [] + + async def claim_due_tasks(self, **_kwargs): + self.claim_calls += 1 + self.claim_started.set() + raise asyncio.CancelledError("poll claim cancelled itself") + + async def release_claim(self, task_id, **kwargs): + self.release_calls.append((task_id, kwargs)) + return True + + +@pytest.mark.asyncio +async def test_inner_claim_cancellation_does_not_kill_poller_or_handoff(caplog): + repo = SelfCancellingClaimRepository() + service = McpTaskService( + repository=repo, + drivers=McpTaskDriverRegistry(), + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_polls=3, + ) + handoff = AsyncMock() + service._finish_cancelled_claim_handoff = handoff + + try: + with caplog.at_level(logging.ERROR): + await service.start() + await repo.claim_started.wait() + async with asyncio.timeout(1): + while not any("MCP task claim operation failed" in record.message for record in caplog.records): + await asyncio.sleep(0) + + assert service._task is not None + assert not service._task.done() + assert repo.claim_calls == 1 + assert repo.release_calls == [] + handoff.assert_not_awaited() + claim_logs = [record for record in caplog.records if "MCP task claim operation failed" in record.message] + assert len(claim_logs) == 1 + assert "poll claim" in claim_logs[0].message + assert "task_id=batch" in claim_logs[0].message + finally: + await service.stop() + + +@pytest.mark.asyncio +async def test_same_tick_outer_claim_cancellation_wins_and_preserves_args(): + caller = None + + async def claim(): + assert caller is not None + caller.cancel("same tick claim cancellation") + raise asyncio.CancelledError("claim cancelled itself") + + service = McpTaskService( + repository=SimpleNamespace(), + drivers=McpTaskDriverRegistry(), + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_polls=3, + ) + caller = asyncio.create_task( + service._claim_with_cancellation_release( + claim(), + action="poll claim", + release=AsyncMock(), + ) + ) + + with pytest.raises(asyncio.CancelledError) as caught: + await caller + + assert caught.value.args == ("same tick claim cancellation",) + assert not service._compensation_tasks diff --git a/backend/tests/test_run_journal.py b/backend/tests/test_run_journal.py index e531b9de0f5..24de4e2522c 100644 --- a/backend/tests/test_run_journal.py +++ b/backend/tests/test_run_journal.py @@ -422,6 +422,432 @@ def no_loop(): events = await store.list_events("t1", "r1") assert any(e["event_type"] == "llm.ai.response" for e in events) + @pytest.mark.anyio + async def test_threshold_flush_cancelled_before_start_requeues_detached_batch(self): + store = MemoryRunEventStore() + journal = RunJournal("r1", "t1", store, flush_threshold=1) + journal.record_delivery() + + flush_task = next(iter(journal._pending_flush_tasks)) + flush_task.cancel() + await asyncio.gather(flush_task, return_exceptions=True) + await asyncio.sleep(0) + + await journal.flush() + events = await store.list_events("t1", "r1") + assert [event["event_type"] for event in events] == ["run.delivery"] + + @pytest.mark.anyio + @pytest.mark.parametrize("background_flush", [False, True], ids=["explicit", "threshold"]) + async def test_cancelled_jsonl_flush_does_not_requeue_committed_batch(self, tmp_path, monkeypatch, background_flush): + from deerflow.runtime.events.store.jsonl import JsonlRunEventStore + + async def inline_to_thread(func, /, *args, **kwargs): + return func(*args, **kwargs) + + monkeypatch.setattr(asyncio, "to_thread", inline_to_thread) + store = JsonlRunEventStore(base_dir=tmp_path / "jsonl") + journal = RunJournal("r1", "t1", store, flush_threshold=1 if background_flush else 100) + write_completed = asyncio.Event() + allow_write_return = asyncio.Event() + original_write_batch = store._write_batch_async + + async def write_then_block_return(thread_id, batch): + records = await original_write_batch(thread_id, batch) + write_completed.set() + await allow_write_return.wait() + return records + + monkeypatch.setattr(store, "_write_batch_async", write_then_block_return) + journal.record_delivery() + + flush_task = next(iter(journal._pending_flush_tasks)) if background_flush else asyncio.create_task(journal.flush()) + await asyncio.wait_for(write_completed.wait(), timeout=1) + flush_task.cancel() + await asyncio.sleep(0) + flush_task.cancel() + allow_write_return.set() + + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(flush_task, timeout=1) + + await asyncio.wait_for(journal.flush(), timeout=1) + events = await store.list_events("t1", "r1") + assert [event["event_type"] for event in events] == ["run.delivery"] + + @pytest.mark.anyio + async def test_cancelled_flush_drains_failed_write_before_requeueing_batch(self): + class FailingStore: + def __init__(self): + self.write_started = asyncio.Event() + self.allow_failure = asyncio.Event() + self.write_cancelled = False + self.write_finished = False + + async def put_batch(self, _batch): + self.write_started.set() + try: + await self.allow_failure.wait() + except asyncio.CancelledError: + self.write_cancelled = True + raise + self.write_finished = True + raise RuntimeError("write failed") + + store = FailingStore() + journal = RunJournal("r1", "t1", store, flush_threshold=100) + journal.record_delivery() + flush_task = asyncio.create_task(journal.flush()) + await store.write_started.wait() + flush_task.cancel() + await asyncio.sleep(0) + flush_task.cancel() + store.allow_failure.set() + + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(flush_task, timeout=1) + + assert store.write_finished is True + assert store.write_cancelled is False + assert [event["event_type"] for event in journal._buffer] == ["run.delivery"] + + @pytest.mark.anyio + @pytest.mark.parametrize("eventual_error", [False, True], ids=["late-success", "late-failure"]) + async def test_cancelled_hung_write_transfers_ownership_without_blind_requeue(self, monkeypatch, caplog, eventual_error): + import deerflow.runtime.journal as journal_module + + monkeypatch.setattr(journal_module, "_CANCELLATION_DRAIN_TIMEOUT_SECONDS", 0.01, raising=False) + + class HangingStore: + def __init__(self): + self.started = asyncio.Event() + self.finish = asyncio.Event() + self.calls = 0 + self.cancelled = False + + async def put_batch(self, _batch): + self.calls += 1 + self.started.set() + try: + await self.finish.wait() + except asyncio.CancelledError: + self.cancelled = True + raise + if eventual_error: + raise RuntimeError("journal write failed late") + return [] + + store = HangingStore() + journal = RunJournal("r1", "t1", store, flush_threshold=100) + journal.record_delivery() + flush_task = asyncio.create_task(journal.flush()) + try: + await store.started.wait() + flush_task.cancel() + + done, _ = await asyncio.wait({flush_task}, timeout=0.2) + assert flush_task in done + result = await asyncio.gather(flush_task, return_exceptions=True) + assert isinstance(result[0], asyncio.CancelledError) + assert journal._buffer == [] + assert len(journal._detached_write_tasks) == 1 + assert store.calls == 1 + assert store.cancelled is False + + store.finish.set() + detached = tuple(journal._detached_write_tasks) + await asyncio.gather(*detached, return_exceptions=True) + await asyncio.sleep(0) + assert journal._detached_write_tasks == {} + if eventual_error: + assert len(journal._buffer) == 1 + assert sum("Detached journal write failed" in record.message for record in caplog.records) == 1 + else: + assert journal._buffer == [] + finally: + store.finish.set() + await asyncio.gather(flush_task, return_exceptions=True) + detached = tuple(getattr(journal, "_detached_write_tasks", ())) + if detached: + await asyncio.gather(*detached, return_exceptions=True) + + @pytest.mark.anyio + async def test_eventual_underlying_write_cancellation_requeues_once(self, monkeypatch, caplog): + import deerflow.runtime.journal as journal_module + + monkeypatch.setattr(journal_module, "_CANCELLATION_DRAIN_TIMEOUT_SECONDS", 0.01, raising=False) + + class SelfCancellingStore: + def __init__(self): + self.started = asyncio.Event() + self.finish = asyncio.Event() + self.calls = 0 + + async def put_batch(self, _batch): + self.calls += 1 + self.started.set() + await self.finish.wait() + raise asyncio.CancelledError + + store = SelfCancellingStore() + journal = RunJournal("r1", "t1", store, flush_threshold=100) + journal.record_delivery() + flush_task = asyncio.create_task(journal.flush()) + try: + await store.started.wait() + flush_task.cancel() + done, _ = await asyncio.wait({flush_task}, timeout=0.2) + assert flush_task in done + result = await asyncio.gather(flush_task, return_exceptions=True) + assert isinstance(result[0], asyncio.CancelledError) + assert journal._buffer == [] + assert len(journal._detached_write_tasks) == 1 + + store.finish.set() + detached = tuple(journal._detached_write_tasks) + await asyncio.gather(*detached, return_exceptions=True) + await asyncio.sleep(0) + assert journal._detached_write_tasks == {} + assert len(journal._buffer) == 1 + assert sum("Detached journal write failed" in record.message for record in caplog.records) == 1 + assert store.calls == 1 + finally: + store.finish.set() + await asyncio.gather(flush_task, return_exceptions=True) + detached = tuple(getattr(journal, "_detached_write_tasks", ())) + if detached: + await asyncio.gather(*detached, return_exceptions=True) + + @pytest.mark.anyio + async def test_detached_write_serializes_threshold_and_explicit_flushes(self, monkeypatch): + import deerflow.runtime.journal as journal_module + + monkeypatch.setattr(journal_module, "_CANCELLATION_DRAIN_TIMEOUT_SECONDS", 0.01, raising=False) + + class CountingStore: + def __init__(self): + self.calls = 0 + self.first_started = asyncio.Event() + self.first_finish = asyncio.Event() + self.second_started = asyncio.Event() + self.second_finish = asyncio.Event() + + async def put_batch(self, _batch): + self.calls += 1 + if self.calls == 1: + self.first_started.set() + await self.first_finish.wait() + else: + self.second_started.set() + await self.second_finish.wait() + return [] + + store = CountingStore() + journal = RunJournal("r1", "t1", store, flush_threshold=100) + journal.record_delivery() + first_flush = asyncio.create_task(journal.flush()) + try: + await store.first_started.wait() + first_flush.cancel() + done, _ = await asyncio.wait({first_flush}, timeout=0.2) + assert first_flush in done + first_result = await asyncio.gather(first_flush, return_exceptions=True) + assert isinstance(first_result[0], asyncio.CancelledError) + assert len(journal._detached_write_tasks) == 1 + + journal.record_delivery() + journal._flush_sync() + await asyncio.sleep(0) + assert store.calls == 1 + assert len(journal._buffer) == 1 + + later_flush = asyncio.create_task(journal.flush()) + await asyncio.sleep(0) + assert store.calls == 1 + + store.first_finish.set() + await asyncio.wait_for(store.second_started.wait(), timeout=0.2) + assert store.calls == 2 + store.second_finish.set() + await later_flush + assert journal._detached_write_tasks == {} + assert journal._buffer == [] + finally: + store.first_finish.set() + store.second_finish.set() + await asyncio.gather(first_flush, return_exceptions=True) + if "later_flush" in locals(): + await asyncio.gather(later_flush, return_exceptions=True) + detached = tuple(getattr(journal, "_detached_write_tasks", ())) + if detached: + await asyncio.gather(*detached, return_exceptions=True) + + @pytest.mark.anyio + async def test_later_flush_cancellation_keeps_detached_predecessor_and_does_not_overtake(self, monkeypatch): + import deerflow.runtime.journal as journal_module + + monkeypatch.setattr(journal_module, "_CANCELLATION_DRAIN_TIMEOUT_SECONDS", 0.01, raising=False) + + class BlockingStore: + def __init__(self): + self.calls = 0 + self.started = asyncio.Event() + self.finish = asyncio.Event() + + async def put_batch(self, _batch): + self.calls += 1 + self.started.set() + await self.finish.wait() + return [] + + store = BlockingStore() + journal = RunJournal("r1", "t1", store, flush_threshold=100) + journal.record_delivery() + first_flush = asyncio.create_task(journal.flush()) + later_flush = None + try: + await store.started.wait() + first_flush.cancel() + done, _ = await asyncio.wait({first_flush}, timeout=0.2) + assert first_flush in done + first_result = await asyncio.gather(first_flush, return_exceptions=True) + assert isinstance(first_result[0], asyncio.CancelledError) + assert len(journal._detached_write_tasks) == 1 + + journal.record_delivery() + later_flush = asyncio.create_task(journal.flush()) + await asyncio.sleep(0) + later_flush.cancel() + done, _ = await asyncio.wait({later_flush}, timeout=0.2) + assert later_flush in done + later_result = await asyncio.gather(later_flush, return_exceptions=True) + assert isinstance(later_result[0], asyncio.CancelledError) + assert store.calls == 1 + assert len(journal._detached_write_tasks) == 1 + assert len(journal._buffer) == 1 + finally: + store.finish.set() + await asyncio.gather(first_flush, return_exceptions=True) + if later_flush is not None: + await asyncio.gather(later_flush, return_exceptions=True) + detached = tuple(getattr(journal, "_detached_write_tasks", ())) + if detached: + await asyncio.gather(*detached, return_exceptions=True) + + @pytest.mark.anyio + async def test_journal_write_cancellation_uses_one_absolute_drain_deadline(self, monkeypatch): + import deerflow.runtime.journal as journal_module + + monkeypatch.setattr(journal_module, "_CANCELLATION_DRAIN_TIMEOUT_SECONDS", 0.01, raising=False) + observed_deadlines = [] + real_wait_for_task_until = journal_module.wait_for_task_until + + async def recording_wait(task, *, deadline): + observed_deadlines.append(deadline) + return await real_wait_for_task_until(task, deadline=deadline) + + monkeypatch.setattr(journal_module, "wait_for_task_until", recording_wait) + + class StalledStore: + def __init__(self): + self.started = asyncio.Event() + self.finish = asyncio.Event() + + async def put_batch(self, _batch): + self.started.set() + await self.finish.wait() + return [] + + store = StalledStore() + journal = RunJournal("r1", "t1", store, flush_threshold=100) + journal.record_delivery() + flush_task = asyncio.create_task(journal.flush()) + try: + await store.started.wait() + flush_task.cancel() + await asyncio.sleep(0) + flush_task.cancel() + done, _ = await asyncio.wait({flush_task}, timeout=0.2) + assert flush_task in done + result = await asyncio.gather(flush_task, return_exceptions=True) + assert isinstance(result[0], asyncio.CancelledError) + assert len(observed_deadlines) == 1 + assert len(journal._detached_write_tasks) == 1 + finally: + store.finish.set() + await asyncio.gather(flush_task, return_exceptions=True) + detached = tuple(getattr(journal, "_detached_write_tasks", ())) + if detached: + await asyncio.gather(*detached, return_exceptions=True) + + @pytest.mark.anyio + @pytest.mark.parametrize("eventual_error", [False, True], ids=["threshold-late-success", "threshold-late-failure"]) + async def test_threshold_detached_predecessor_blocks_foreground_flush_until_resolution(self, monkeypatch, eventual_error): + import deerflow.runtime.journal as journal_module + + monkeypatch.setattr(journal_module, "_CANCELLATION_DRAIN_TIMEOUT_SECONDS", 0.01, raising=False) + + class OrderedStore: + def __init__(self): + self.calls: list[list[str]] = [] + self.first_started = asyncio.Event() + self.first_finish = asyncio.Event() + + async def put_batch(self, batch): + self.calls.append([event["event_type"] for event in batch]) + if len(self.calls) == 1: + self.first_started.set() + await self.first_finish.wait() + if eventual_error: + raise RuntimeError("threshold journal write failed late") + return [] + return [] + + store = OrderedStore() + journal = RunJournal("r1", "t1", store, flush_threshold=1) + journal._put(event_type="first", category="trace", content="first") + threshold_task = next(iter(journal._pending_flush_tasks)) + foreground = None + try: + await store.first_started.wait() + assert threshold_task.done() is False + + journal._flush_threshold = 100 + journal._put(event_type="second", category="trace", content="second") + assert [event["event_type"] for event in journal._buffer] == ["second"] + + foreground = asyncio.create_task(journal.flush()) + await asyncio.sleep(0) + assert foreground.done() is False + + threshold_task.cancel() + threshold_result = await asyncio.gather(threshold_task, return_exceptions=True) + assert isinstance(threshold_result[0], asyncio.CancelledError) + await asyncio.sleep(0) + + done, _ = await asyncio.wait({foreground}, timeout=0.05) + assert foreground not in done + assert store.calls == [["first"]] + assert len(journal._detached_write_tasks) == 1 + + store.first_finish.set() + await asyncio.wait_for(foreground, timeout=0.2) + if eventual_error: + assert store.calls == [["first"], ["first", "second"]] + else: + assert store.calls == [["first"], ["second"]] + assert journal._buffer == [] + assert journal._detached_write_tasks == {} + assert journal._pending_flush_tasks == set() + finally: + store.first_finish.set() + await asyncio.gather(threshold_task, return_exceptions=True) + if foreground is not None: + await asyncio.gather(foreground, return_exceptions=True) + detached = tuple(getattr(journal, "_detached_write_tasks", ())) + if detached: + await asyncio.gather(*detached, return_exceptions=True) + class TestIdentifyCaller: def test_lead_agent_tag(self, journal_setup): diff --git a/backend/tests/test_run_manager.py b/backend/tests/test_run_manager.py index 61abacef72d..febd625c25f 100644 --- a/backend/tests/test_run_manager.py +++ b/backend/tests/test_run_manager.py @@ -9,6 +9,7 @@ import pytest from sqlalchemy.exc import DatabaseError as SQLAlchemyDatabaseError +import deerflow.runtime.runs.manager as manager_module from deerflow.config.run_ownership_config import RunOwnershipConfig from deerflow.runtime import DisconnectMode, RunManager, RunStatus, ThreadOperationKind from deerflow.runtime.events.store.memory import MemoryRunEventStore @@ -129,6 +130,55 @@ async def update_lease(self, run_id, *, owner_worker_id, lease_expires_at): return False +class BlockingFinalizationRunStore(MemoryRunStore): + def __init__(self) -> None: + super().__init__() + self.finalization_started = asyncio.Event() + + async def finalize_if_not_cancelled(self, *args, **kwargs): + self.finalization_started.set() + await asyncio.Event().wait() + + +class FailingFinalizationRunStore(MemoryRunStore): + def __init__(self) -> None: + super().__init__() + self.finalization_failed = asyncio.Event() + + async def finalize_if_not_cancelled(self, *args, **kwargs): + self.finalization_failed.set() + raise RuntimeError("finalization unavailable") + + +class CancellationResistantLock: + def __init__(self, outcome: str = "acquired") -> None: + self.acquire_started = asyncio.Event() + self.allow_acquire = asyncio.Event() + self.outcome = outcome + self.cancel_count = 0 + self.release_count = 0 + self.acquired = False + + async def acquire(self) -> bool: + self.acquire_started.set() + try: + await self.allow_acquire.wait() + except asyncio.CancelledError: + self.cancel_count += 1 + await self.allow_acquire.wait() + if self.outcome == "exception": + raise RuntimeError("late lock waiter failure") + if self.outcome == "cancelled": + raise asyncio.CancelledError("late lock waiter cancellation") + self.acquired = True + return True + + def release(self) -> None: + assert self.acquired + self.acquired = False + self.release_count += 1 + + async def _stored_statuses(store: MemoryRunStore, *run_ids: str) -> dict[str, Any]: rows = {} for run_id in run_ids: @@ -137,6 +187,1040 @@ async def _stored_statuses(store: MemoryRunStore, *run_ids: str) -> dict[str, An return rows +@pytest.mark.anyio +async def test_repeated_cancellation_during_finalization_still_fences_local_run(): + store = BlockingFinalizationRunStore() + manager = RunManager( + store=store, + run_ownership_config=RunOwnershipConfig( + lease_seconds=30, + grace_seconds=10, + heartbeat_enabled=True, + ), + ) + record = await manager.create("thread-1") + await manager.set_status(record.run_id, RunStatus.running) + task = asyncio.create_task(manager.set_status_if_not_cancelled(record.run_id, RunStatus.success)) + await store.finalization_started.wait() + await manager._lock.acquire() + task.cancel() + await asyncio.sleep(0) + task.cancel() + manager._lock.release() + + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, timeout=1) + + assert record.ownership_lost is True + assert record.abort_event.is_set() is True + assert record.status == RunStatus.error + + +@pytest.mark.anyio +async def test_finalization_fence_failure_preserves_caller_cancellation(monkeypatch, caplog): + store = BlockingFinalizationRunStore() + manager = RunManager( + store=store, + run_ownership_config=RunOwnershipConfig( + lease_seconds=30, + grace_seconds=10, + heartbeat_enabled=True, + ), + ) + record = await manager.create("thread-1") + await manager.set_status(record.run_id, RunStatus.running) + + async def fail_fence(*_args, **_kwargs): + raise RuntimeError("fence unavailable") + + monkeypatch.setattr(manager, "_mark_ownership_lost", fail_fence) + task = asyncio.create_task(manager.set_status_if_not_cancelled(record.run_id, RunStatus.success)) + await store.finalization_started.wait() + task.cancel() + + with caplog.at_level(logging.ERROR), pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, timeout=1) + + assert "fence unavailable" in caplog.text + + +@pytest.mark.anyio +async def test_cancellation_during_normal_error_finalization_fence_still_fences_local_run(): + store = FailingFinalizationRunStore() + manager = RunManager( + store=store, + run_ownership_config=RunOwnershipConfig( + lease_seconds=30, + grace_seconds=10, + heartbeat_enabled=True, + ), + persistence_retry_policy=PersistenceRetryPolicy(max_attempts=1, initial_delay=0), + ) + record = await manager.create("thread-1") + await manager.set_status(record.run_id, RunStatus.running) + await manager._lock.acquire() + task = asyncio.create_task(manager.set_status_if_not_cancelled(record.run_id, RunStatus.success)) + await store.finalization_failed.wait() + await asyncio.sleep(0) + task.cancel() + manager._lock.release() + + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, timeout=1) + + assert record.ownership_lost is True + assert record.abort_event.is_set() is True + + +@pytest.mark.anyio +async def test_hung_finalization_fence_transfers_to_background(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(manager_module, "_CANCELLATION_DRAIN_TIMEOUT_SECONDS", 0.01) + store = BlockingFinalizationRunStore() + manager = RunManager( + store=store, + run_ownership_config=RunOwnershipConfig( + lease_seconds=30, + grace_seconds=10, + heartbeat_enabled=True, + ), + ) + record = await manager.create("thread-1") + await manager.set_status(record.run_id, RunStatus.running) + fence_started = asyncio.Event() + finish_fence = asyncio.Event() + fence_calls = 0 + original_mark = manager._mark_ownership_lost + + async def hung_mark(*args, **kwargs): + nonlocal fence_calls + fence_calls += 1 + fence_started.set() + await finish_fence.wait() + await original_mark(*args, **kwargs) + + monkeypatch.setattr(manager, "_mark_ownership_lost", hung_mark) + caller = asyncio.create_task(manager.set_status_if_not_cancelled(record.run_id, RunStatus.success)) + await store.finalization_started.wait() + caller.cancel() + await fence_started.wait() + + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(caller, timeout=0.2) + assert len(manager._cancellation_cleanup_tasks) == 1 + + finish_fence.set() + await asyncio.gather(*tuple(manager._cancellation_cleanup_tasks), return_exceptions=True) + await asyncio.sleep(0) + assert record.ownership_lost is True + assert fence_calls == 1 + assert not manager._cancellation_cleanup_tasks + + +@pytest.mark.anyio +async def test_hung_cancelled_admission_cleanup_transfers_to_background(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(manager_module, "_CANCELLATION_DRAIN_TIMEOUT_SECONDS", 0.01) + 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() + cleanup_started = asyncio.Event() + finish_cleanup = asyncio.Event() + cleanup_calls = 0 + original_persist = manager._persist_status + original_close = manager._close_cancelled_admission + + async def block_old_persist(record, status, **kwargs): + if record.run_id == old.run_id: + old_persist_started.set() + await release_old_persist.wait() + return await original_persist(record, status, **kwargs) + + async def hung_close(record): + nonlocal cleanup_calls + cleanup_calls += 1 + cleanup_started.set() + await finish_cleanup.wait() + await original_close(record) + + monkeypatch.setattr(manager, "_persist_status", block_old_persist) + monkeypatch.setattr(manager, "_close_cancelled_admission", hung_close) + caller = asyncio.create_task(manager.create_or_reject("thread-1", multitask_strategy="interrupt")) + await old_persist_started.wait() + replacement = next(record for record in manager._runs.values() if record.run_id != old.run_id) + caller.cancel() + release_old_persist.set() + await cleanup_started.wait() + + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(caller, timeout=0.2) + assert len(manager._cancellation_cleanup_tasks) == 1 + + finish_cleanup.set() + await asyncio.gather(*tuple(manager._cancellation_cleanup_tasks), return_exceptions=True) + await asyncio.sleep(0) + assert replacement.status == RunStatus.interrupted + assert cleanup_calls == 1 + assert not manager._cancellation_cleanup_tasks + + +@pytest.mark.anyio +@pytest.mark.parametrize("cleanup_error", [RuntimeError("fence unavailable"), asyncio.CancelledError()], ids=["failure", "self-cancel"]) +async def test_late_cancellation_cleanup_outcome_is_consumed_without_replacing_cancel( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + cleanup_error: BaseException, +): + monkeypatch.setattr(manager_module, "_CANCELLATION_DRAIN_TIMEOUT_SECONDS", 0.01) + store = BlockingFinalizationRunStore() + manager = RunManager( + store=store, + run_ownership_config=RunOwnershipConfig( + lease_seconds=30, + grace_seconds=10, + heartbeat_enabled=True, + ), + ) + record = await manager.create("thread-1") + await manager.set_status(record.run_id, RunStatus.running) + finish_cleanup = asyncio.Event() + + async def late_cleanup(*_args, **_kwargs): + await finish_cleanup.wait() + raise cleanup_error + + monkeypatch.setattr(manager, "_mark_ownership_lost", late_cleanup) + caller = asyncio.create_task(manager.set_status_if_not_cancelled(record.run_id, RunStatus.success)) + await store.finalization_started.wait() + caller.cancel() + + with caplog.at_level(logging.ERROR), pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(caller, timeout=0.2) + assert len(manager._cancellation_cleanup_tasks) == 1 + + finish_cleanup.set() + cleanup = next(iter(manager._cancellation_cleanup_tasks)) + await asyncio.gather(cleanup, return_exceptions=True) + await asyncio.sleep(0) + assert not manager._cancellation_cleanup_tasks + assert caplog.text.count(f"run_id={record.run_id}") == 1 + assert "Run cancellation cleanup" in caplog.text + + +@pytest.mark.anyio +async def test_repeated_cancellation_uses_one_cleanup_task_and_absolute_deadline(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(manager_module, "_CANCELLATION_DRAIN_TIMEOUT_SECONDS", 0.03) + store = BlockingFinalizationRunStore() + manager = RunManager( + store=store, + run_ownership_config=RunOwnershipConfig( + lease_seconds=30, + grace_seconds=10, + heartbeat_enabled=True, + ), + ) + record = await manager.create("thread-1") + await manager.set_status(record.run_id, RunStatus.running) + finish_cleanup = asyncio.Event() + cleanup_started = asyncio.Event() + cleanup_calls = 0 + observed = [] + original_wait = manager_module.wait_for_task_until + + async def record_wait(task, *, deadline): + observed.append((task, deadline)) + return await original_wait(task, deadline=deadline) + + async def late_cleanup(*_args, **_kwargs): + nonlocal cleanup_calls + cleanup_calls += 1 + cleanup_started.set() + await finish_cleanup.wait() + + monkeypatch.setattr(manager_module, "wait_for_task_until", record_wait) + monkeypatch.setattr(manager, "_mark_ownership_lost", late_cleanup) + caller = asyncio.create_task(manager.set_status_if_not_cancelled(record.run_id, RunStatus.success)) + await store.finalization_started.wait() + caller.cancel() + await cleanup_started.wait() + await asyncio.sleep(0.005) + caller.cancel() + + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(caller, timeout=0.2) + assert len(observed) == 1 + assert len(manager._cancellation_cleanup_tasks) == 1 + cleanup = observed[0][0] + finish_cleanup.set() + await asyncio.gather(cleanup, return_exceptions=True) + await asyncio.sleep(0) + assert cleanup_calls == 1 + assert not manager._cancellation_cleanup_tasks + + +@pytest.mark.anyio +async def test_shutdown_observes_supervised_cleanup_only_within_budget(caplog: pytest.LogCaptureFixture): + manager = RunManager() + finish_cleanup = asyncio.Event() + + async def wait_for_cleanup() -> None: + await finish_cleanup.wait() + + cleanup = asyncio.create_task(wait_for_cleanup()) + manager._track_cancellation_cleanup(cleanup, action="test cleanup", run_id="run-1") + loop = asyncio.get_running_loop() + started = loop.time() + + with caplog.at_level(logging.WARNING): + await manager.shutdown(timeout=0.01) + + assert loop.time() - started < 0.2 + assert cleanup.done() is False + assert cleanup in manager._cancellation_cleanup_tasks + assert "cancellation cleanup task" in caplog.text + finish_cleanup.set() + await cleanup + await asyncio.sleep(0) + assert not manager._cancellation_cleanup_tasks + + +@pytest.mark.anyio +async def test_shutdown_observes_cleanup_registered_by_cancelled_run_before_deadline(): + manager = RunManager() + cleanup_finished = asyncio.Event() + cleanup_registered = asyncio.Event() + release_producer = manager._begin_cancellation_cleanup_producer() + + async def run() -> None: + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + cleanup = asyncio.create_task(cleanup_finished.wait()) + manager._track_cancellation_cleanup(cleanup, action="dynamic cleanup", run_id="run-1") + cleanup_registered.set() + finally: + release_producer() + + record = await manager.create("thread-1") + record.task = asyncio.create_task(run()) + shutdown_task = asyncio.create_task(manager.shutdown(timeout=0.2)) + await cleanup_registered.wait() + assert shutdown_task.done() is False + cleanup_finished.set() + await shutdown_task + assert not manager._cancellation_cleanup_tasks + + +@pytest.mark.anyio +async def test_shutdown_keeps_dynamic_cleanup_supervised_after_deadline(caplog: pytest.LogCaptureFixture): + manager = RunManager() + cleanup_registered = asyncio.Event() + cleanup_finished = asyncio.Event() + release_producer = manager._begin_cancellation_cleanup_producer() + + async def run() -> None: + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + cleanup = asyncio.create_task(cleanup_finished.wait()) + manager._track_cancellation_cleanup(cleanup, action="stalled dynamic cleanup", run_id="run-1") + cleanup_registered.set() + finally: + release_producer() + + record = await manager.create("thread-1") + record.task = asyncio.create_task(run()) + with caplog.at_level(logging.WARNING): + shutdown_task = asyncio.create_task(manager.shutdown(timeout=0.05)) + await cleanup_registered.wait() + await shutdown_task + + assert any("cancellation cleanup task" in message for message in caplog.messages) + cleanup = next(iter(manager._cancellation_cleanup_tasks)) + assert cleanup.done() is False + cleanup_finished.set() + await cleanup + await asyncio.sleep(0) + assert not manager._cancellation_cleanup_tasks + + +@pytest.mark.anyio +async def test_shutdown_waits_for_delayed_run_completion_cleanup_registration(): + manager = RunManager() + cleanup_finished = asyncio.Event() + cleanup_registered = asyncio.Event() + loop = asyncio.get_running_loop() + release_producer = manager._begin_cancellation_cleanup_producer() + + async def register_cleanup() -> None: + await asyncio.sleep(0) + cleanup = asyncio.create_task(cleanup_finished.wait()) + manager._track_cancellation_cleanup(cleanup, action="delayed cleanup", run_id="run-1") + cleanup_registered.set() + release_producer() + + async def run() -> None: + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + return + + record = await manager.create("thread-1") + record.task = asyncio.create_task(run()) + + def schedule_register() -> None: + asyncio.create_task(register_cleanup()) + + def second_barrier() -> None: + loop.call_soon(schedule_register) + + def first_barrier() -> None: + loop.call_soon(second_barrier) + + def after_run(_completed: asyncio.Future[None]) -> None: + loop.call_soon(first_barrier) + + record.task.add_done_callback(after_run) + shutdown_task = asyncio.create_task(manager.shutdown(timeout=0.05)) + await asyncio.wait_for(cleanup_registered.wait(), timeout=0.2) + assert shutdown_task.done() is False + cleanup = next(iter(manager._cancellation_cleanup_tasks)) + cleanup_finished.set() + await shutdown_task + await asyncio.sleep(0) + assert cleanup.done() is True + assert not manager._cancellation_cleanup_tasks + + +@pytest.mark.anyio +async def test_shutdown_keeps_delayed_stalled_cleanup_supervised_until_deadline(caplog: pytest.LogCaptureFixture): + manager = RunManager() + cleanup_finished = asyncio.Event() + cleanup_registered = asyncio.Event() + loop = asyncio.get_running_loop() + release_producer = manager._begin_cancellation_cleanup_producer() + + async def register_cleanup() -> None: + await asyncio.sleep(0) + cleanup = asyncio.create_task(cleanup_finished.wait()) + manager._track_cancellation_cleanup(cleanup, action="delayed stalled cleanup", run_id="run-1") + cleanup_registered.set() + release_producer() + + async def run() -> None: + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + return + + record = await manager.create("thread-1") + record.task = asyncio.create_task(run()) + + def schedule_register() -> None: + asyncio.create_task(register_cleanup()) + + def second_barrier() -> None: + loop.call_soon(schedule_register) + + def first_barrier() -> None: + loop.call_soon(second_barrier) + + def after_run(_completed: asyncio.Future[None]) -> None: + loop.call_soon(first_barrier) + + record.task.add_done_callback(after_run) + with caplog.at_level(logging.WARNING): + shutdown_task = asyncio.create_task(manager.shutdown(timeout=0.05)) + await asyncio.wait_for(cleanup_registered.wait(), timeout=0.2) + assert shutdown_task.done() is False + await shutdown_task + + cleanup = next(iter(manager._cancellation_cleanup_tasks)) + assert cleanup.done() is False + assert any("cancellation cleanup task" in message for message in caplog.messages) + cleanup_finished.set() + await cleanup + await asyncio.sleep(0) + assert not manager._cancellation_cleanup_tasks + + +@pytest.mark.anyio +async def test_shutdown_waits_for_registered_cleanup_producer_through_deep_callback_chain(): + manager = RunManager() + cleanup_finished = asyncio.Event() + cleanup_registered = asyncio.Event() + loop = asyncio.get_running_loop() + release_producer = manager._begin_cancellation_cleanup_producer() + + def callback(depth: int) -> None: + if depth < 100: + loop.call_soon(callback, depth + 1) + return + cleanup = asyncio.create_task(cleanup_finished.wait()) + manager._track_cancellation_cleanup(cleanup, action="deep delayed cleanup", run_id="run-1") + cleanup_registered.set() + release_producer() + + loop.call_soon(callback, 0) + shutdown_task = asyncio.create_task(manager.shutdown(timeout=0.05)) + await asyncio.wait_for(cleanup_registered.wait(), timeout=0.2) + assert shutdown_task.done() is False + cleanup_finished.set() + await shutdown_task + assert not manager._cancellation_cleanup_tasks + + +@pytest.mark.anyio +async def test_shutdown_waits_for_cleanup_producer_that_finishes_without_registration(): + manager = RunManager() + producer_released = asyncio.Event() + loop = asyncio.get_running_loop() + release_producer = manager._begin_cancellation_cleanup_producer() + + def callback(depth: int) -> None: + if depth < 100: + loop.call_soon(callback, depth + 1) + return + release_producer() + producer_released.set() + + loop.call_soon(callback, 0) + shutdown_task = asyncio.create_task(manager.shutdown(timeout=0.2)) + await asyncio.wait_for(producer_released.wait(), timeout=0.2) + await shutdown_task + assert not manager._cancellation_cleanup_tasks + + +@pytest.mark.anyio +async def test_shutdown_cancellation_propagates_while_waiting_for_cleanup_producer(): + manager = RunManager() + release_producer = manager._begin_cancellation_cleanup_producer() + manager._cancellation_cleanup_state_changed.clear() + shutdown_task = asyncio.create_task(manager.shutdown(timeout=1.0)) + await asyncio.sleep(0) + assert shutdown_task.done() is False + + shutdown_task.cancel("shutdown cancelled") + with pytest.raises(asyncio.CancelledError) as caught: + await shutdown_task + assert caught.value.args == ("shutdown cancelled",) + assert manager._cancellation_cleanup_producers + release_producer() + await asyncio.sleep(0) + + +@pytest.mark.anyio +async def test_shutdown_bounds_initial_manager_lock_wait(caplog: pytest.LogCaptureFixture): + manager = RunManager() + await manager._lock.acquire() + started = asyncio.get_running_loop().time() + + try: + with caplog.at_level(logging.WARNING): + await asyncio.wait_for(manager.shutdown(timeout=0.01), timeout=0.2) + finally: + manager._lock.release() + + assert asyncio.get_running_loop().time() - started < 0.2 + assert "could not acquire manager lock" in caplog.text + await asyncio.wait_for(manager._lock.acquire(), timeout=0.1) + manager._lock.release() + await asyncio.sleep(0) + assert not getattr(manager._lock, "_waiters", ()) + + +@pytest.mark.anyio +async def test_lock_waiter_settlement_has_absolute_deadline_and_tracks_late_acquire(): + manager = RunManager() + lock = CancellationResistantLock() + manager._lock = lock + loop = asyncio.get_running_loop() + + waiter_task = asyncio.create_task(manager._acquire_lock_until(loop.time() + 0.02)) + await lock.acquire_started.wait() + started = loop.time() + + result = await asyncio.wait_for(asyncio.shield(waiter_task), timeout=0.2) + + assert result is False + assert loop.time() - started < 0.1 + assert lock.cancel_count == 1 + assert len(manager._lock_waiters) == 1 + + lock.allow_acquire.set() + await asyncio.sleep(0) + await asyncio.sleep(0) + + assert lock.release_count == 1 + assert not manager._lock_waiters + + +@pytest.mark.anyio +async def test_repeated_lock_waiter_cancellation_preserves_original_signal_and_deadline(): + manager = RunManager() + lock = CancellationResistantLock() + manager._lock = lock + loop = asyncio.get_running_loop() + waiter_task = asyncio.create_task(manager._acquire_lock_until(loop.time() + 0.03)) + await lock.acquire_started.wait() + + started = loop.time() + waiter_task.cancel("first cancellation") + await asyncio.sleep(0) + waiter_task.cancel("second cancellation") + + with pytest.raises(asyncio.CancelledError) as caught: + await asyncio.wait_for(waiter_task, timeout=0.2) + + assert caught.value.args == ("first cancellation",) + assert loop.time() - started < 0.1 + assert lock.cancel_count == 1 + assert len(manager._lock_waiters) == 1 + + lock.allow_acquire.set() + await asyncio.sleep(0) + await asyncio.sleep(0) + assert lock.release_count == 1 + assert not manager._lock_waiters + + +@pytest.mark.anyio +@pytest.mark.parametrize("outcome", ["exception", "cancelled"]) +async def test_late_lock_waiter_outcome_is_observed_once(outcome, caplog: pytest.LogCaptureFixture): + manager = RunManager() + lock = CancellationResistantLock(outcome) + manager._lock = lock + loop = asyncio.get_running_loop() + + with caplog.at_level(logging.WARNING): + result = await manager._acquire_lock_until(loop.time() + 0.02) + + assert result is False + assert len(manager._lock_waiters) == 1 + waiter = next(iter(manager._lock_waiters)) + lock.allow_acquire.set() + if outcome == "exception": + with pytest.raises(RuntimeError, match="late lock waiter failure"): + await asyncio.wait_for(asyncio.shield(waiter), timeout=0.2) + else: + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(asyncio.shield(waiter), timeout=0.2) + await asyncio.sleep(0) + + assert not manager._lock_waiters + assert "Manager lock waiter" in caplog.text + + +@pytest.mark.anyio +async def test_lock_waiter_same_tick_timeout_releases_at_most_once(monkeypatch: pytest.MonkeyPatch): + manager = RunManager() + lock = CancellationResistantLock() + manager._lock = lock + + async def settle(_waiter, *, deadline): + _waiter.cancel() + lock.allow_acquire.set() + lock.acquired = True + return True + + monkeypatch.setattr(manager, "_cancel_and_settle_lock_waiter", settle) + result = await manager._acquire_lock_until(asyncio.get_running_loop().time() - 1) + + assert result is False + assert lock.release_count == 0 + + # Exercise the timeout branch with an already-expired deadline while the + # waiter is still pending in a controlled same-tick race. + result = await manager._acquire_lock_until(asyncio.get_running_loop().time() + 0.001) + assert result is False + assert lock.release_count == 1 + assert lock.acquired is False + + +@pytest.mark.anyio +async def test_shutdown_bounds_post_wait_manager_lock_wait(caplog: pytest.LogCaptureFixture): + manager = RunManager() + run_started = asyncio.Event() + lock_held = asyncio.Event() + release_lock = asyncio.Event() + + async def run() -> None: + run_started.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + await manager._lock.acquire() + lock_held.set() + await release_lock.wait() + manager._lock.release() + + record = await manager.create("thread-1") + record.task = asyncio.create_task(run()) + await run_started.wait() + shutdown_task = asyncio.create_task(manager.shutdown(timeout=0.03)) + + try: + await asyncio.wait_for(lock_held.wait(), timeout=0.2) + with caplog.at_level(logging.WARNING): + await asyncio.wait_for(shutdown_task, timeout=0.2) + finally: + release_lock.set() + await record.task + + assert "could not acquire manager lock" in caplog.text + await asyncio.wait_for(manager._lock.acquire(), timeout=0.1) + manager._lock.release() + assert not getattr(manager._lock, "_waiters", ()) + + +@pytest.mark.anyio +async def test_shutdown_observes_completed_run_after_post_wait_lock_timeout(): + store = MemoryRunStore() + manager = RunManager(store=store) + run_started = asyncio.Event() + holder_started = asyncio.Event() + release_holder = asyncio.Event() + holder_task: asyncio.Task[None] | None = None + + async def hold_lock() -> None: + await manager._lock.acquire() + holder_started.set() + await release_holder.wait() + manager._lock.release() + + async def run() -> None: + nonlocal holder_task + run_started.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + holder_task = asyncio.create_task(hold_lock()) + await holder_started.wait() + raise RuntimeError("run failed during shutdown") + + record = await manager.create("thread-1") + record.task = asyncio.create_task(run()) + await run_started.wait() + + try: + await asyncio.wait_for(manager.shutdown(timeout=0.03), timeout=0.2) + assert record.task.done() + assert getattr(record.task, "_log_traceback", True) is False + assert record.status == RunStatus.pending + stored = await store.get(record.run_id) + assert stored is not None + assert stored["status"] == RunStatus.pending.value + finally: + release_holder.set() + if holder_task is not None: + await holder_task + + +@pytest.mark.anyio +async def test_shutdown_cancellation_observes_completed_run_during_post_wait_lock(): + store = MemoryRunStore() + manager = RunManager(store=store) + run_started = asyncio.Event() + holder_started = asyncio.Event() + release_holder = asyncio.Event() + second_lock_wait_started = asyncio.Event() + holder_task: asyncio.Task[None] | None = None + acquire_calls = 0 + original_acquire = manager._acquire_lock_until + + async def observe_second_lock_wait(deadline: float) -> bool: + nonlocal acquire_calls + acquire_calls += 1 + if acquire_calls == 2: + second_lock_wait_started.set() + return await original_acquire(deadline) + + async def hold_lock() -> None: + await manager._lock.acquire() + holder_started.set() + await release_holder.wait() + manager._lock.release() + + async def run() -> None: + nonlocal holder_task + run_started.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + holder_task = asyncio.create_task(hold_lock()) + await holder_started.wait() + raise RuntimeError("run failed during shutdown cancellation") + + manager._acquire_lock_until = observe_second_lock_wait + record = await manager.create("thread-1") + record.task = asyncio.create_task(run()) + await run_started.wait() + shutdown_task = asyncio.create_task(manager.shutdown(timeout=1.0)) + + try: + await asyncio.wait_for(second_lock_wait_started.wait(), timeout=0.2) + shutdown_task.cancel("caller cancelled shutdown") + with pytest.raises(asyncio.CancelledError) as caught: + await shutdown_task + assert caught.value.args == ("caller cancelled shutdown",) + assert record.task.done() + assert getattr(record.task, "_log_traceback", True) is False + assert record.status == RunStatus.pending + stored = await store.get(record.run_id) + assert stored is not None + assert stored["status"] == RunStatus.pending.value + finally: + release_holder.set() + if holder_task is not None: + await holder_task + + +@pytest.mark.anyio +@pytest.mark.parametrize("late_failure", [False, True], ids=["success", "failure"]) +async def test_shutdown_supervises_late_status_persistence_without_cancelling_or_retrying( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + late_failure: bool, +): + manager = RunManager(store=MemoryRunStore()) + run_started = asyncio.Event() + release_persist = asyncio.Event() + persist_started = asyncio.Event() + persist_calls = 0 + + async def run() -> None: + run_started.set() + await asyncio.Event().wait() + + async def stubborn_persist(_record: Any, _status: RunStatus, **_kwargs: Any) -> bool: + nonlocal persist_calls + persist_calls += 1 + persist_started.set() + try: + await release_persist.wait() + except asyncio.CancelledError: + raise AssertionError("shutdown must not cancel status persistence") + if late_failure: + raise RuntimeError("late shutdown persistence failed") + return True + + monkeypatch.setattr(manager, "_persist_status", stubborn_persist) + record = await manager.create("thread-1") + record.task = asyncio.create_task(run()) + await run_started.wait() + + with caplog.at_level(logging.WARNING): + shutdown_task = asyncio.create_task(manager.shutdown(timeout=0.03)) + await asyncio.wait_for(persist_started.wait(), timeout=0.2) + persistence_task = next(iter(manager._shutdown_persistence_tasks)) + await asyncio.wait_for(shutdown_task, timeout=0.15) + + assert persistence_task.done() is False + assert persistence_task in manager._shutdown_persistence_tasks + assert persist_calls == 1 + release_persist.set() + await persistence_task + await asyncio.sleep(0) + + assert not manager._shutdown_persistence_tasks + assert persist_calls == 1 + if late_failure: + assert caplog.text.count("late shutdown persistence failed") == 1 + + +@pytest.mark.anyio +async def test_heartbeat_does_not_schedule_orphans_after_stop_during_renewal( + monkeypatch: pytest.MonkeyPatch, +): + manager = RunManager( + run_ownership_config=RunOwnershipConfig( + lease_seconds=5, + grace_seconds=10, + heartbeat_enabled=True, + ), + ) + renewal_started = asyncio.Event() + renewal_cancelled = asyncio.Event() + release_renewal = asyncio.Event() + renewal_calls = 0 + scheduled_orphans: list[None] = [] + real_asyncio = manager_module.asyncio + + class FastAsyncio: + def __getattr__(self, name: str) -> Any: + return getattr(real_asyncio, name) + + async def wait_for(self, awaitable: Any, timeout: float) -> Any: + if timeout == 1: + awaitable.close() + raise TimeoutError + return await real_asyncio.wait_for(awaitable, timeout) + + async def renew_leases() -> None: + nonlocal renewal_calls + renewal_calls += 1 + if renewal_calls != 3: + return + renewal_started.set() + try: + await release_renewal.wait() + except asyncio.CancelledError: + renewal_cancelled.set() + await release_renewal.wait() + + monkeypatch.setattr(manager_module, "asyncio", FastAsyncio()) + monkeypatch.setattr(manager, "_renew_leases", renew_leases) + monkeypatch.setattr( + manager, + "_schedule_orphan_reconciliation", + lambda: scheduled_orphans.append(None), + ) + + await manager.start_heartbeat() + task = manager._heartbeat_task + assert task is not None + await renewal_started.wait() + + await manager.stop_heartbeat(timeout=0.01) + await asyncio.wait_for(renewal_cancelled.wait(), timeout=0.2) + release_renewal.set() + await task + await asyncio.sleep(0) + + assert renewal_calls == 3 + assert scheduled_orphans == [] + + +@pytest.mark.anyio +async def test_stop_heartbeat_stubborn_task_keeps_one_background_owner_after_deadline( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +): + manager = RunManager( + run_ownership_config=RunOwnershipConfig( + lease_seconds=30, + grace_seconds=10, + heartbeat_enabled=True, + ), + ) + heartbeat_started = asyncio.Event() + heartbeat_cancelled = asyncio.Event() + finish_heartbeat = asyncio.Event() + cancel_calls = 0 + + async def stubborn_heartbeat() -> None: + nonlocal cancel_calls + heartbeat_started.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + cancel_calls += 1 + heartbeat_cancelled.set() + await finish_heartbeat.wait() + raise RuntimeError("late heartbeat failure") + + monkeypatch.setattr(manager, "_heartbeat_loop", stubborn_heartbeat) + await manager.start_heartbeat() + task = manager._heartbeat_task + assert task is not None + await heartbeat_started.wait() + started = asyncio.get_running_loop().time() + + with caplog.at_level(logging.WARNING): + await manager.stop_heartbeat(timeout=0.01) + elapsed = asyncio.get_running_loop().time() - started + assert elapsed < 0.2 + await asyncio.wait_for(heartbeat_cancelled.wait(), timeout=0.2) + assert heartbeat_cancelled.is_set() + assert cancel_calls == 1 + assert manager._heartbeat_task is task + assert "background" in caplog.text + + started = asyncio.get_running_loop().time() + await manager.stop_heartbeat(timeout=0.01) + assert asyncio.get_running_loop().time() - started < 0.05 + assert cancel_calls == 1 + + finish_heartbeat.set() + await asyncio.gather(task, return_exceptions=True) + await asyncio.sleep(0) + assert manager._heartbeat_task is None + assert "late heartbeat failure" in caplog.text + + +@pytest.mark.anyio +async def test_orphan_recovery_stubborn_task_keeps_one_background_owner_after_deadline(caplog: pytest.LogCaptureFixture): + manager = RunManager() + recovery_started = asyncio.Event() + recovery_cancelled = asyncio.Event() + finish_recovery = asyncio.Event() + cancel_calls = 0 + + async def stubborn_recovery() -> None: + nonlocal cancel_calls + recovery_started.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + cancel_calls += 1 + recovery_cancelled.set() + await finish_recovery.wait() + raise RuntimeError("late orphan recovery failure") + + task = asyncio.create_task(stubborn_recovery()) + manager._orphan_recovery_task = task + task.add_done_callback(manager._orphan_reconciliation_done) + await recovery_started.wait() + started = asyncio.get_running_loop().time() + + with caplog.at_level(logging.WARNING): + await manager._drain_orphan_recovery_task(timeout=0.01) + elapsed = asyncio.get_running_loop().time() - started + assert elapsed < 0.2 + await asyncio.wait_for(recovery_cancelled.wait(), timeout=0.2) + assert recovery_cancelled.is_set() + assert cancel_calls == 1 + assert manager._orphan_recovery_task is task + assert "background" in caplog.text + + started = asyncio.get_running_loop().time() + await manager._drain_orphan_recovery_task(timeout=0.01) + assert asyncio.get_running_loop().time() - started < 0.05 + assert cancel_calls == 1 + + finish_recovery.set() + await asyncio.gather(task, return_exceptions=True) + await asyncio.sleep(0) + assert manager._orphan_recovery_task is None + assert "late orphan recovery failure" in caplog.text + + +@pytest.mark.anyio +async def test_stop_heartbeat_consumes_completed_failure_once(caplog: pytest.LogCaptureFixture): + manager = RunManager( + run_ownership_config=RunOwnershipConfig( + lease_seconds=30, + grace_seconds=10, + heartbeat_enabled=True, + ), + ) + + async def failed_heartbeat() -> None: + raise RuntimeError("heartbeat failure") + + manager._heartbeat_loop = failed_heartbeat + await manager.start_heartbeat() + await asyncio.sleep(0) + + with caplog.at_level(logging.WARNING): + await manager.stop_heartbeat(timeout=0) + await asyncio.sleep(0) + + assert caplog.messages.count("Run lease heartbeat failed; its task has stopped") == 1 + assert manager._heartbeat_task is None + + @pytest.mark.anyio async def test_reservation_delete_failure_preserves_body_error_and_clears_local_record(caplog): store = FailingDeleteRunStore() diff --git a/backend/tests/test_runtime_cancellation.py b/backend/tests/test_runtime_cancellation.py new file mode 100644 index 00000000000..929540cfa0e --- /dev/null +++ b/backend/tests/test_runtime_cancellation.py @@ -0,0 +1,91 @@ +import asyncio +from types import SimpleNamespace + +import pytest + +import deerflow.runtime.cancellation as cancellation +from deerflow.runtime.cancellation import wait_for_task_until + + +@pytest.mark.anyio +async def test_wait_for_task_until_reports_completion(): + child = asyncio.create_task(asyncio.sleep(0, result="done")) + + completed = await wait_for_task_until(child, deadline=asyncio.get_running_loop().time() + 1) + + assert completed is True + assert child.result() == "done" + + +@pytest.mark.anyio +async def test_wait_for_task_until_times_out_without_cancelling_child(): + event = asyncio.Event() + child = asyncio.create_task(event.wait()) + + completed = await wait_for_task_until(child, deadline=asyncio.get_running_loop().time() + 0.01) + + assert completed is False + assert child.done() is False + event.set() + await child + + +@pytest.mark.anyio +async def test_wait_for_task_until_zero_budget_returns_immediately(): + event = asyncio.Event() + child = asyncio.create_task(event.wait()) + + completed = await wait_for_task_until(child, deadline=asyncio.get_running_loop().time()) + + assert completed is False + assert child.done() is False + child.cancel() + with pytest.raises(asyncio.CancelledError): + await child + + +@pytest.mark.anyio +async def test_wait_for_task_until_repeated_cancellation_keeps_original_deadline(monkeypatch): + clock = iter((0.0, 0.01, 0.02, 0.05)) + clock_loop = SimpleNamespace(time=lambda: next(clock, 0.05)) + + wait_timeouts = [] + entered_first_wait = asyncio.Event() + entered_second_wait = asyncio.Event() + + async def fake_wait(tasks, *, timeout): + del tasks + wait_timeouts.append(timeout) + if len(wait_timeouts) == 1: + entered_first_wait.set() + await asyncio.Future() + if len(wait_timeouts) == 2: + entered_second_wait.set() + await asyncio.Future() + return set(), set() + + monkeypatch.setattr( + cancellation, + "asyncio", + SimpleNamespace( + CancelledError=asyncio.CancelledError, + get_running_loop=lambda: clock_loop, + wait=fake_wait, + ), + ) + event = asyncio.Event() + child = asyncio.create_task(event.wait()) + waiter = asyncio.create_task(wait_for_task_until(child, deadline=0.05)) + + await entered_first_wait.wait() + waiter.cancel() + await entered_second_wait.wait() + waiter.cancel() + assert waiter.cancelling() == 2 + completed = await waiter + + assert completed is False + assert wait_timeouts == pytest.approx([0.05, 0.04, 0.03]) + assert child.done() is False + event.set() + await child