From 906c3d4554c20e6e72251f160677f4d5e93c8c37 Mon Sep 17 00:00:00 2001 From: spud <92900806+jamespud@users.noreply.github.com> Date: Sun, 20 Sep 2026 19:11:42 +0800 Subject: [PATCH] fix(mcp): make durable task claims cancellation-safe (#4966) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(mcp): re-scope to MCP task claim lifecycle only Keep PR #4966 a small, closed MCP lease/cancellation state-machine change and move RunJournal and Run lifecycle work into dedicated follow-ups. This branch contains only the MCP task claim lifecycle: - mcp task release/snapshot fencing by owner + per-claim lease token - phase-level single-flight poll/cancel/notification owners with retained handoff - routine cancellation no longer persisted as a task failure diagnostic - bounded ordinary release ownership retention past the drain deadline - 0018_mcp_task_lease_tokens migration + migration/bootstrap head assertions - wait_for_task_until helper (MCP uses it); worker-specific capture helper moved to the run-finalization follow-up RunJournal (journal.py + test_run_journal.py) and run lifecycle (manager/worker/store/run sql + run tests) are preserved on backup/cancellation-safety-full and will be raised as separate follow-ups. * fix(mcp): unblock claims after ambiguous handoff resolves A phase-level single-flight owner only guards an ambiguous claim outcome. Once the claim resolves, the phase owner is released immediately; the handoff may continue releasing returned rows as bounded, service-owned background work (transferred to _compensation_tasks on timeout). Per-claim token fencing rejects a late release against a newer claim generation, so a stuck release no longer locks the whole phase until process restart. - README: drop the stale progress-snapshot sentence from the bounded ordinary release description. - service: pop the identity-checked phase owner as soon as the claim outcome is known, then release returned rows with the bounded path; carry the release in _compensation_tasks if it exceeds the drain deadline. - mcp/AGENTS.md: document that only an unresolved claim outcome (not the handoff) blocks later phase scans, and that returned-row releases may continue in the background once the owner is released. - tests: pin that the phase owner is released before a stuck release finishes while the release stays service strong-owned. * refactor(mcp): remove unused single-record claim wrappers _poll_one, _cancel_one, and _notify_one are unreachable in production: the worker always processes claimed records through _run_claimed_batch, so these wrappers preserved a second, dead single-record lifecycle (state is None) whose only observable behavior was a wrapper-specific cancellation release. Remove the three wrappers and migrate the regressions that guarded their cancel/release invariants to exercise the production _run_claimed_batch path (operation=_*_one_claimed, release=_release_*_after_cancellation). The single wrapper-only "state is None" contract (test_poll_release_hang_without_batch) is deleted; all 11 remaining invariants (CancelledError preservation, repeated cancellation, poll-only token-fenced lease release, notification claimed vs dispatched phase release, hung compensation -> service ownership, and background compensation exactly-once observation) are now covered through the real batch lifecycle. * fix(mcp): fence claim-owned mutations against stale generations The per-claim token check in the ORM release/apply paths was only in the SELECT; the final write went out by primary key. On SQLite (where with_for_update() is a no-op) a mutation from an older claim generation could therefore clear a claim that a newer generation had reclaimed after lease expiry — the exact distributed lease-fencing failure the per-claim token was meant to prevent. Make every claim-owned mutation a single atomic conditional UPDATE with the owner and per-claim token in the WHERE clause (rowcount 0 => stale, return False, no mutation): - release_claim: atomic fence; record the poll-failure event after the fence wins (same transaction, holding the write lock). - apply_snapshot / apply_cancel_snapshot: atomic fence; record the event after. - finish_notification_run: atomic fence; use a CASE on event_version >> dispatch_version to keep a newer event pending for redelivery instead of swallowing it as delivered. Add one regression per path: a stale generation's release/apply/finish after a same-worker reclaim is rejected and never clears the newer claim. * test(mcp): pin the migration chain head to the lease-token revision 0026_mcp_task_lease_tokens becomes the alembic head, so the chain-head pin in the 0025 repair test had to move on. Follow the 0023 precedent there (single head plus expected predecessor) instead of pinning a literal head, and give the new revision its own migration test, which owns the pin and covers the nullable claim-token columns on upgrade and their removal on downgrade. * refactor(mcp): close cancellation cleanup leftovers * fix(mcp): retain cancelled release diagnostics * test(mcp): remove obsolete settled compensation case * test(mcp): cover interleaved lease reclaim races --------- Co-authored-by: Willem Jiang --- README.md | 2 +- backend/app/mcp_tasks/service.py | 1042 +++++- .../packages/harness/deerflow/mcp/AGENTS.md | 2 +- .../deerflow/persistence/mcp_tasks/model.py | 2 + .../deerflow/persistence/mcp_tasks/sql.py | 286 +- .../deerflow/persistence/migrations/AGENTS.md | 9 +- .../versions/0026_mcp_task_lease_tokens.py | 31 + .../harness/deerflow/runtime/cancellation.py | 24 + backend/tests/AGENTS.md | 10 + backend/tests/test_agent_guidance_check.py | 12 + backend/tests/test_mcp_task_repository.py | 623 +++- backend/tests/test_mcp_task_service.py | 2807 ++++++++++++++++- ...st_migration_0025_repair_run_change_seq.py | 11 +- ...st_migration_0026_mcp_task_lease_tokens.py | 51 + backend/tests/test_runtime_cancellation.py | 91 + 15 files changed, 4768 insertions(+), 235 deletions(-) create mode 100644 backend/packages/harness/deerflow/persistence/migrations/versions/0026_mcp_task_lease_tokens.py create mode 100644 backend/packages/harness/deerflow/runtime/cancellation.py create mode 100644 backend/tests/test_migration_0026_mcp_task_lease_tokens.py create mode 100644 backend/tests/test_runtime_cancellation.py diff --git a/README.md b/README.md index 2a29f965f..769f91a6a 100644 --- a/README.md +++ b/README.md @@ -609,7 +609,7 @@ OpenViking memory backend; it does not replace automatic turn capture or recall. The Gateway can adapt an MCP server's ordinary `submit` / `status` / `cancel` tools into durable background tasks. The Agent sees only the configured submit tool and a DeerFlow-local task ID; remote IDs are persisted before the submit call returns, while status and cancel stay internal to the runtime. Polling uses cross-worker leases, exponential retry backoff, scoped MCP sessions, bounded result storage, and restart recovery. A status-tool `isError` is retained as a bounded diagnostic and retried; servers report a permanent remote-task outcome through a normal structured result with `status: "failed"`. Remote poll hints are finite positive numbers capped at 24 hours, artifact-reference JSON is limited to 64 KiB, and task/server identifiers are validated against their durable SQL column limits before persistence. Input-required and terminal updates wake the current chat through idempotent Agent runs, while `list_background_tasks` and `cancel_background_task` let the Agent manage tasks without asking users for remote handles. Current-thread tasks are available through `GET /api/threads/{thread_id}/mcp-tasks`, its detail endpoint, and `POST /api/threads/{thread_id}/mcp-tasks/{task_id}/cancel`; when the task runtime actually starts, the Web UI exposes the same safe local view from the chat header with live status refresh, cancellation, and on-demand result, artifact, input-request, status-error, and cancellation-retry details. Default-disabled and memory-backend deployments hide that UI and do not poll the task endpoints. A failed remote cancellation remains queued with backoff, and its latest bounded error and attempt count stay visible in the expanded task card. Enable `mcp_tasks` in `config.yaml`, configure `task_toolsets` with exact raw tool names in `extensions_config.json`, and use a SQL database backend (`sqlite` or `postgres`). Task-enabled server connection, authentication, interceptor, timeout, or binding changes require a Gateway restart so Agent tool discovery and background calls cannot use different configuration versions. `input_required` is notification-only for now: DeerFlow can display the request but cannot yet submit the user's answer back to the remote task. -Notification launch and failed Agent-run deliveries use capped exponential backoff with a visible attempt count and stop after five failed attempts. A permanently rejected target such as a deleted chat is dead-lettered immediately instead of retried forever or recreated. Cancellation endpoints return after durably recording the request; the background service owns the potentially slow remote MCP call and its retry schedule. +Notification launch and failed Agent-run deliveries use capped exponential backoff with a visible attempt count and stop after five failed attempts. When a bounded ordinary release exceeds its drain deadline, the service retains ownership until it settles. A permanently rejected target such as a deleted chat is dead-lettered immediately instead of retried forever or recreated. Cancellation endpoints return after durably recording the request; the background service owns the potentially slow remote MCP call and its retry schedule. Notification runs keep their trusted delivery instruction separate from the framed, untrusted remote event payload. The process-started task runtime—not a hot config read—controls whether the task-management tools are exposed, so changing `mcp_tasks` requires a Gateway restart. When a skill's `allowed-tools` policy is active, `list_background_tasks` and `cancel_background_task` must be declared explicitly like other business tools. See the [MCP Server Guide](backend/docs/MCP_SERVER.md) for detailed instructions. diff --git a/backend/app/mcp_tasks/service.py b/backend/app/mcp_tasks/service.py index 6a318cb92..16a414950 100644 --- a/backend/app/mcp_tasks/service.py +++ b/backend/app/mcp_tasks/service.py @@ -6,7 +6,8 @@ import logging 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 @@ from deerflow.mcp.tasks import ( 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,33 @@ logger = logging.getLogger(__name__) _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 + + +@dataclass(slots=True) +class _ClaimOwner: + claim_task: asyncio.Future[list[dict[str, Any]]] + handoff_task: asyncio.Task[None] | None = None + + +_current_batch_record: ContextVar[_BatchRecordState | None] = ContextVar( + "mcp_task_current_batch_record", + default=None, +) def _bound_error(error: str | None) -> str | None: @@ -42,6 +70,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 +118,11 @@ class McpTaskService: 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._claim_owners: dict[str, _ClaimOwner] = {} self._stop = asyncio.Event() @property @@ -86,6 +133,160 @@ class McpTaskService: 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: + def finalize(completed: asyncio.Future[Any]) -> None: + self._observe_batch_release_task( + state, + completed, + ordinary=ordinary, + action=action, + ) + + task.add_done_callback(finalize) + + def _retain_batch_release_task(self, task: asyncio.Future[Any]) -> None: + """Transfer a timed-out ordinary release to service ownership.""" + if task in self._compensation_tasks: + return + self._compensation_tasks.add(task) + task.add_done_callback(self._compensation_tasks.discard) + + 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: + release_task = asyncio.ensure_future(release()) + try: + await asyncio.wait_for( + asyncio.shield(release_task), + timeout=_CANCELLATION_DRAIN_TIMEOUT_SECONDS, + ) + except asyncio.CancelledError: + self._track_compensation_task( + release_task, + action=action, + task_id=str(record.get("id") or "unknown"), + ) + raise + except TimeoutError: + self._track_compensation_task( + release_task, + action=action, + task_id=str(record.get("id") or "unknown"), + ) + logger.warning( + "Timed out after %.1f seconds waiting for MCP task release; it continues in the background (%s, task_id=%s)", + _CANCELLATION_DRAIN_TIMEOUT_SECONDS, + action, + record.get("id"), + ) + 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.wait_for( + asyncio.shield(task), + timeout=_CANCELLATION_DRAIN_TIMEOUT_SECONDS, + ) + except TimeoutError: + # The release is still in flight past the drain deadline on the + # uncancelled path; it stays tracked by the batch state and settles + # in the background instead of blocking the poller. + self._retain_batch_release_task(task) + self._observe_batch_release_task(state, task, ordinary=True, action=action) + logger.warning( + "Timed out after %.1f seconds waiting for MCP task release; it continues in the background (%s, task_id=%s)", + _CANCELLATION_DRAIN_TIMEOUT_SECONDS, + action, + record.get("id"), + ) + return + except asyncio.CancelledError: + self._observe_batch_release_task(state, task, ordinary=True, action=action) + if state.ordinary_release_terminal and not asyncio.current_task().cancelling(): + # The release cancelled itself and the caller is not cancelling: + # consume it once and return without re-raising. + 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 +375,9 @@ class McpTaskService: ) 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,48 +391,299 @@ class McpTaskService: 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( + lambda: self._repository.claim_due_tasks( + now=now, + lease_owner=self._lease_owner, + lease_seconds=self._lease_seconds, + limit=self._max_concurrent_polls, + ), + phase="poll", + 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): - if isinstance(result, BaseException): - logger.error( - "Unexpected MCP task poll failure (task_id=%s); the lease will expire for recovery", - record.get("id"), - exc_info=(type(result), result, result.__traceback__), - ) + for state, result in zip(states, results, strict=True): + if not isinstance(result, BaseException) or isinstance(result, asyncio.CancelledError): + continue + logger.error( + "Unexpected MCP task poll failure (task_id=%s); the lease will expire for recovery", + state.record.get("id"), + exc_info=(type(result), result, result.__traceback__), + ) await self._run_notifications(now=datetime.now(UTC)) @@ -299,26 +748,34 @@ class McpTaskService: 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( + lambda: claim( + now=now, + lease_owner=self._lease_owner, + lease_seconds=self._lease_seconds, + limit=self._max_concurrent_polls, + ), + phase="cancel", + 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): - if isinstance(result, BaseException): - logger.error( - "Unexpected MCP task cancellation failure (task_id=%s); the lease will expire for recovery", - record.get("id"), - exc_info=(type(result), result, result.__traceback__), - ) + for state, result in zip(states, results, strict=True): + if not isinstance(result, BaseException) or isinstance(result, asyncio.CancelledError): + continue + logger.error( + "Unexpected MCP task cancellation failure (task_id=%s); the lease will expire for recovery", + state.record.get("id"), + exc_info=(type(result), result, result.__traceback__), + ) - async def _cancel_one(self, record: dict[str, Any]) -> None: + 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: @@ -330,6 +787,7 @@ class McpTaskService: await self._repository.apply_cancel_snapshot( record["id"], lease_owner=self._lease_owner, + lease_token=record["lease_token"], status=snapshot.status.value, result=snapshot.result, result_preview=snapshot.result_preview, @@ -343,52 +801,54 @@ class McpTaskService: 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, + lease_token=record["lease_token"], + 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( + lambda: 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, + ), + phase="notification", + 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: + 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)) @@ -397,6 +857,7 @@ class McpTaskService: await self._repository.dead_letter_notification( task_id, lease_owner=self._lease_owner, + notification_lease_token=record["notification_lease_token"], dispatch_version=dispatch_version, error=_bound_error(f"Notification delivery stopped after {notification_attempts} failed attempts: {previous_error}"), count_failure=False, @@ -412,6 +873,7 @@ class McpTaskService: await self._repository.finish_notification_run( task_id, lease_owner=self._lease_owner, + notification_lease_token=record["notification_lease_token"], dispatch_version=dispatch_version, delivered=False, next_notification_at=now + timedelta(seconds=self._notification_retry_seconds(record)), @@ -422,6 +884,7 @@ class McpTaskService: await self._repository.finish_notification_run( task_id, lease_owner=self._lease_owner, + notification_lease_token=record["notification_lease_token"], dispatch_version=dispatch_version, delivered=True, next_notification_at=None, @@ -432,6 +895,7 @@ class McpTaskService: await self._repository.finish_notification_run( task_id, lease_owner=self._lease_owner, + notification_lease_token=record["notification_lease_token"], dispatch_version=dispatch_version, delivered=False, next_notification_at=now + timedelta(seconds=self._notification_retry_seconds(record)), @@ -442,6 +906,7 @@ class McpTaskService: await self._repository.defer_dispatched_notification( task_id, lease_owner=self._lease_owner, + notification_lease_token=record["notification_lease_token"], dispatch_version=dispatch_version, next_notification_at=now + timedelta(seconds=self._poll_interval_seconds), now=now, @@ -463,6 +928,7 @@ class McpTaskService: await self._repository.dead_letter_notification( task_id, lease_owner=self._lease_owner, + notification_lease_token=record["notification_lease_token"], dispatch_version=dispatch_version, error=_bound_error(str(exc) or type(exc).__name__), count_failure=True, @@ -470,27 +936,40 @@ class McpTaskService: ) 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, + notification_lease_token=record["notification_lease_token"], + 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, + notification_lease_token=record["notification_lease_token"], + 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( task_id, lease_owner=self._lease_owner, + notification_lease_token=record["notification_lease_token"], dispatch_version=dispatch_version, run_id=result["run_id"], now=now, @@ -503,14 +982,307 @@ class McpTaskService: self._max_poll_backoff_seconds, ) - async def _poll_one(self, record: dict, *, now: datetime) -> None: + async def _claim_with_cancellation_release( + self, + claim: Callable[[], Awaitable[list[dict[str, Any]]]], + *, + phase: str, + action: str, + release: Callable[[dict[str, Any]], Awaitable[None]], + ) -> list[dict[str, Any]]: + if phase in self._claim_owners: + logger.warning( + "Skipping MCP %s claim because the previous claim/handoff is still unresolved", + phase, + ) + return [] + + claim_task = asyncio.ensure_future(claim()) + owner = _ClaimOwner(claim_task=claim_task) + self._claim_owners[phase] = owner + try: + records = await asyncio.wait_for( + asyncio.shield(claim_task), + timeout=_CANCELLATION_DRAIN_TIMEOUT_SECONDS, + ) + except TimeoutError: + self._start_claim_owner_handoff( + owner, + phase=phase, + action=action, + release=release, + ) + logger.warning( + "Timed out after %.1f seconds waiting for MCP task claim; it continues in the background (%s, task_id=batch)", + _CANCELLATION_DRAIN_TIMEOUT_SECONDS, + action, + ) + return [] + 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) + if self._claim_owners.get(phase) is owner: + self._claim_owners.pop(phase, None) + return [] + handoff = self._start_claim_owner_handoff( + owner, + phase=phase, + action=action, + release=release, + ) + 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 + except Exception: + if self._claim_owners.get(phase) is owner: + self._claim_owners.pop(phase, None) + raise + else: + if self._claim_owners.get(phase) is owner: + self._claim_owners.pop(phase, None) + return records + + def _start_claim_owner_handoff( + self, + owner: _ClaimOwner, + *, + phase: str, + action: str, + release: Callable[[dict[str, Any]], Awaitable[None]], + ) -> asyncio.Task[None]: + if owner.handoff_task is None: + owner.handoff_task = asyncio.create_task( + self._finish_cancelled_claim_handoff( + owner.claim_task, + owner=owner, + phase=phase, + action=action, + release=release, + ), + name=f"mcp-{action.replace(' ', '-')}-handoff", + ) + self._track_compensation_task(owner.handoff_task, action=f"finish {action} handoff", task_id="batch") + return owner.handoff_task + + async def _finish_cancelled_claim_handoff( + self, + claim_task: asyncio.Future[list[dict[str, Any]]], + *, + owner: _ClaimOwner, + phase: str, + action: str, + release: Callable[[dict[str, Any]], Awaitable[None]], + ) -> None: + try: + 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 + # The claim's durable outcome is now known. Release the phase owner + # immediately: per-claim token fencing already rejects a stale release + # against a newer claim, so the phase no longer needs this owner to + # guard the ambiguous claim. Returned rows are released as bounded, + # service-owned background work, so a stuck release cannot lock the + # whole phase until process restart. + if self._claim_owners.get(phase) is owner: + self._claim_owners.pop(phase, None) + if records: + await self._release_claimed_records(records, release=release) + finally: + if self._claim_owners.get(phase) is owner: + self._claim_owners.pop(phase, None) + + 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); the lease will expire for recovery", + 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, + notification_lease_token=record["notification_lease_token"], + 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.wait_for( + asyncio.shield(task), + timeout=_CANCELLATION_DRAIN_TIMEOUT_SECONDS, + ) + except TimeoutError: + self._track_compensation_task( + task, + action="release notification failure", + task_id=record["id"], + ) + logger.warning( + "Timed out after %.1f seconds waiting for MCP task notification release; it continues in the background (task_id=%s)", + _CANCELLATION_DRAIN_TIMEOUT_SECONDS, + record["id"], + ) + return + 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, + lease_token=record["lease_token"], + next_cancel_at=datetime.now(UTC), + error=record.get("last_cancel_error"), + ), + 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, + notification_lease_token=record["notification_lease_token"], + next_notification_at=datetime.now(UTC), + error=record.get("notification_error"), + count_failure=False, + ) + action = "release dispatched notification lease" + else: + compensation = self._repository.release_notification_claim( + task_id, + lease_owner=self._lease_owner, + notification_lease_token=record["notification_lease_token"], + next_notification_at=datetime.now(UTC), + error=record.get("notification_error"), + 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, + lease_token=record["lease_token"], + ), + 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 +1309,16 @@ class McpTaskService: 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) @@ -553,6 +1334,7 @@ class McpTaskService: applied = await self._repository.apply_snapshot( record["id"], lease_owner=self._lease_owner, + lease_token=record["lease_token"], status=snapshot.status.value, result=snapshot.result, result_preview=snapshot.result_preview, @@ -589,6 +1371,7 @@ class McpTaskService: await self._repository.release_claim( record["id"], lease_owner=self._lease_owner, + lease_token=record["lease_token"], next_poll_at=now + timedelta(seconds=retry_seconds), error=bounded_error, tracking_degraded_after_errors=self._tracking_degraded_after_errors, @@ -644,20 +1427,61 @@ class McpTaskService: 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 2c32e7464..ff2f3c28e 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. Every claim carries a fresh per-claim token, and every poll/cancel/notification release or snapshot apply is fenced by the owner plus that token, so a release that completes after the same service reclaims the row cannot clear the newer lease. Poll, cancel, and notification claims also have separate phase-level single-flight owners that protect only an ambiguous claim outcome: after an uncancelled timeout or caller cancellation, later scans do not start another claim while the outcome is unknown. Once the claim resolves, the phase owner is released immediately, and any returned rows are released through bounded, service-owned background work; per-claim token fencing prevents a late release from mutating a newer claim generation, so a stuck release does not stall the whole phase. While the claim outcome remains unresolved, later scans skip the phase and emit a warning rather than overlapping an ambiguous database claim. Once the claim resolves, the handoff may continue releasing returned rows in the background without blocking later scans. Lease expiry remains the cross-process crash fallback rather than the normal recovery for a live service's late claim. 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. Routine cancellation releases clear only the lease/token and preserve any pre-existing `last_cancel_error` or `notification_error`; cancellation itself is not persisted as a task failure diagnostic. Cancelling an in-flight poll releases only the owner- and per-claim-token-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/model.py b/backend/packages/harness/deerflow/persistence/mcp_tasks/model.py index 73c851e9b..3746817ad 100644 --- a/backend/packages/harness/deerflow/persistence/mcp_tasks/model.py +++ b/backend/packages/harness/deerflow/persistence/mcp_tasks/model.py @@ -48,6 +48,7 @@ class McpTaskRow(Base): next_notification_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True) notification_lease_owner: Mapped[str | None] = mapped_column(String(128), nullable=True) notification_lease_expires_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True) + notification_lease_token: Mapped[str | None] = mapped_column(String(64), nullable=True) next_poll_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), index=True, nullable=True) last_polled_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True) last_poll_error: Mapped[str | None] = mapped_column(Text, nullable=True) @@ -55,6 +56,7 @@ class McpTaskRow(Base): consecutive_poll_error_count: Mapped[int] = mapped_column(Integer, default=0) lease_owner: Mapped[str | None] = mapped_column(String(128), nullable=True) lease_expires_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True) + lease_token: Mapped[str | None] = mapped_column(String(64), nullable=True) cancel_requested_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True) cancel_attempt_count: Mapped[int] = mapped_column(Integer, default=0, server_default="0") next_cancel_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True) diff --git a/backend/packages/harness/deerflow/persistence/mcp_tasks/sql.py b/backend/packages/harness/deerflow/persistence/mcp_tasks/sql.py index 8766f57ac..604fc3178 100644 --- a/backend/packages/harness/deerflow/persistence/mcp_tasks/sql.py +++ b/backend/packages/harness/deerflow/persistence/mcp_tasks/sql.py @@ -2,10 +2,11 @@ from __future__ import annotations import hashlib import json +import uuid from datetime import UTC, datetime, timedelta from typing import Any -from sqlalchemy import or_, select, update +from sqlalchemy import case, or_, select, update from sqlalchemy.exc import IntegrityError from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker @@ -33,6 +34,11 @@ _TIMESTAMP_FIELDS = ( _INFLIGHT_NOTIFICATION_STATUSES = frozenset({"claimed", "dispatched", "retry"}) +def _new_claim_token() -> str: + """Return a fresh per-claim token used to fence releases against reclaims.""" + return uuid.uuid4().hex + + def _notification_event(row: McpTaskRow, *, tracking_degraded: bool) -> dict[str, Any] | None: if row.status not in _ATTENTION_STATUS_VALUES and not tracking_degraded: return None @@ -235,6 +241,7 @@ class McpTaskRepository: for row in rows: row.lease_owner = lease_owner row.lease_expires_at = lease_expires_at + row.lease_token = _new_claim_token() row.poll_attempt_count += 1 row.updated_at = now await session.commit() @@ -245,6 +252,7 @@ class McpTaskRepository: task_id: str, *, lease_owner: str, + lease_token: str, status: str, result: Any | None, result_preview: str | None, @@ -256,37 +264,44 @@ class McpTaskRepository: polled_at: datetime, ) -> bool: async with self._sf() as session: - stmt = ( - select(McpTaskRow) + # Atomic fence: a poll result from an older generation must not + # overwrite a claim a newer generation reclaimed after lease expiry. + values: dict[str, Any] = { + "status": status, + "result": result, + "result_preview": result_preview, + "result_truncated": result_truncated, + "result_artifact": result_artifact, + "error": error, + "input_required": input_required, + "next_poll_at": next_poll_at, + "last_polled_at": polled_at, + "last_poll_error": None, + "consecutive_poll_error_count": 0, + "lease_owner": None, + "lease_expires_at": None, + "lease_token": None, + "updated_at": polled_at, + } + if status in _TERMINAL_STATUS_VALUES: + values["completed_at"] = polled_at + update_result = await session.execute( + update(McpTaskRow) .where( McpTaskRow.id == task_id, McpTaskRow.lease_owner == lease_owner, + McpTaskRow.lease_token == lease_token, McpTaskRow.lease_expires_at >= polled_at, McpTaskRow.status.not_in(_TERMINAL_STATUS_VALUES), McpTaskRow.cancel_requested_at.is_(None), ) - .with_for_update() + .values(**values) ) - row = (await session.execute(stmt)).scalar_one_or_none() - if row is None: + if not update_result.rowcount: return False - row.status = status - row.result = result - row.result_preview = result_preview - row.result_truncated = result_truncated - row.result_artifact = result_artifact - row.error = error - row.input_required = input_required - row.next_poll_at = next_poll_at - row.last_polled_at = polled_at - row.last_poll_error = None - row.consecutive_poll_error_count = 0 - row.lease_owner = None - row.lease_expires_at = None - row.updated_at = polled_at - if status in _TERMINAL_STATUS_VALUES: - row.completed_at = polled_at - _record_event_if_changed(row, tracking_degraded=False, now=polled_at) + row = (await session.execute(select(McpTaskRow).where(McpTaskRow.id == task_id))).scalar_one_or_none() + if row is not None: + _record_event_if_changed(row, tracking_degraded=False, now=polled_at) await session.commit() return True @@ -295,30 +310,77 @@ class McpTaskRepository: task_id: str, *, lease_owner: str, + lease_token: str, next_poll_at: datetime, error: str, tracking_degraded_after_errors: int = 3, ) -> bool: async with self._sf() as session: - stmt = select(McpTaskRow).where(McpTaskRow.id == task_id, McpTaskRow.lease_owner == lease_owner).with_for_update() - row = (await session.execute(stmt)).scalar_one_or_none() - if row is None: - return False now = datetime.now(UTC) - row.next_poll_at = next_poll_at - row.last_poll_error = error - row.consecutive_poll_error_count = int(row.consecutive_poll_error_count or 0) + 1 - row.lease_owner = None - row.lease_expires_at = None - row.updated_at = now - _record_event_if_changed( - row, - tracking_degraded=row.consecutive_poll_error_count >= tracking_degraded_after_errors, - now=now, + # Atomic fence: only clear the claim if the owner AND per-claim token + # still match. ``with_for_update()`` is a no-op on SQLite, so the old + # select-then-write could clear a claim that a newer generation had + # reclaimed after lease expiry. A conditional UPDATE makes the fence + # atomic: a stale release (rowcount 0) is a no-op and never mutates a + # newer claim. + update_result = await session.execute( + update(McpTaskRow) + .where( + McpTaskRow.id == task_id, + McpTaskRow.lease_owner == lease_owner, + McpTaskRow.lease_token == lease_token, + ) + .values( + next_poll_at=next_poll_at, + last_poll_error=error, + consecutive_poll_error_count=McpTaskRow.consecutive_poll_error_count + 1, + lease_owner=None, + lease_expires_at=None, + lease_token=None, + updated_at=now, + ) ) + if not update_result.rowcount: + return False + # The fence won and we hold the write lock, so read the released row + # consistently and record the poll-failure tracking event. + row = (await session.execute(select(McpTaskRow).where(McpTaskRow.id == task_id))).scalar_one_or_none() + if row is not None: + _record_event_if_changed( + row, + tracking_degraded=int(row.consecutive_poll_error_count or 0) >= tracking_degraded_after_errors, + now=now, + ) await session.commit() return True + async def release_poll_claim_after_cancellation( + self, + task_id: str, + *, + lease_owner: str, + lease_token: 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, + McpTaskRow.lease_token == lease_token, + ) + .values( + lease_owner=None, + lease_expires_at=None, + lease_token=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, @@ -351,6 +413,7 @@ class McpTaskRepository: # lease so it cannot trigger a concurrent remote cancellation. row.lease_owner = None row.lease_expires_at = None + row.lease_token = None row.updated_at = requested_at await session.commit() return self._row_to_dict(row) @@ -380,6 +443,7 @@ class McpTaskRepository: for row in rows: row.lease_owner = lease_owner row.lease_expires_at = expires_at + row.lease_token = _new_claim_token() row.cancel_attempt_count = int(row.cancel_attempt_count or 0) + 1 row.updated_at = now await session.commit() @@ -390,6 +454,7 @@ class McpTaskRepository: task_id: str, *, lease_owner: str, + lease_token: str, status: str, result: Any | None, result_preview: str | None, @@ -402,34 +467,38 @@ class McpTaskRepository: if status not in _TERMINAL_STATUS_VALUES: raise ValueError("A cancellation response must report a terminal task status") async with self._sf() as session: - stmt = ( - select(McpTaskRow) + update_result = await session.execute( + update(McpTaskRow) .where( McpTaskRow.id == task_id, McpTaskRow.lease_owner == lease_owner, + McpTaskRow.lease_token == lease_token, McpTaskRow.lease_expires_at >= completed_at, McpTaskRow.status.not_in(_TERMINAL_STATUS_VALUES), ) - .with_for_update() + .values( + status=status, + result=result, + result_preview=result_preview, + result_truncated=result_truncated, + result_artifact=result_artifact, + error=error, + input_required=input_required, + next_poll_at=None, + next_cancel_at=None, + last_cancel_error=None, + lease_owner=None, + lease_expires_at=None, + lease_token=None, + completed_at=completed_at, + updated_at=completed_at, + ) ) - row = (await session.execute(stmt)).scalar_one_or_none() - if row is None: + if not update_result.rowcount: return False - row.status = status - row.result = result - row.result_preview = result_preview - row.result_truncated = result_truncated - row.result_artifact = result_artifact - row.error = error - row.input_required = input_required - row.next_poll_at = None - row.next_cancel_at = None - row.last_cancel_error = None - row.lease_owner = None - row.lease_expires_at = None - row.completed_at = completed_at - row.updated_at = completed_at - _record_event_if_changed(row, tracking_degraded=False, now=completed_at) + row = (await session.execute(select(McpTaskRow).where(McpTaskRow.id == task_id))).scalar_one_or_none() + if row is not None: + _record_event_if_changed(row, tracking_degraded=False, now=completed_at) await session.commit() return True @@ -438,17 +507,23 @@ class McpTaskRepository: task_id: str, *, lease_owner: str, + lease_token: str, next_cancel_at: datetime, - error: str, + error: str | None, ) -> bool: stmt = ( update(McpTaskRow) - .where(McpTaskRow.id == task_id, McpTaskRow.lease_owner == lease_owner) + .where( + McpTaskRow.id == task_id, + McpTaskRow.lease_owner == lease_owner, + McpTaskRow.lease_token == lease_token, + ) .values( next_cancel_at=next_cancel_at, last_cancel_error=error, lease_owner=None, lease_expires_at=None, + lease_token=None, updated_at=datetime.now(UTC), ) ) @@ -485,6 +560,7 @@ class McpTaskRepository: for row in rows: row.notification_lease_owner = lease_owner row.notification_lease_expires_at = expires_at + row.notification_lease_token = _new_claim_token() rebuild_snapshot = row.notification_status in ("pending", "claimed") or (row.notification_status == "retry" and row.dispatch_version != row.event_version) if rebuild_snapshot: if row.dispatch_version != row.event_version: @@ -506,6 +582,7 @@ class McpTaskRepository: task_id: str, *, lease_owner: str, + notification_lease_token: str, dispatch_version: int, run_id: str, now: datetime, @@ -515,6 +592,7 @@ class McpTaskRepository: .where( McpTaskRow.id == task_id, McpTaskRow.notification_lease_owner == lease_owner, + McpTaskRow.notification_lease_token == notification_lease_token, McpTaskRow.notification_lease_expires_at >= now, McpTaskRow.dispatch_version == dispatch_version, McpTaskRow.notification_status.in_(("claimed", "retry")), @@ -526,6 +604,7 @@ class McpTaskRepository: next_notification_at=now, notification_lease_owner=None, notification_lease_expires_at=None, + notification_lease_token=None, updated_at=now, ) ) @@ -539,8 +618,9 @@ class McpTaskRepository: task_id: str, *, lease_owner: str, + notification_lease_token: str, next_notification_at: datetime, - error: str, + error: str | None, replace_with_latest: bool, count_failure: bool = False, ) -> bool: @@ -550,6 +630,7 @@ class McpTaskRepository: "next_notification_at": next_notification_at, "notification_lease_owner": None, "notification_lease_expires_at": None, + "notification_lease_token": None, "updated_at": datetime.now(UTC), } if replace_with_latest: @@ -559,7 +640,15 @@ class McpTaskRepository: ) if count_failure: values["notification_attempt_count"] = McpTaskRow.notification_attempt_count + 1 - stmt = update(McpTaskRow).where(McpTaskRow.id == task_id, McpTaskRow.notification_lease_owner == lease_owner).values(**values) + stmt = ( + update(McpTaskRow) + .where( + McpTaskRow.id == task_id, + McpTaskRow.notification_lease_owner == lease_owner, + McpTaskRow.notification_lease_token == notification_lease_token, + ) + .values(**values) + ) async with self._sf() as session: result = await session.execute(stmt) await session.commit() @@ -570,57 +659,67 @@ class McpTaskRepository: task_id: str, *, lease_owner: str, + notification_lease_token: str, dispatch_version: int, delivered: bool, next_notification_at: datetime | None, error: str | None, now: datetime, ) -> bool: + if delivered: + # A newer event may have arrived after this dispatch was queued; keep + # it pending for redelivery instead of swallowing it as delivered. + newer = McpTaskRow.event_version > dispatch_version + values: dict[str, Any] = { + "notified_version": dispatch_version, + "notification_status": case((newer, "pending"), else_="delivered"), + "dispatch_version": None, + "dispatch_attempt": 0, + "dispatch_event": None, + "notification_run_id": None, + "notification_error": None, + "notification_attempt_count": 0, + "next_notification_at": case((newer, now), else_=None), + } + else: + values = { + "notification_status": "retry", + "dispatch_attempt": McpTaskRow.dispatch_attempt + 1, + "notification_attempt_count": McpTaskRow.notification_attempt_count + 1, + "notification_run_id": None, + "notification_error": error, + "next_notification_at": next_notification_at, + } + values.update( + notification_lease_owner=None, + notification_lease_expires_at=None, + notification_lease_token=None, + updated_at=now, + ) async with self._sf() as session: - stmt = ( - select(McpTaskRow) + result = await session.execute( + update(McpTaskRow) .where( McpTaskRow.id == task_id, McpTaskRow.notification_lease_owner == lease_owner, + McpTaskRow.notification_lease_token == notification_lease_token, McpTaskRow.notification_lease_expires_at >= now, McpTaskRow.dispatch_version == dispatch_version, McpTaskRow.notification_status == "dispatched", ) - .with_for_update() + .values(**values) ) - row = (await session.execute(stmt)).scalar_one_or_none() - if row is None: - return False - if delivered: - row.notified_version = dispatch_version - row.notification_status = "pending" if row.event_version > dispatch_version else "delivered" - row.dispatch_version = None - row.dispatch_attempt = 0 - row.dispatch_event = None - row.notification_run_id = None - row.notification_error = None - row.notification_attempt_count = 0 - row.next_notification_at = now if row.event_version > dispatch_version else None - else: - row.notification_status = "retry" - row.dispatch_attempt = int(row.dispatch_attempt or 0) + 1 - row.notification_attempt_count = int(row.notification_attempt_count or 0) + 1 - row.notification_run_id = None - row.notification_error = error - row.next_notification_at = next_notification_at - row.notification_lease_owner = None - row.notification_lease_expires_at = None - row.updated_at = now await session.commit() - return True + return bool(result.rowcount) async def release_notification_lease( self, task_id: str, *, lease_owner: str, + notification_lease_token: str, next_notification_at: datetime, - error: str, + error: str | None, count_failure: bool = False, ) -> bool: """Release unexpected notification work without changing its phase.""" @@ -629,6 +728,7 @@ class McpTaskRepository: "next_notification_at": next_notification_at, "notification_lease_owner": None, "notification_lease_expires_at": None, + "notification_lease_token": None, "updated_at": datetime.now(UTC), } if count_failure: @@ -638,6 +738,7 @@ class McpTaskRepository: .where( McpTaskRow.id == task_id, McpTaskRow.notification_lease_owner == lease_owner, + McpTaskRow.notification_lease_token == notification_lease_token, ) .values(**values) ) @@ -651,6 +752,7 @@ class McpTaskRepository: task_id: str, *, lease_owner: str, + notification_lease_token: str, dispatch_version: int, error: str, count_failure: bool, @@ -660,6 +762,7 @@ class McpTaskRepository: base_filters = ( McpTaskRow.id == task_id, McpTaskRow.notification_lease_owner == lease_owner, + McpTaskRow.notification_lease_token == notification_lease_token, McpTaskRow.notification_lease_expires_at >= now, McpTaskRow.dispatch_version == dispatch_version, McpTaskRow.notification_status.in_(("claimed", "retry", "dispatched")), @@ -670,6 +773,7 @@ class McpTaskRepository: "next_notification_at": None, "notification_lease_owner": None, "notification_lease_expires_at": None, + "notification_lease_token": None, "dispatch_version": None, "dispatch_attempt": 0, "dispatch_event": None, @@ -695,6 +799,7 @@ class McpTaskRepository: next_notification_at=now, notification_lease_owner=None, notification_lease_expires_at=None, + notification_lease_token=None, dispatch_version=None, dispatch_attempt=0, dispatch_event=None, @@ -710,6 +815,7 @@ class McpTaskRepository: task_id: str, *, lease_owner: str, + notification_lease_token: str, dispatch_version: int, next_notification_at: datetime, now: datetime, @@ -720,6 +826,7 @@ class McpTaskRepository: .where( McpTaskRow.id == task_id, McpTaskRow.notification_lease_owner == lease_owner, + McpTaskRow.notification_lease_token == notification_lease_token, McpTaskRow.notification_lease_expires_at >= now, McpTaskRow.dispatch_version == dispatch_version, McpTaskRow.notification_status == "dispatched", @@ -728,6 +835,7 @@ class McpTaskRepository: next_notification_at=next_notification_at, notification_lease_owner=None, notification_lease_expires_at=None, + notification_lease_token=None, updated_at=now, ) ) diff --git a/backend/packages/harness/deerflow/persistence/migrations/AGENTS.md b/backend/packages/harness/deerflow/persistence/migrations/AGENTS.md index 6c5663e2b..123851c17 100644 --- a/backend/packages/harness/deerflow/persistence/migrations/AGENTS.md +++ b/backend/packages/harness/deerflow/persistence/migrations/AGENTS.md @@ -24,10 +24,12 @@ The empty-DB path keeps using `create_all` because `Base.metadata` is the only a `0020_threads_meta_project_id` → `0021_batch_acceptance` → `0019_thread_incarnations` → `0022_scheduled_occurrence_seq` → `0023_run_change_seq` → `0023_user_preferences` → -`0024_project_documents` → `0025_repair_run_change_seq` (current head). The preference +`0024_project_documents` → `0025_repair_run_change_seq` → +`0026_mcp_task_lease_tokens` (current head). The preference revision adds a separate owner/key table with a cascading users foreign key and does not alter users; the project-documents revision adds a new owner-scoped -shelf table, so the bootstrap forward-compat floor is unchanged. +shelf table, and the MCP lease-token revision adds two nullable token columns to +`mcp_tasks`, so the bootstrap forward-compat floor is unchanged. The incarnation revision deliberately retains the exact id audited by the rollback-floor binary; Alembic orders revisions by `down_revision`, not by the numeric prefix. @@ -152,7 +154,7 @@ on installs that never enabled it. The convention is: - `migrations/versions/0014_managed_subagents.py` — creates the deployment-level managed Subagent catalog table - `migrations/versions/0015_scheduled_task_enqueue.py` — interrupts legacy transient queued rows, adds durable scheduled-run launch leases and attempt counts, expands the one-active-occurrence index to `queued`/`launching`/`running`, and migrates the overlap policy from `skip` to `enqueue`; chains after `0014_managed_subagents` - `migrations/versions/0016_subagent_batches.py` — creates durable native-subagent batch and item tables, including owner/submission idempotency, item identity, lease/recovery state, and result fields -- `migrations/versions/0017_personal_access_tokens.py` — creates the personal access token table for programmatic API access +- `migrations/versions/0017_personal_access_tokens.py` — creates the personal access token table for programmatic API access; chains after `0016_subagent_batches` - `migrations/versions/0018_oauth_identity_pg_partial.py` — converts `idx_users_oauth_identity` to a partial index on Postgres (`postgresql_where`), matching what `UserRow.__table_args__` already builds via `create_all`; `0001_baseline` never applied the predicate on Postgres, so every `alembic upgrade head`-provisioned deployment carried a full index until this revision. Postgres-only, idempotent (checks `pg_index.indpred` directly), no-op on SQLite (already partial via `sqlite_where`) and on a DB where the index doesn't exist yet. Originally generated as 0017 and renumbered to 0018 after 0017_personal_access_tokens merged first and kept that slot - `migrations/versions/0019_projects.py` — creates the `projects` table (id/user_id/name/instructions/presentation/status + timestamps) for the Projects Phase-1 organization feature; chains after `0018_oauth_identity_pg_partial` - `migrations/versions/0020_threads_meta_project_id.py` — adds nullable `threads_meta.project_id` plus `ix_threads_meta_project_id` (no FK by design: project delete clears membership first, and the reserved `deerflow_project_id` metadata key stays in sync); chains after `0019_projects` @@ -162,6 +164,7 @@ on installs that never enabled it. The convention is: - `migrations/versions/0023_run_change_seq.py` — adds `runs.change_seq`, its global singleton allocation clock, and owner-aware cursor indexes. Legacy rows remain at zero and page by run id; lifecycle, cancellation, and model-name mutations allocate monotonically increasing positions in their own transaction. Atomic replacement uses one position for every affected row. Progress snapshots and lease heartbeats do not advance the clock. `0023_user_preferences` follows this revision. Its migration test verifies membership in the single-head chain and the expected predecessor rather than pinning the latest head, so later migrations can extend the chain. - `migrations/versions/0024_project_documents.py` — creates the `project_documents` shelf table (id/project_id/user_id/name/stored_relpath/sha256/size_bytes, nullable promotion provenance and trash fields, timestamps) with indexes on project_id, user_id, sha256 and trashed_at; no DB-level foreign key on project_id by design (project delete trashes the shelf inside its own transaction). New table, so the bootstrap forward-compat floor is unchanged; chains after `0023_user_preferences` (renumbered from 0023 after the rebase) - `migrations/versions/0025_repair_run_change_seq.py` — heals databases that skipped `0023_run_change_seq` because it was inserted ahead of the already-shipped `0023_user_preferences` (#5516): re-applies the guarded `run_change_clock` table, `runs.change_seq` column, and cursor indexes on upgrade; no-ops on healthy shapes; chains after `0024_project_documents`. Its downgrade is intentionally a no-op: the repaired objects belong to ancestor `0023_run_change_seq`, remain required at 0024, and must retain their existing change positions. Only the original 0023 downgrade removes them. `tests/test_run_change_repair_history.py` reconstructs both pre-insertion published descendants and verifies historical upgrade, unchanged healthy positions, and usable run-store writes after downgrade and re-upgrade +- `migrations/versions/0026_mcp_task_lease_tokens.py` — chains after `0025_repair_run_change_seq` and adds nullable `mcp_tasks.lease_token` / `notification_lease_token` columns so every poll, cancel, and notification mutation can be fenced to the exact claim generation - `persistence/bootstrap.py` — `bootstrap_schema(engine, backend=...)`, the three-branch provisioning decision, locked revision validation, and the narrow 0019 forward-compatibility exception - `extensions/loader.py::load_extensions` — registers each spec's `table_prefix` with `register_extension_table_prefix()` - Tests: `tests/test_persistence_bootstrap.py` (branches), `tests/test_persistence_bootstrap_concurrency.py` (concurrency), `tests/test_persistence_bootstrap_regression.py` (issue #3682), `tests/test_persistence_migrations_env.py` (filter, including extension-owned tables), `tests/test_extension_loader.py::TestTablePrefixRegistration` (spec-to-filter wiring), `tests/blocking_io/test_persistence_bootstrap.py` (asyncio.to_thread anchor), `tests/test_migration_0004_run_ownership_dedupe.py` + `tests/test_migration_0007_scheduled_run_active_dedupe.py` (dedupe-before-unique-index pre-steps), `tests/test_migration_0025_repair_run_change_seq.py` (issue #5516 skipped-revision heal) diff --git a/backend/packages/harness/deerflow/persistence/migrations/versions/0026_mcp_task_lease_tokens.py b/backend/packages/harness/deerflow/persistence/migrations/versions/0026_mcp_task_lease_tokens.py new file mode 100644 index 000000000..167917165 --- /dev/null +++ b/backend/packages/harness/deerflow/persistence/migrations/versions/0026_mcp_task_lease_tokens.py @@ -0,0 +1,31 @@ +"""fence MCP task claims by per-claim lease tokens. + +Revision ID: 0026_mcp_task_lease_tokens +Revises: 0025_repair_run_change_seq +Create Date: 2026-08-27 +""" + +from __future__ import annotations + +from collections.abc import Sequence + +import sqlalchemy as sa + +revision: str = "0026_mcp_task_lease_tokens" +down_revision: str | Sequence[str] | None = "0025_repair_run_change_seq" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +def upgrade() -> None: + from deerflow.persistence.migrations._helpers import safe_add_column + + safe_add_column("mcp_tasks", sa.Column("lease_token", sa.String(length=64), nullable=True)) + safe_add_column("mcp_tasks", sa.Column("notification_lease_token", sa.String(length=64), nullable=True)) + + +def downgrade() -> None: + from deerflow.persistence.migrations._helpers import safe_drop_column + + safe_drop_column("mcp_tasks", "notification_lease_token") + safe_drop_column("mcp_tasks", "lease_token") diff --git a/backend/packages/harness/deerflow/runtime/cancellation.py b/backend/packages/harness/deerflow/runtime/cancellation.py new file mode 100644 index 000000000..6e2cd99e8 --- /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/tests/AGENTS.md b/backend/tests/AGENTS.md index b6cea4f93..b30a789cd 100644 --- a/backend/tests/AGENTS.md +++ b/backend/tests/AGENTS.md @@ -2,6 +2,16 @@ Backend tests must preserve the runtime invariants they exercise without changing production execution topology. +## MCP claim fencing + +`test_mcp_task_repository.py` covers same-worker reclaim during an in-flight +release, poll/cancel snapshot, or notification completion. Use explicit events +to pause the old operation at the persistence boundary, reclaim via the real +repository, then verify the entire new row remains unchanged. Reclaiming before +the old operation starts does not catch SQLite SELECT/ORM-flush races. Keep the +old completion timestamp within its original lease so expiry cannot mask a +missing token fence; always drain paused tasks and restore session patches. + ## Executor starvation tests `test_executor_starvation.py` covers the deterministic starvation semantics from RFC #4560: diff --git a/backend/tests/test_agent_guidance_check.py b/backend/tests/test_agent_guidance_check.py index af1c07e4b..e9c9d60f9 100644 --- a/backend/tests/test_agent_guidance_check.py +++ b/backend/tests/test_agent_guidance_check.py @@ -173,6 +173,18 @@ def test_local_guidance_files_contain_the_split_original_sections() -> None: assert "Before changing files in this directory" not in text, relative_text +def test_mcp_task_lease_token_migration_is_documented() -> None: + guidance = (REPO_ROOT / "backend" / "packages" / "harness" / "deerflow" / "persistence" / "migrations" / "AGENTS.md").read_text(encoding="utf-8") + + for required in ( + "0026_mcp_task_lease_tokens.py", + "0016_subagent_batches", + "lease_token", + "notification_lease_token", + ): + assert required in guidance + + def test_repository_exposes_one_local_and_one_ci_entrypoint() -> None: makefile = (REPO_ROOT / "Makefile").read_text(encoding="utf-8") workflow = (REPO_ROOT / ".github" / "workflows" / "lint-check.yml").read_text(encoding="utf-8") diff --git a/backend/tests/test_mcp_task_repository.py b/backend/tests/test_mcp_task_repository.py index 87db91e66..eb429f2d6 100644 --- a/backend/tests/test_mcp_task_repository.py +++ b/backend/tests/test_mcp_task_repository.py @@ -1,3 +1,4 @@ +import asyncio import contextlib import sqlite3 from datetime import UTC, datetime, timedelta @@ -6,6 +7,7 @@ import pytest import pytest_asyncio from sqlalchemy import event from sqlalchemy.exc import IntegrityError +from sqlalchemy.ext.asyncio import AsyncSession from deerflow.config.database_config import DatabaseConfig from deerflow.persistence.engine import close_engine, get_engine, get_session_factory, init_engine_from_config @@ -57,6 +59,147 @@ async def _create_working_task( ) +@contextlib.asynccontextmanager +async def _pause_claim_mutation(monkeypatch, operation): + """Pause an old mutation while another session reclaims its row. + + The former SELECT/ORM-flush implementation must pause after its ownership + read has loaded the old row. The atomic implementation pauses before its + conditional UPDATE. Both leave the competing claim free to commit using + the production SQLite engine, without replacing any persistence logic. + """ + entered = asyncio.Event() + resume = asyncio.Event() + original_execute = AsyncSession.execute + intercepted = False + + async def execute(session, statement, *args, **kwargs): + nonlocal intercepted + if asyncio.current_task() is not task or intercepted: + return await original_execute(session, statement, *args, **kwargs) + intercepted = True + if statement.is_select: + result = await original_execute(session, statement, *args, **kwargs) + entered.set() + await resume.wait() + if statement.is_select: + return result + return await original_execute(session, statement, *args, **kwargs) + + with monkeypatch.context() as patch: + patch.setattr(AsyncSession, "execute", execute) + task = asyncio.create_task(operation) + try: + await asyncio.wait_for(entered.wait(), timeout=5) + yield task, resume + finally: + resume.set() + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("operation", ["release_claim", "apply_snapshot", "apply_cancel_snapshot"]) +async def test_interleaved_reclaim_fences_inflight_poll_and_cancel_mutations(tmp_path, monkeypatch, operation): + repo = await _make_repo(tmp_path) + now = datetime(2026, 1, 1, tzinfo=UTC) + task_id = "interleaved-claim" + await _create_working_task(repo, task_id=task_id, now=now) + claim = repo.claim_due_tasks + if operation == "apply_cancel_snapshot": + await repo.request_cancel(task_id, user_id="user-1", thread_id="thread-1", requested_at=now) + claim = repo.claim_cancel_requests + first = await claim(now=now, lease_owner="worker-1", lease_seconds=60, limit=1) + kwargs = {"lease_owner": "worker-1", "lease_token": first[0]["lease_token"]} + if operation == "release_claim": + kwargs.update(next_poll_at=now + timedelta(seconds=30), error="old poll failed") + else: + kwargs.update( + status="cancelled" if operation == "apply_cancel_snapshot" else "completed", + result={"stale": True}, + result_preview="old result", + result_truncated=False, + result_artifact=None, + error=None, + input_required=None, + ) + if operation == "apply_snapshot": + kwargs.update(next_poll_at=None, polled_at=now) + else: + kwargs.update(completed_at=now) + + async with _pause_claim_mutation(monkeypatch, getattr(repo, operation)(task_id, **kwargs)) as (pending, resume): + # Advance only the claim clock, not the stale operation's completion + # timestamp: expiry must not reject it before the token fence is tested. + second = await asyncio.wait_for(claim(now=now + timedelta(seconds=61), lease_owner="worker-1", lease_seconds=60, limit=1), timeout=5) + assert len(second) == 1 + assert second[0]["lease_token"] != first[0]["lease_token"] + before = await repo.get(task_id, user_id="user-1") + resume.set() + applied = await asyncio.wait_for(pending, timeout=5) + + # Check the entire row, including scheduling, errors, results and event + # versions, not just the new lease: stale work must have no side effects. + assert await repo.get(task_id, user_id="user-1") == before + assert applied is False + + +@pytest.mark.asyncio +@pytest.mark.parametrize("delivered", [True, False], ids=["success", "failure"]) +async def test_interleaved_reclaim_fences_inflight_notification_completion(tmp_path, monkeypatch, delivered): + repo = await _make_repo(tmp_path) + now = datetime(2026, 1, 1, tzinfo=UTC) + task_id = "interleaved-notification" + await _create_working_task(repo, task_id=task_id, now=now) + poll = await repo.claim_due_tasks(now=now, lease_owner="poller", lease_seconds=60, limit=1) + assert await repo.apply_snapshot( + task_id, + lease_owner="poller", + lease_token=poll[0]["lease_token"], + status="completed", + result={"done": True}, + result_preview=None, + result_truncated=False, + result_artifact=None, + error=None, + input_required=None, + next_poll_at=None, + polled_at=now, + ) + claim_kwargs = {"lease_owner": "notifier", "lease_seconds": 60, "limit": 1, "tracking_degraded_after_errors": 3} + launch = await repo.claim_notification_work(now=now, **claim_kwargs) + assert await repo.mark_notification_dispatched( + task_id, + lease_owner="notifier", + notification_lease_token=launch[0]["notification_lease_token"], + dispatch_version=launch[0]["dispatch_version"], + run_id="notification-run", + now=now, + ) + first = await repo.claim_notification_work(now=now, **claim_kwargs) + operation = repo.finish_notification_run( + task_id, + lease_owner="notifier", + notification_lease_token=first[0]["notification_lease_token"], + dispatch_version=first[0]["dispatch_version"], + delivered=delivered, + next_notification_at=None if delivered else now + timedelta(seconds=30), + error=None if delivered else "old notification failed", + now=now, + ) + async with _pause_claim_mutation(monkeypatch, operation) as (pending, resume): + second = await asyncio.wait_for(repo.claim_notification_work(now=now + timedelta(seconds=61), **claim_kwargs), timeout=5) + assert len(second) == 1 + assert second[0]["notification_lease_token"] != first[0]["notification_lease_token"] + before = await repo.get(task_id, user_id="user-1") + resume.set() + applied = await asyncio.wait_for(pending, timeout=5) + + assert await repo.get(task_id, user_id="user-1") == before + assert applied is False + + @pytest.mark.asyncio async def test_legacy_task_writer_leaves_thread_incarnation_null(tmp_path): repo = await _make_repo(tmp_path) @@ -301,7 +444,7 @@ async def test_apply_snapshot_requires_current_lease_owner_and_terminalizes_task repo = await _make_repo(tmp_path) now = datetime.now(UTC) await _create_working_task(repo, task_id="task-2", now=now) - await repo.claim_due_tasks( + claimed = await repo.claim_due_tasks( now=now, lease_owner="worker-new", lease_seconds=60, @@ -311,6 +454,7 @@ async def test_apply_snapshot_requires_current_lease_owner_and_terminalizes_task stale_applied = await repo.apply_snapshot( "task-2", lease_owner="worker-old", + lease_token=claimed[0]["lease_token"], status="failed", result=None, result_preview=None, @@ -326,6 +470,7 @@ async def test_apply_snapshot_requires_current_lease_owner_and_terminalizes_task applied = await repo.apply_snapshot( "task-2", lease_owner="worker-new", + lease_token=claimed[0]["lease_token"], status="completed", result={"report": "ready"}, result_preview=None, @@ -365,7 +510,7 @@ async def test_apply_snapshot_rejects_result_after_same_workers_lease_expires(tm repo = await _make_repo(tmp_path) now = datetime.now(UTC) await _create_working_task(repo, task_id="task-expired", now=now) - await repo.claim_due_tasks( + claimed = await repo.claim_due_tasks( now=now, lease_owner="worker-1", lease_seconds=60, @@ -375,6 +520,7 @@ async def test_apply_snapshot_rejects_result_after_same_workers_lease_expires(tm applied = await repo.apply_snapshot( "task-expired", lease_owner="worker-1", + lease_token=claimed[0]["lease_token"], status="completed", result={"report": "stale"}, result_preview=None, @@ -398,7 +544,7 @@ async def test_input_required_is_persisted_and_remains_scheduled_for_slow_pollin repo = await _make_repo(tmp_path) now = datetime.now(UTC) await _create_working_task(repo, task_id="task-3", now=now) - await repo.claim_due_tasks( + claimed = await repo.claim_due_tasks( now=now, lease_owner="worker-1", lease_seconds=60, @@ -408,6 +554,7 @@ async def test_input_required_is_persisted_and_remains_scheduled_for_slow_pollin applied = await repo.apply_snapshot( "task-3", lease_owner="worker-1", + lease_token=claimed[0]["lease_token"], status="input_required", result=None, result_preview=None, @@ -432,7 +579,7 @@ async def test_release_claim_retries_transient_poll_failure(tmp_path): repo = await _make_repo(tmp_path) now = datetime.now(UTC) await _create_working_task(repo, task_id="task-4", now=now) - await repo.claim_due_tasks( + claimed = await repo.claim_due_tasks( now=now, lease_owner="worker-1", lease_seconds=60, @@ -443,6 +590,7 @@ async def test_release_claim_retries_transient_poll_failure(tmp_path): released = await repo.release_claim( "task-4", lease_owner="worker-1", + lease_token=claimed[0]["lease_token"], next_poll_at=retry_at, error="temporary network failure", ) @@ -456,6 +604,231 @@ async def test_release_claim_retries_transient_poll_failure(tmp_path): assert stored["lease_owner"] is None +@pytest.mark.asyncio +async def test_release_claim_after_same_worker_reclaim_cannot_clear_new_claim(tmp_path): + """A stale release from an older generation must be a no-op once the same + worker reclaims the task with a fresh per-claim token (token fencing).""" + repo = await _make_repo(tmp_path) + now = datetime.now(UTC) + await _create_working_task(repo, task_id="task-fence", now=now) + claimed = await repo.claim_due_tasks(now=now, lease_owner="worker-1", lease_seconds=60, limit=10) + old_token = claimed[0]["lease_token"] + + reclaim_at = now + timedelta(seconds=61) # after the 60s lease expires + reclaimed = await repo.claim_due_tasks(now=reclaim_at, lease_owner="worker-1", lease_seconds=61, limit=10) + assert reclaimed + new_token = reclaimed[0]["lease_token"] + assert new_token != old_token + + # The stale release (old owner + old token) must not clear the new claim. + released = await repo.release_claim( + "task-fence", + lease_owner="worker-1", + lease_token=old_token, + next_poll_at=reclaim_at + timedelta(seconds=30), + error="stale release", + ) + assert released is False + + stored = await repo.get("task-fence", user_id="user-1") + assert stored is not None + assert stored["lease_owner"] == "worker-1" + assert stored["lease_token"] == new_token + assert stored["lease_expires_at"] is not None + + +@pytest.mark.asyncio +async def test_apply_snapshot_after_same_worker_reclaim_cannot_clear_new_claim(tmp_path): + """A poll snapshot from an older generation must not overwrite a newer claim.""" + repo = await _make_repo(tmp_path) + now = datetime.now(UTC) + await _create_working_task(repo, task_id="task-apply-fence", now=now) + claimed = await repo.claim_due_tasks(now=now, lease_owner="worker-1", lease_seconds=60, limit=10) + old_token = claimed[0]["lease_token"] + + reclaim_at = now + timedelta(seconds=61) + reclaimed = await repo.claim_due_tasks(now=reclaim_at, lease_owner="worker-1", lease_seconds=61, limit=10) + new_token = reclaimed[0]["lease_token"] + assert new_token != old_token + + applied = await repo.apply_snapshot( + "task-apply-fence", + lease_owner="worker-1", + lease_token=old_token, + status="completed", + result=None, + result_preview=None, + result_truncated=False, + result_artifact=None, + error=None, + input_required=None, + next_poll_at=None, + polled_at=reclaim_at, + ) + assert applied is False + + stored = await repo.get("task-apply-fence", user_id="user-1") + assert stored is not None + assert stored["lease_owner"] == "worker-1" + assert stored["lease_token"] == new_token + assert stored["status"] == "working" + + +@pytest.mark.asyncio +async def test_apply_cancel_snapshot_after_same_worker_reclaim_cannot_clear_new_claim(tmp_path): + """A cancel snapshot from an older generation must not overwrite a newer claim.""" + repo = await _make_repo(tmp_path) + now = datetime.now(UTC) + await _create_working_task(repo, task_id="task-cancel-fence", now=now) + claimed = await repo.claim_due_tasks(now=now, lease_owner="worker-1", lease_seconds=60, limit=10) + old_token = claimed[0]["lease_token"] + + reclaim_at = now + timedelta(seconds=61) + reclaimed = await repo.claim_due_tasks(now=reclaim_at, lease_owner="worker-1", lease_seconds=61, limit=10) + new_token = reclaimed[0]["lease_token"] + assert new_token != old_token + + applied = await repo.apply_cancel_snapshot( + "task-cancel-fence", + lease_owner="worker-1", + lease_token=old_token, + status="cancelled", + result=None, + result_preview=None, + result_truncated=False, + result_artifact=None, + error=None, + input_required=None, + completed_at=reclaim_at, + ) + assert applied is False + + stored = await repo.get("task-cancel-fence", user_id="user-1") + assert stored is not None + assert stored["lease_owner"] == "worker-1" + assert stored["lease_token"] == new_token + assert stored["status"] == "working" + + +@pytest.mark.asyncio +async def test_finish_notification_run_after_reclaim_cannot_clear_new_claim(tmp_path): + """A stale notification finish must not clear a newer notification lease.""" + repo = await _make_repo(tmp_path) + now = datetime.now(UTC) + await _create_working_task(repo, task_id="task-notify-fence", now=now) + poll_claim = await repo.claim_due_tasks(now=now, lease_owner="poller", lease_seconds=60, limit=1) + await repo.apply_snapshot( + "task-notify-fence", + lease_owner="poller", + lease_token=poll_claim[0]["lease_token"], + status="input_required", + result=None, + result_preview=None, + result_truncated=False, + result_artifact=None, + error=None, + input_required={"prompt": "Approve?"}, + next_poll_at=now, + polled_at=now, + ) + first = await repo.claim_notification_work( + now=now, + lease_owner="notifier", + lease_seconds=60, + limit=1, + tracking_degraded_after_errors=3, + ) + await repo.mark_notification_dispatched( + "task-notify-fence", + lease_owner="notifier", + notification_lease_token=first[0]["notification_lease_token"], + dispatch_version=first[0]["dispatch_version"], + run_id="notify-run-1", + now=now, + ) + reclaimed = await repo.claim_notification_work( + now=now, + lease_owner="notifier", + lease_seconds=60, + limit=1, + tracking_degraded_after_errors=3, + ) + assert reclaimed + new_notify_token = reclaimed[0]["notification_lease_token"] + + finished = await repo.finish_notification_run( + "task-notify-fence", + lease_owner="notifier", + notification_lease_token="stale-notify-token", + dispatch_version=reclaimed[0]["dispatch_version"], + delivered=True, + next_notification_at=None, + error=None, + now=now, + ) + assert finished is False + + stored = await repo.get("task-notify-fence", user_id="user-1") + assert stored is not None + assert stored["notification_lease_owner"] == "notifier" + assert stored["notification_lease_token"] == new_notify_token + + +@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) + claimed = 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", + lease_token=claimed[0]["lease_token"], + 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 + + reclaimed = 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", + lease_token=reclaimed[0]["lease_token"], + ) + + 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) + claimed = 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", + lease_token=claimed[0]["lease_token"], + ) + + 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) @@ -463,10 +836,11 @@ async def test_consecutive_poll_error_count_increments_and_resets_on_success(tmp await _create_working_task(repo, task_id="task-6", now=now) for expected_errors in (1, 2): - await repo.claim_due_tasks(now=now, lease_owner="worker-1", lease_seconds=60, limit=10) + claimed = await repo.claim_due_tasks(now=now, lease_owner="worker-1", lease_seconds=60, limit=10) await repo.release_claim( "task-6", lease_owner="worker-1", + lease_token=claimed[0]["lease_token"], next_poll_at=now - timedelta(seconds=1), error="temporary network failure", ) @@ -474,10 +848,11 @@ async def test_consecutive_poll_error_count_increments_and_resets_on_success(tmp assert stored is not None assert stored["consecutive_poll_error_count"] == expected_errors - await repo.claim_due_tasks(now=now, lease_owner="worker-1", lease_seconds=60, limit=10) + claimed = await repo.claim_due_tasks(now=now, lease_owner="worker-1", lease_seconds=60, limit=10) applied = await repo.apply_snapshot( "task-6", lease_owner="worker-1", + lease_token=claimed[0]["lease_token"], status="working", result=None, result_preview=None, @@ -500,10 +875,11 @@ async def test_notification_snapshot_is_versioned_and_not_overwritten_in_flight( repo = await _make_repo(tmp_path) now = datetime.now(UTC) await _create_working_task(repo, task_id="task-notify", now=now) - await repo.claim_due_tasks(now=now, lease_owner="poller", lease_seconds=60, limit=1) + poll_claim = await repo.claim_due_tasks(now=now, lease_owner="poller", lease_seconds=60, limit=1) await repo.apply_snapshot( "task-notify", lease_owner="poller", + lease_token=poll_claim[0]["lease_token"], status="input_required", result=None, result_preview=None, @@ -525,10 +901,11 @@ async def test_notification_snapshot_is_versioned_and_not_overwritten_in_flight( assert first[0]["dispatch_version"] == 1 assert first[0]["dispatch_event"]["input_required"] == {"prompt": "Approve?"} - await repo.claim_due_tasks(now=now, lease_owner="poller", lease_seconds=60, limit=1) + poll_claim = await repo.claim_due_tasks(now=now, lease_owner="poller", lease_seconds=60, limit=1) await repo.apply_snapshot( "task-notify", lease_owner="poller", + lease_token=poll_claim[0]["lease_token"], status="completed", result={"done": True}, result_preview=None, @@ -548,11 +925,12 @@ async def test_notification_snapshot_is_versioned_and_not_overwritten_in_flight( await repo.mark_notification_dispatched( "task-notify", lease_owner="notifier", + notification_lease_token=first[0]["notification_lease_token"], dispatch_version=1, run_id="notify-run-1", now=now, ) - await repo.claim_notification_work( + dispatched_claim = await repo.claim_notification_work( now=now, lease_owner="notifier", lease_seconds=60, @@ -562,6 +940,7 @@ async def test_notification_snapshot_is_versioned_and_not_overwritten_in_flight( await repo.finish_notification_run( "task-notify", lease_owner="notifier", + notification_lease_token=dispatched_claim[0]["notification_lease_token"], dispatch_version=1, delivered=True, next_notification_at=None, @@ -584,10 +963,11 @@ async def test_notification_retry_rebuilds_a_newer_event_and_resets_its_budget(t repo = await _make_repo(tmp_path) now = datetime.now(UTC) await _create_working_task(repo, task_id="task-retry-latest", now=now) - await repo.claim_due_tasks(now=now, lease_owner="poller", lease_seconds=60, limit=1) + poll_claim = await repo.claim_due_tasks(now=now, lease_owner="poller", lease_seconds=60, limit=1) await repo.apply_snapshot( "task-retry-latest", lease_owner="poller", + lease_token=poll_claim[0]["lease_token"], status="input_required", result=None, result_preview=None, @@ -608,11 +988,12 @@ async def test_notification_retry_rebuilds_a_newer_event_and_resets_its_budget(t await repo.mark_notification_dispatched( "task-retry-latest", lease_owner="notifier", + notification_lease_token=first[0]["notification_lease_token"], dispatch_version=first[0]["dispatch_version"], run_id="notify-run-1", now=now, ) - await repo.claim_notification_work( + dispatched_claim = await repo.claim_notification_work( now=now, lease_owner="notifier", lease_seconds=60, @@ -623,6 +1004,7 @@ async def test_notification_retry_rebuilds_a_newer_event_and_resets_its_budget(t await repo.finish_notification_run( "task-retry-latest", lease_owner="notifier", + notification_lease_token=dispatched_claim[0]["notification_lease_token"], dispatch_version=first[0]["dispatch_version"], delivered=False, next_notification_at=retry_at, @@ -635,10 +1017,11 @@ async def test_notification_retry_rebuilds_a_newer_event_and_resets_its_budget(t assert failed["dispatch_attempt"] == 1 assert failed["notification_attempt_count"] == 1 - await repo.claim_due_tasks(now=now, lease_owner="poller", lease_seconds=60, limit=1) + poll_claim = await repo.claim_due_tasks(now=now, lease_owner="poller", lease_seconds=60, limit=1) await repo.apply_snapshot( "task-retry-latest", lease_owner="poller", + lease_token=poll_claim[0]["lease_token"], status="completed", result={"done": True}, result_preview=None, @@ -668,10 +1051,11 @@ async def test_unexpected_notification_failure_releases_lease_without_changing_p repo = await _make_repo(tmp_path) now = datetime.now(UTC) await _create_working_task(repo, task_id="task-notify-release", now=now) - await repo.claim_due_tasks(now=now, lease_owner="poller", lease_seconds=60, limit=1) + poll_claim = await repo.claim_due_tasks(now=now, lease_owner="poller", lease_seconds=60, limit=1) await repo.apply_snapshot( "task-notify-release", lease_owner="poller", + lease_token=poll_claim[0]["lease_token"], status="input_required", result=None, result_preview=None, @@ -682,7 +1066,7 @@ async def test_unexpected_notification_failure_releases_lease_without_changing_p next_poll_at=now, polled_at=now, ) - await repo.claim_notification_work( + claimed = await repo.claim_notification_work( now=now, lease_owner="notifier", lease_seconds=60, @@ -694,6 +1078,7 @@ async def test_unexpected_notification_failure_releases_lease_without_changing_p assert await repo.release_notification_lease( "task-notify-release", lease_owner="notifier", + notification_lease_token=claimed[0]["notification_lease_token"], next_notification_at=retry_at, error="run store unavailable", ) @@ -711,10 +1096,11 @@ async def test_notification_launch_failure_counts_and_reclaims_latest_snapshot(t repo = await _make_repo(tmp_path) now = datetime.now(UTC) await _create_working_task(repo, task_id="task-launch-retry", now=now) - await repo.claim_due_tasks(now=now, lease_owner="poller", lease_seconds=60, limit=1) + poll_claim = await repo.claim_due_tasks(now=now, lease_owner="poller", lease_seconds=60, limit=1) await repo.apply_snapshot( "task-launch-retry", lease_owner="poller", + lease_token=poll_claim[0]["lease_token"], status="input_required", result=None, result_preview=None, @@ -737,6 +1123,7 @@ async def test_notification_launch_failure_counts_and_reclaims_latest_snapshot(t assert await repo.release_notification_claim( "task-launch-retry", lease_owner="notifier", + notification_lease_token=first[0]["notification_lease_token"], next_notification_at=retry_at, error="run store unavailable", replace_with_latest=True, @@ -765,10 +1152,11 @@ async def test_permanent_notification_failure_is_not_reclaimed(tmp_path): repo = await _make_repo(tmp_path) now = datetime.now(UTC) await _create_working_task(repo, task_id="task-dead-letter", now=now) - await repo.claim_due_tasks(now=now, lease_owner="poller", lease_seconds=60, limit=1) + poll_claim = await repo.claim_due_tasks(now=now, lease_owner="poller", lease_seconds=60, limit=1) await repo.apply_snapshot( "task-dead-letter", lease_owner="poller", + lease_token=poll_claim[0]["lease_token"], status="completed", result={"done": True}, result_preview=None, @@ -790,6 +1178,7 @@ async def test_permanent_notification_failure_is_not_reclaimed(tmp_path): assert await repo.dead_letter_notification( "task-dead-letter", lease_owner="notifier", + notification_lease_token=claimed[0]["notification_lease_token"], dispatch_version=claimed[0]["dispatch_version"], error="Thread deleted-thread not found", count_failure=True, @@ -818,10 +1207,11 @@ async def test_dispatched_notification_can_be_dead_lettered_after_retry_budget(t repo = await _make_repo(tmp_path) now = datetime.now(UTC) await _create_working_task(repo, task_id="task-dispatched-budget", now=now) - await repo.claim_due_tasks(now=now, lease_owner="poller", lease_seconds=60, limit=1) + poll_claim = await repo.claim_due_tasks(now=now, lease_owner="poller", lease_seconds=60, limit=1) await repo.apply_snapshot( "task-dispatched-budget", lease_owner="poller", + lease_token=poll_claim[0]["lease_token"], status="completed", result={"done": True}, result_preview=None, @@ -843,6 +1233,7 @@ async def test_dispatched_notification_can_be_dead_lettered_after_retry_budget(t assert await repo.mark_notification_dispatched( "task-dispatched-budget", lease_owner="notifier", + notification_lease_token=first[0]["notification_lease_token"], dispatch_version=dispatch_version, run_id="notify-run-1", now=now, @@ -859,6 +1250,7 @@ async def test_dispatched_notification_can_be_dead_lettered_after_retry_budget(t assert await repo.dead_letter_notification( "task-dispatched-budget", lease_owner="budget-checker", + notification_lease_token=claimed[0]["notification_lease_token"], dispatch_version=dispatch_version, error="Notification delivery stopped after 5 failed attempts", count_failure=False, @@ -876,10 +1268,11 @@ async def test_dead_lettering_dispatched_snapshot_preserves_newer_event(tmp_path repo = await _make_repo(tmp_path) now = datetime.now(UTC) await _create_working_task(repo, task_id="task-dispatched-latest", now=now) - await repo.claim_due_tasks(now=now, lease_owner="poller", lease_seconds=60, limit=1) + poll_claim = await repo.claim_due_tasks(now=now, lease_owner="poller", lease_seconds=60, limit=1) await repo.apply_snapshot( "task-dispatched-latest", lease_owner="poller", + lease_token=poll_claim[0]["lease_token"], status="input_required", result=None, result_preview=None, @@ -901,15 +1294,17 @@ async def test_dead_lettering_dispatched_snapshot_preserves_newer_event(tmp_path assert await repo.mark_notification_dispatched( "task-dispatched-latest", lease_owner="notifier", + notification_lease_token=first[0]["notification_lease_token"], dispatch_version=dispatch_version, run_id="notify-run-1", now=now, ) - await repo.claim_due_tasks(now=now, lease_owner="poller", lease_seconds=60, limit=1) + poll_claim = await repo.claim_due_tasks(now=now, lease_owner="poller", lease_seconds=60, limit=1) await repo.apply_snapshot( "task-dispatched-latest", lease_owner="poller", + lease_token=poll_claim[0]["lease_token"], status="completed", result={"done": True}, result_preview=None, @@ -931,6 +1326,7 @@ async def test_dead_lettering_dispatched_snapshot_preserves_newer_event(tmp_path assert await repo.dead_letter_notification( "task-dispatched-latest", lease_owner="budget-checker", + notification_lease_token=claimed[0]["notification_lease_token"], dispatch_version=dispatch_version, error="old snapshot exhausted its retry budget", count_failure=False, @@ -958,7 +1354,7 @@ async def test_cancel_request_stops_polling_and_rejects_stale_poll_result(tmp_pa repo = await _make_repo(tmp_path) now = datetime.now(UTC) await _create_working_task(repo, task_id="task-cancel", now=now) - await repo.claim_due_tasks(now=now, lease_owner="stale-poller", lease_seconds=60, limit=1) + stale_poll_claim = await repo.claim_due_tasks(now=now, lease_owner="stale-poller", lease_seconds=60, limit=1) requested = await repo.request_cancel( "task-cancel", @@ -972,6 +1368,7 @@ async def test_cancel_request_stops_polling_and_rejects_stale_poll_result(tmp_pa await repo.apply_snapshot( "task-cancel", lease_owner="stale-poller", + lease_token=stale_poll_claim[0]["lease_token"], status="completed", result={"stale": True}, result_preview=None, @@ -1006,6 +1403,7 @@ async def test_cancel_request_stops_polling_and_rejects_stale_poll_result(tmp_pa assert await repo.apply_cancel_snapshot( "task-cancel", lease_owner="canceller", + lease_token=claimed[0]["lease_token"], status="cancelled", result=None, result_preview=None, @@ -1019,3 +1417,188 @@ async def test_cancel_request_stops_polling_and_rejects_stale_poll_result(tmp_pa assert stored is not None assert stored["status"] == "cancelled" assert stored["notification_status"] == "pending" + + +@pytest.mark.asyncio +async def test_late_poll_release_after_same_worker_reclaim_is_fenced(tmp_path): + repo = await _make_repo(tmp_path) + now = datetime.now(UTC) + await _create_working_task(repo, task_id="task-late-poll", now=now) + + first = await repo.claim_due_tasks( + now=now, + lease_owner="worker-same", + lease_seconds=1, + limit=10, + ) + assert [row["id"] for row in first] == ["task-late-poll"] + first_token = first[0]["lease_token"] + assert first_token + + reclaimed = await repo.claim_due_tasks( + now=now + timedelta(seconds=5), + lease_owner="worker-same", + lease_seconds=60, + limit=10, + ) + assert [row["id"] for row in reclaimed] == ["task-late-poll"] + assert reclaimed[0]["lease_token"] != first_token + + released = await repo.release_poll_claim_after_cancellation( + "task-late-poll", + lease_owner="worker-same", + lease_token=first_token, + ) + assert released is False + + stored = await repo.get("task-late-poll", user_id="user-1") + assert stored is not None + assert stored["lease_owner"] == "worker-same" + assert stored["lease_token"] == reclaimed[0]["lease_token"] + + +@pytest.mark.asyncio +async def test_late_cancel_release_after_same_worker_reclaim_is_fenced(tmp_path): + repo = await _make_repo(tmp_path) + now = datetime.now(UTC) + await _create_working_task(repo, task_id="task-late-cancel", now=now) + await repo.request_cancel( + "task-late-cancel", + user_id="user-1", + thread_id="thread-1", + requested_at=now, + ) + + first = await repo.claim_cancel_requests( + now=now, + lease_owner="worker-same", + lease_seconds=1, + limit=1, + ) + assert [row["id"] for row in first] == ["task-late-cancel"] + first_token = first[0]["lease_token"] + assert first_token + + reclaimed = await repo.claim_cancel_requests( + now=now + timedelta(seconds=5), + lease_owner="worker-same", + lease_seconds=60, + limit=1, + ) + assert [row["id"] for row in reclaimed] == ["task-late-cancel"] + assert reclaimed[0]["lease_token"] != first_token + + released = await repo.release_cancel_claim( + "task-late-cancel", + lease_owner="worker-same", + lease_token=first_token, + next_cancel_at=now + timedelta(seconds=30), + error="cancelled", + ) + assert released is False + + stored = await repo.get("task-late-cancel", user_id="user-1") + assert stored is not None + assert stored["lease_owner"] == "worker-same" + assert stored["lease_token"] == reclaimed[0]["lease_token"] + + +@pytest.mark.asyncio +async def test_late_notification_release_after_same_worker_reclaim_is_fenced(tmp_path): + repo = await _make_repo(tmp_path) + now = datetime.now(UTC) + await _create_working_task(repo, task_id="task-late-notify", now=now) + poll_claim = await repo.claim_due_tasks(now=now, lease_owner="poller", lease_seconds=60, limit=1) + await repo.apply_snapshot( + "task-late-notify", + lease_owner="poller", + lease_token=poll_claim[0]["lease_token"], + status="completed", + result={"done": True}, + result_preview=None, + result_truncated=False, + result_artifact=None, + error=None, + input_required=None, + next_poll_at=None, + polled_at=now, + ) + + first = await repo.claim_notification_work( + now=now, + lease_owner="notifier-same", + lease_seconds=1, + limit=1, + tracking_degraded_after_errors=3, + ) + assert [row["id"] for row in first] == ["task-late-notify"] + first_token = first[0]["notification_lease_token"] + assert first_token + + reclaimed = await repo.claim_notification_work( + now=now + timedelta(seconds=5), + lease_owner="notifier-same", + lease_seconds=60, + limit=1, + tracking_degraded_after_errors=3, + ) + assert [row["id"] for row in reclaimed] == ["task-late-notify"] + assert reclaimed[0]["notification_lease_token"] != first_token + + released = await repo.release_notification_lease( + "task-late-notify", + lease_owner="notifier-same", + notification_lease_token=first_token, + next_notification_at=now + timedelta(seconds=30), + error="cancelled", + ) + assert released is False + + stored = await repo.get("task-late-notify", user_id="user-1") + assert stored is not None + assert stored["notification_lease_owner"] == "notifier-same" + assert stored["notification_lease_token"] == reclaimed[0]["notification_lease_token"] + + +@pytest.mark.asyncio +async def test_late_snapshot_apply_after_same_worker_reclaim_is_fenced(tmp_path): + repo = await _make_repo(tmp_path) + now = datetime.now(UTC) + await _create_working_task(repo, task_id="task-late-apply", now=now) + + first = await repo.claim_due_tasks( + now=now, + lease_owner="worker-same", + lease_seconds=1, + limit=10, + ) + first_token = first[0]["lease_token"] + reclaimed = await repo.claim_due_tasks( + now=now + timedelta(seconds=5), + lease_owner="worker-same", + lease_seconds=60, + limit=10, + ) + assert [row["id"] for row in reclaimed] == ["task-late-apply"] + assert reclaimed[0]["lease_token"] != first_token + + applied = await repo.apply_snapshot( + "task-late-apply", + lease_owner="worker-same", + lease_token=first_token, + status="completed", + result={"stale": True}, + result_preview=None, + result_truncated=False, + result_artifact=None, + error=None, + input_required=None, + next_poll_at=None, + polled_at=now + timedelta(seconds=5), + ) + assert applied is False + + stored = await repo.get("task-late-apply", user_id="user-1") + assert stored is not None + assert stored["lease_owner"] == "worker-same" + assert stored["lease_token"] == reclaimed[0]["lease_token"] diff --git a/backend/tests/test_mcp_task_service.py b/backend/tests/test_mcp_task_service.py index 54a4c94f9..531efb149 100644 --- a/backend/tests/test_mcp_task_service.py +++ b/backend/tests/test_mcp_task_service.py @@ -1,5 +1,7 @@ import asyncio +import gc import logging +import weakref from datetime import UTC, datetime, timedelta from types import SimpleNamespace from unittest.mock import AsyncMock @@ -48,6 +50,10 @@ class FakeRepository: 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): @@ -168,6 +174,8 @@ def _claimed_row(*, driver_name="fake"): "status": "working", "driver_data": {"status_tool": "status"}, "lease_owner": "ignored-by-service-fixture", + "lease_token": "lease-token-1", + "notification_lease_token": "notify-lease-token-1", } @@ -347,7 +355,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( @@ -644,13 +652,13 @@ async def test_notification_delivery_waits_for_successful_agent_run(): "dispatch_event": {"status": "completed"}, } - await service._notify_one(claimed, now=now) + await service._notify_one_claimed(claimed, now=now) repo.mark_notification_dispatched.assert_awaited_once() repo.finish_notification_run.assert_not_awaited() get_run.return_value = SimpleNamespace(status=RunStatus.success) - await service._notify_one( + await service._notify_one_claimed( { **claimed, "notification_status": "dispatched", @@ -679,7 +687,7 @@ async def test_missing_dispatched_notification_run_retries_delivery(): ) now = datetime.now(UTC) - await service._notify_one( + await service._notify_one_claimed( { **_claimed_row(), "notification_status": "dispatched", @@ -768,7 +776,7 @@ async def test_notification_busy_thread_replaces_claim_with_latest_event(): ) now = datetime.now(UTC) - await service._notify_one( + await service._notify_one_claimed( { **_claimed_row(), "notification_status": "claimed", @@ -801,7 +809,7 @@ async def test_notification_launch_failure_backs_off_and_replaces_with_latest_ev ) now = datetime.now(UTC) - await service._notify_one( + await service._notify_one_claimed( { **_claimed_row(), "notification_status": "claimed", @@ -836,7 +844,7 @@ async def test_permanently_rejected_notification_is_dead_lettered(): ) now = datetime.now(UTC) - await service._notify_one( + await service._notify_one_claimed( { **_claimed_row(), "notification_status": "claimed", @@ -874,7 +882,7 @@ async def test_notification_retry_budget_dead_letters_before_creating_another_ru ) now = datetime.now(UTC) - await service._notify_one( + await service._notify_one_claimed( { **_claimed_row(), "notification_status": "retry", @@ -912,7 +920,7 @@ async def test_dispatched_notification_retry_budget_dead_letters_before_hydratin ) now = datetime.now(UTC) - await service._notify_one( + await service._notify_one_claimed( { **_claimed_row(), "notification_status": "dispatched", @@ -1417,3 +1425,2784 @@ 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 + + +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"} + + +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"] + + +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 == [] + + +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, + phase="probe", + 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 +@pytest.mark.parametrize("phase", ["poll", "cancel", "notification"]) +async def test_single_flight_claim_releases_late_uncancelled_claim(phase, monkeypatch): + monkeypatch.setattr(service_module, "_CANCELLATION_DRAIN_TIMEOUT_SECONDS", 0.01) + + class BlockingClaimRepository: + def __init__(self): + self.claim_calls = {"poll": 0, "cancel": 0, "notification": 0} + self.claim_started = asyncio.Event() + self.claim_gate = asyncio.Event() + self.released: list[tuple[str, str, dict]] = [] + + async def _claim(self, claim_phase): + self.claim_calls[claim_phase] += 1 + self.claim_started.set() + await self.claim_gate.wait() + return [{**_claimed_row(), "notification_status": "pending"}] + + async def claim_due_tasks(self, **_kwargs): + return await self._claim("poll") + + async def claim_cancel_requests(self, **_kwargs): + return await self._claim("cancel") + + async def claim_notification_work(self, **_kwargs): + return await self._claim("notification") + + async def release_poll_claim_after_cancellation(self, task_id, **kwargs): + self.released.append(("poll", task_id, kwargs)) + return True + + async def release_cancel_claim(self, task_id, **kwargs): + self.released.append(("cancel", task_id, kwargs)) + return True + + async def release_notification_claim(self, task_id, **kwargs): + self.released.append(("notification", task_id, kwargs)) + return True + + repo = BlockingClaimRepository() + service = McpTaskService( + repository=repo, + drivers=McpTaskDriverRegistry(), + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_polls=3, + ) + claim_factories = { + "poll": lambda: repo.claim_due_tasks(), + "cancel": lambda: repo.claim_cancel_requests(), + "notification": lambda: repo.claim_notification_work(), + } + releases = { + "poll": service._release_poll_after_cancellation, + "cancel": service._release_cancel_after_cancellation, + "notification": service._release_notification_after_cancellation, + } + + async def claim_once(): + return await service._claim_with_cancellation_release( + claim_factories[phase], + phase=phase, + action=f"{phase} claim", + release=releases[phase], + ) + + try: + assert await claim_once() == [] + await repo.claim_started.wait() + assert await claim_once() == [] + assert await claim_once() == [] + assert repo.claim_calls[phase] == 1 + assert list(service._claim_owners) == [phase] + + repo.claim_gate.set() + async with asyncio.timeout(0.2): + while not repo.released: + await asyncio.sleep(0) + + assert [(released_phase, task_id) for released_phase, task_id, _ in repo.released] == [(phase, "task-1")] + assert repo.released[0][2]["lease_owner"] == service._lease_owner + token_key = "notification_lease_token" if phase == "notification" else "lease_token" + assert repo.released[0][2][token_key] == ("notify-lease-token-1" if phase == "notification" else "lease-token-1") + + claimed = await claim_once() + assert [record["id"] for record in claimed] == ["task-1"] + assert repo.claim_calls[phase] == 2 + assert not service._claim_owners + finally: + repo.claim_gate.set() + owners = tuple(getattr(service, "_claim_owners", {}).values()) + handoffs = [owner.handoff_task for owner in owners if owner.handoff_task is not None] + if handoffs: + await asyncio.gather(*handoffs, return_exceptions=True) + + +@pytest.mark.asyncio +async def test_single_flight_claim_skip_logs_unresolved_owner(caplog): + service = McpTaskService( + repository=SimpleNamespace(), + drivers=McpTaskDriverRegistry(), + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_polls=3, + ) + claim_task = asyncio.get_running_loop().create_future() + service._claim_owners["poll"] = service_module._ClaimOwner(claim_task=claim_task) + + try: + with caplog.at_level(logging.WARNING): + result = await service._claim_with_cancellation_release( + lambda: pytest.fail("an unresolved owner must suppress a new claim"), + phase="poll", + action="poll claim", + release=AsyncMock(), + ) + + assert result == [] + assert "previous claim/handoff is still unresolved" in caplog.text + finally: + claim_task.cancel() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("phase", ["poll", "cancel", "notification"]) +async def test_single_flight_claim_releases_owner_before_stuck_release(phase, monkeypatch): + monkeypatch.setattr(service_module, "_CANCELLATION_DRAIN_TIMEOUT_SECONDS", 0.01) + + class BlockingClaimAndReleaseRepository: + def __init__(self): + self.claim_calls = {"poll": 0, "cancel": 0, "notification": 0} + self.claim_started = asyncio.Event() + self.claim_gate = asyncio.Event() + self.release_started = asyncio.Event() + self.release_gate = asyncio.Event() + + async def _claim(self, claim_phase): + self.claim_calls[claim_phase] += 1 + self.claim_started.set() + await self.claim_gate.wait() + return [{**_claimed_row(), "notification_status": "pending"}] + + async def claim_due_tasks(self, **_kwargs): + return await self._claim("poll") + + async def claim_cancel_requests(self, **_kwargs): + return await self._claim("cancel") + + async def claim_notification_work(self, **_kwargs): + return await self._claim("notification") + + async def _release(self): + self.release_started.set() + await self.release_gate.wait() + return True + + async def release_poll_claim_after_cancellation(self, _task_id, **_kwargs): + return await self._release() + + async def release_cancel_claim(self, _task_id, **_kwargs): + return await self._release() + + async def release_notification_claim(self, _task_id, **_kwargs): + return await self._release() + + repo = BlockingClaimAndReleaseRepository() + service = McpTaskService( + repository=repo, + drivers=McpTaskDriverRegistry(), + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_polls=3, + ) + claim_factories = { + "poll": lambda: repo.claim_due_tasks(), + "cancel": lambda: repo.claim_cancel_requests(), + "notification": lambda: repo.claim_notification_work(), + } + releases = { + "poll": service._release_poll_after_cancellation, + "cancel": service._release_cancel_after_cancellation, + "notification": service._release_notification_after_cancellation, + } + + async def claim_once(): + return await service._claim_with_cancellation_release( + claim_factories[phase], + phase=phase, + action=f"{phase} claim", + release=releases[phase], + ) + + try: + assert await claim_once() == [] + await repo.claim_started.wait() + + repo.claim_gate.set() + await repo.release_started.wait() + await asyncio.sleep(0.02) + + # The claim's durable outcome is now known, so the phase owner is released + # even though the release is still blocked. A new claim can proceed; the + # stuck release continues in the background (per-claim token fencing + # rejects it if it settles late). + assert list(service._claim_owners) == [] + assert service._compensation_tasks, "a stuck release must not be abandoned once the phase owner is released" + claimed = await claim_once() + assert [record["id"] for record in claimed] == ["task-1"] + assert repo.claim_calls[phase] == 2 + + repo.release_gate.set() + async with asyncio.timeout(0.2): + while service._claim_owners: + await asyncio.sleep(0) + finally: + repo.claim_gate.set() + repo.release_gate.set() + owners = tuple(getattr(service, "_claim_owners", {}).values()) + handoffs = [owner.handoff_task for owner in owners if owner.handoff_task is not None] + if handoffs: + await asyncio.gather(*handoffs, return_exceptions=True) + await _wait_for_compensation_tasks_to_clear(service) + + +@pytest.mark.asyncio +async def test_routine_cancel_release_preserves_existing_diagnostic(): + release = AsyncMock(return_value=True) + service = McpTaskService( + repository=SimpleNamespace(release_cancel_claim=release), + drivers=McpTaskDriverRegistry(), + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_polls=3, + ) + + await service._release_cancel_after_cancellation( + { + "id": "task-1", + "lease_token": "lease-1", + "last_cancel_error": "remote cancellation failed", + } + ) + + assert release.await_args.kwargs["error"] == "remote cancellation failed" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("notification_status", ["pending", "dispatched"]) +async def test_routine_notification_release_preserves_existing_diagnostic(notification_status): + release_claim = AsyncMock(return_value=True) + release_lease = AsyncMock(return_value=True) + service = McpTaskService( + repository=SimpleNamespace( + release_notification_claim=release_claim, + release_notification_lease=release_lease, + ), + drivers=McpTaskDriverRegistry(), + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_polls=3, + ) + + await service._release_notification_after_cancellation( + { + "id": "task-1", + "notification_lease_token": "notify-lease-1", + "notification_status": notification_status, + "notification_error": "notification launch failed", + } + ) + + release = release_lease if notification_status == "dispatched" else release_claim + assert release.await_args.kwargs["error"] == "notification launch failed" + assert release.await_args.kwargs.get("count_failure", False) is False + + +@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() + + +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 +@pytest.mark.parametrize("phase", ["poll", "cancel", "notification"]) +async def test_ordinary_batch_release_timeout_has_service_owned_strong_root(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() + await caller + + assert len(service._compensation_tasks) == 1 + release_task = next(iter(service._compensation_tasks)) + release_ref = weakref.ref(release_task) + + del release_task + del caller + gc.collect() + assert release_ref() is not None + + repo.release_gate.set() + await _wait_for_compensation_tasks_to_clear(service) + await asyncio.sleep(0) + assert not service._compensation_tasks + + +@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, + phase="poll", + 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 + + +@pytest.mark.asyncio +async def test_poll_retry_release_hang_does_not_block_run_once(monkeypatch): + import app.mcp_tasks.service as service_module + + monkeypatch.setattr(service_module, "_CANCELLATION_DRAIN_TIMEOUT_SECONDS", 0.01, raising=False) + + class HangingReleaseRepo(FakeRepository): + def __init__(self): + super().__init__([_claimed_row()]) + self.release_started = asyncio.Event() + self.finish_release = asyncio.Event() + self.release_calls = 0 + + async def release_claim(self, task_id, **kwargs): + self.release_calls += 1 + self.release_started.set() + await self.finish_release.wait() + return True + + repo = HangingReleaseRepo() + 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: + await asyncio.wait_for(service.run_once(now=datetime.now(UTC)), timeout=0.2) + finally: + repo.finish_release.set() + await asyncio.sleep(0) + + assert repo.release_calls == 1 + + +@pytest.mark.asyncio +async def test_notification_failure_release_hang_does_not_block(monkeypatch): + import app.mcp_tasks.service as service_module + + monkeypatch.setattr(service_module, "_CANCELLATION_DRAIN_TIMEOUT_SECONDS", 0.01, raising=False) + + class HangingNotificationReleaseRepo: + def __init__(self, records): + self.records = list(records) + self.release_started = asyncio.Event() + self.finish_release = asyncio.Event() + self.release_calls = 0 + + async def claim_notification_work(self, **_kwargs): + return list(self.records) + + async def release_notification_lease(self, task_id, **kwargs): + self.release_calls += 1 + self.release_started.set() + await self.finish_release.wait() + return True + + record = { + **_claimed_row(), + "notification_status": "dispatched", + "notification_run_id": "run-broken", + "dispatch_version": 3, + } + repo = HangingNotificationReleaseRepo([record]) + 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")), + ) + try: + await asyncio.wait_for(service._run_notifications(now=datetime.now(UTC)), timeout=0.2) + finally: + repo.finish_release.set() + await asyncio.sleep(0) + + assert repo.release_calls == 1 + + +@pytest.mark.asyncio +async def test_claim_hang_does_not_block_run_once(monkeypatch): + import app.mcp_tasks.service as service_module + + monkeypatch.setattr(service_module, "_CANCELLATION_DRAIN_TIMEOUT_SECONDS", 0.01, raising=False) + + class HangingClaimRepo(FakeRepository): + def __init__(self): + super().__init__() + self.claim_started = asyncio.Event() + self.finish_claim = asyncio.Event() + + async def claim_due_tasks(self, **_kwargs): + self.claim_started.set() + await self.finish_claim.wait() + return [] + + repo = HangingClaimRepo() + service = McpTaskService( + repository=repo, + drivers=McpTaskDriverRegistry(), + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_polls=3, + ) + try: + await asyncio.wait_for(service.run_once(now=datetime.now(UTC)), timeout=0.2) + finally: + repo.finish_claim.set() + await asyncio.sleep(0) + + assert repo.claim_started.is_set() + + +class FailingReleaseRepository(FakeRepository): + def __init__(self): + super().__init__([_claimed_row()]) + self.release_called = False + + async def release_claim(self, task_id, **kwargs): + self.release_called = True + raise RuntimeError("release db down") + + +def test_failing_release_does_not_leak_unretrieved_shield_exception(tmp_path): + import os + import subprocess + import sys + import textwrap + + backend_dir = os.path.dirname(os.path.dirname(__file__)) + probe = textwrap.dedent( + """ + import asyncio, gc + from datetime import UTC, datetime + from app.mcp_tasks.service import McpTaskService + from deerflow.mcp.tasks import McpTaskDriverRegistry + + class Repo: + def __init__(self): + self.row = { + "id": "task-1", "user_id": "u", "thread_id": "t", "run_id": None, + "tool_call_id": "c", "server_name": "s", "driver_name": "fake", + "remote_task_id": "r", "task_name": "n", "status": "working", + "driver_data": {}, "lease_owner": "o", "lease_token": "tok-1", + "consecutive_poll_error_count": 0, + } + self.claimed = False + + async def claim_due_tasks(self, **_kw): + if self.claimed: + return [] + self.claimed = True + return [dict(self.row)] + + async def release_claim(self, task_id, **kw): + raise RuntimeError("release db down") + + class Driver: + async def get_status(self, task): + raise RuntimeError("poll down") + + async def scenario(): + repo = Repo() + drivers = McpTaskDriverRegistry() + drivers.register("fake", Driver()) + service = McpTaskService( + repository=repo, drivers=drivers, + poll_interval_seconds=5, lease_seconds=120, max_concurrent_polls=3, + ) + await service.run_once(now=datetime.now(UTC)) + del service, repo, drivers + for _ in range(20): + gc.collect() + await asyncio.sleep(0) + + captured = [] + loop = asyncio.new_event_loop() + try: + loop.set_exception_handler(lambda _loop, context: captured.append(context.get("message", ""))) + loop.run_until_complete(scenario()) + finally: + loop.close() + if any("Future exception was never retrieved" in message for message in captured): + print("UNRETRIEVED") + """ + ) + env = {**os.environ, "PYTHONPATH": backend_dir} + result = subprocess.run( + [sys.executable, "-c", probe], + capture_output=True, + text=True, + cwd=backend_dir, + env=env, + ) + assert "UNRETRIEVED" not in result.stdout, result.stderr + + +@pytest.mark.asyncio +async def test_poll_cancellation_releases_only_the_current_poll_lease(): + repo = PollPersistenceRepo() + driver = HangingDriver() + drivers = McpTaskDriverRegistry() + drivers.register("fake", driver) + now = datetime.now(UTC) + service = McpTaskService( + repository=repo, + drivers=drivers, + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_polls=3, + ) + + task = asyncio.create_task( + service._run_claimed_batch( + [_claimed_row()], + operation=lambda record: service._poll_one_claimed(record, now=now), + release=service._release_poll_after_cancellation, + action="poll", + ) + ) + await driver.started.wait() + task.cancel() + + with pytest.raises(asyncio.CancelledError): + await task + + # The cancellation must travel through the batch ownership -> poll + # cancellation release path, not the ordinary poll-error release (which is + # recorded separately in ``released``). + assert repo.cancelled_releases == [("task-1", {"lease_owner": service._lease_owner, "lease_token": "lease-token-1"})] + 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) + now = datetime.now(UTC) + service = McpTaskService( + repository=repo, + drivers=drivers, + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_polls=3, + ) + + task = asyncio.create_task( + service._run_claimed_batch( + [_claimed_row()], + operation=lambda record: service._poll_one_claimed(record, now=now), + release=service._release_poll_after_cancellation, + action="poll", + ) + ) + await driver.started.wait() + task.cancel("shutdown") + + with caplog.at_level(logging.ERROR), pytest.raises(asyncio.CancelledError) as caught: + await task + + # The original cancellation identity/args must survive the release failure. + assert caught.value.args == ("shutdown",) + assert repo.cancelled_releases == [("task-1", {"lease_owner": service._lease_owner, "lease_token": "lease-token-1"})] + assert repo.released == [] + assert "poll release unavailable" in caplog.text + + +@pytest.mark.asyncio +async def test_cancel_batch_releases_claim_when_cancelled(): + """A cancel that lands mid-batch-cancel must still release the cancel claim.""" + repo = CancellationBlockingApplyRepo() + driver = FakeDriver() + 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() + + task = asyncio.create_task( + service._run_claimed_batch( + [record], + operation=service._cancel_one_claimed, + release=service._release_cancel_after_cancellation, + action="cancel", + ) + ) + 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"] + + +@pytest.mark.asyncio +async def test_cancel_batch_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._run_claimed_batch( + [_claimed_row()], + operation=service._cancel_one_claimed, + release=service._release_cancel_after_cancellation, + action="cancel", + ) + ) + await repo.apply_started.wait() + task.cancel("shutdown") + with caplog.at_level(logging.ERROR), pytest.raises(asyncio.CancelledError) as caught: + await task + assert caught.value.args == ("shutdown",) + assert repo.release_cancel_calls + assert "release unavailable" in caplog.text + + +@pytest.mark.asyncio +async def test_cancel_batch_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._run_claimed_batch( + [_claimed_row()], + operation=service._cancel_one_claimed, + release=service._release_cancel_after_cancellation, + action="cancel", + ) + ) + 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_failure_release_is_not_restarted_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._run_claimed_batch( + [_claimed_row()], + operation=service._cancel_one_claimed, + release=service._release_cancel_after_cancellation, + action="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 + assert len(repo.release_cancel_calls) == 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] + now = datetime.now(UTC) + + task = asyncio.create_task( + service._run_claimed_batch( + [record], + operation=lambda r: service._notify_one_claimed(r, now=now), + release=service._release_notification_after_cancellation, + action="notification", + ) + ) + await repo.mark_started.wait() + task.cancel("shutdown") + + with caplog.at_level(logging.ERROR), pytest.raises(asyncio.CancelledError) as caught: + await task + assert caught.value.args == ("shutdown",) + assert repo.release_calls + assert "notification release unavailable" in caplog.text + + +@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), + 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(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"}, + } + now = datetime.now(UTC) + + task = asyncio.create_task( + service._run_claimed_batch( + [record], + operation=lambda r: service._notify_one_claimed(r, now=now), + release=service._release_notification_after_cancellation, + action="notification", + ) + ) + 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, + } + now = datetime.now(UTC) + + task = asyncio.create_task( + service._run_claimed_batch( + [record], + operation=lambda r: service._notify_one_claimed(r, now=now), + release=service._release_notification_after_cancellation, + action="notification", + ) + ) + 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() + + +@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) + now = datetime.now(UTC) + 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._run_claimed_batch( + [_claimed_row()], + operation=lambda r: service._poll_one_claimed(r, now=now), + release=service._release_poll_after_cancellation, + action="poll", + ) + ) + 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"] + # Both the batch-cancellation handoff and the individual release task are + # transferred to service-owned background ownership once they exceed the + # drain deadline. + assert service._compensation_tasks + + 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) + now = datetime.now(UTC) + 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._run_claimed_batch( + [_claimed_row()], + operation=lambda r: service._poll_one_claimed(r, now=now), + release=service._release_poll_after_cancellation, + action="poll", + ) + ) + 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 service._compensation_tasks + + 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 any("release remained unavailable" in failure.getMessage() for failure in failures) + assert not service._compensation_tasks diff --git a/backend/tests/test_migration_0025_repair_run_change_seq.py b/backend/tests/test_migration_0025_repair_run_change_seq.py index 28c125a81..7cff68861 100644 --- a/backend/tests/test_migration_0025_repair_run_change_seq.py +++ b/backend/tests/test_migration_0025_repair_run_change_seq.py @@ -15,10 +15,11 @@ import sqlite3 import pytest import sqlalchemy as sa from alembic import command +from alembic.script import ScriptDirectory import deerflow.persistence.models # noqa: F401 -- registers ORM models from deerflow.persistence.base import Base -from deerflow.persistence.bootstrap import _get_alembic_config, _get_head_revision +from deerflow.persistence.bootstrap import _MIGRATIONS_DIR, _get_alembic_config, _get_head_revision from deerflow.persistence.engine import close_engine, get_session_factory, init_engine from deerflow.persistence.run import RunRepository @@ -66,8 +67,12 @@ def _table_and_column_state(db_path) -> tuple[bool, bool, set[str], str | None]: return "run_change_clock" in tables, "change_seq" in run_columns, run_indexes, version_row[0] if version_row else None -async def test_0025_is_the_chain_head(): - assert _get_head_revision() == REVISION +async def test_0025_chains_into_the_single_head(): + script = ScriptDirectory(str(_MIGRATIONS_DIR)) + assert len(script.get_heads()) == 1 + # Later migrations may advance the head without removing this revision. + assert REVISION in {revision.revision for revision in script.walk_revisions()} + assert script.get_revision(REVISION).down_revision == PREVIOUS async def test_0025_repairs_schema_skipped_by_the_0023_insertion(tmp_path): diff --git a/backend/tests/test_migration_0026_mcp_task_lease_tokens.py b/backend/tests/test_migration_0026_mcp_task_lease_tokens.py new file mode 100644 index 000000000..0f2220568 --- /dev/null +++ b/backend/tests/test_migration_0026_mcp_task_lease_tokens.py @@ -0,0 +1,51 @@ +"""Migration tests for 0026_mcp_task_lease_tokens. + +Adds the nullable per-claim token columns ``McpTaskRepository`` uses to fence +poll, cancel, and notification mutations to the exact claim generation. This +file owns the chain-head pin, moved on from +``test_migration_0025_repair_run_change_seq`` with this revision. +""" + +from __future__ import annotations + +import asyncio + +import pytest +import sqlalchemy as sa +from alembic import command +from sqlalchemy.ext.asyncio import create_async_engine + +from deerflow.persistence import bootstrap + +pytestmark = pytest.mark.asyncio + +REVISION = "0026_mcp_task_lease_tokens" +PREVIOUS = "0025_repair_run_change_seq" +TOKEN_COLUMNS = {"lease_token", "notification_lease_token"} + + +async def test_0026_is_the_chain_head(): + assert bootstrap._get_head_revision() == REVISION + + +async def test_0026_adds_nullable_claim_tokens_and_downgrades(tmp_path): + engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'lease-tokens.db'}") + cfg = bootstrap._get_alembic_config(engine) + + async def token_columns() -> dict[str, bool]: + async with engine.connect() as conn: + columns = await conn.run_sync(lambda sync: sa.inspect(sync).get_columns("mcp_tasks")) + return {column["name"]: column["nullable"] for column in columns if column["name"] in TOKEN_COLUMNS} + + try: + await asyncio.to_thread(bootstrap._upgrade, cfg, PREVIOUS) + assert await token_columns() == {} + + await asyncio.to_thread(bootstrap._upgrade, cfg, "head") + # Nullable so rows written before this revision stay valid. + assert await token_columns() == dict.fromkeys(TOKEN_COLUMNS, True) + + await asyncio.to_thread(command.downgrade, cfg, PREVIOUS) + assert await token_columns() == {} + finally: + await engine.dispose() diff --git a/backend/tests/test_runtime_cancellation.py b/backend/tests/test_runtime_cancellation.py new file mode 100644 index 000000000..929540cfa --- /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