deer-flow/backend/app/mcp_tasks/service.py

from __future__ import annotations

import asyncio
import json
import logging
import socket
import uuid
from collections.abc import Awaitable, Callable
from contextvars import ContextVar
from dataclasses import dataclass, replace
from datetime import UTC, datetime, timedelta
from typing import Any

from app.mcp_tasks.errors import PermanentNotificationError
from deerflow.constants import (
    MCP_TASK_POLL_AFTER_MAX_SECONDS,
    MCP_TASK_REMOTE_ID_MAX_LENGTH,
    MCP_TASK_RESULT_ARTIFACT_MAX_BYTES,
)
from deerflow.mcp.tasks import (
    McpTaskDriverRegistry,
    McpTaskProtocolError,
    TaskReference,
    TaskSnapshot,
    TaskStatus,
    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

logger = logging.getLogger(__name__)

_MAX_PERSISTED_ERROR_CHARS = 4_000
_MAX_INPUT_REQUIRED_BYTES = 65_536
_MAX_NOTIFICATION_ATTEMPTS = 5
_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:
    if error is None:
        return 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."""

    def __init__(
        self,
        *,
        repository,
        drivers: McpTaskDriverRegistry,
        poll_interval_seconds: int,
        lease_seconds: int,
        max_concurrent_polls: int,
        max_poll_backoff_seconds: int = 300,
        input_required_poll_interval_seconds: int = 60,
        tracking_degraded_after_errors: int = 3,
        max_result_bytes: int = 65_536,
        result_preview_max_chars: int = 2_000,
        launch_notification: Callable[..., Awaitable[dict[str, Any]]] | None = None,
        get_run: Callable[..., Awaitable[Any | None]] | None = None,
    ) -> None:
        self._repository = repository
        self._drivers = drivers
        self._poll_interval_seconds = poll_interval_seconds
        self._lease_seconds = lease_seconds
        self._max_concurrent_polls = max_concurrent_polls
        self._max_poll_backoff_seconds = max_poll_backoff_seconds
        self._input_required_poll_interval_seconds = input_required_poll_interval_seconds
        self._tracking_degraded_after_errors = tracking_degraded_after_errors
        self._max_result_bytes = max_result_bytes
        self._result_preview_max_chars = result_preview_max_chars
        self._launch_notification = launch_notification
        self._get_run = get_run
        self._lease_owner = f"{socket.gethostname()}:{uuid.uuid4().hex}"
        self._task: asyncio.Task[None] | None = None
        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
    def drivers(self) -> McpTaskDriverRegistry:
        return self._drivers

    @property
    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,
        *,
        driver_name: str,
        request: TaskSubmitRequest,
        now: datetime | None = None,
    ) -> dict:
        """Submit through one driver and persist the remote handle before returning."""
        driver = self._drivers.get(driver_name)
        if driver is None:
            raise LookupError(f"No MCP task driver registered as {driver_name!r}")

        submitted_at = now or datetime.now(UTC)
        local_task_id = request.local_task_id or f"mcp-task-{uuid.uuid4().hex}"
        driver_request = replace(request, local_task_id=local_task_id)
        submission = await driver.submit(driver_request)
        driver_data = {**request.driver_data, **submission.driver_data}
        task_reference = TaskReference(
            local_task_id=local_task_id,
            user_id=request.user_id,
            thread_id=request.thread_id,
            server_name=request.server_name,
            remote_task_id=submission.remote_task_id,
            thread_incarnation=request.thread_incarnation,
            driver_data=driver_data,
        )
        try:
            if len(submission.remote_task_id) > MCP_TASK_REMOTE_ID_MAX_LENGTH:
                raise McpTaskProtocolError(f"MCP task remote_task_id must not exceed {MCP_TASK_REMOTE_ID_MAX_LENGTH} characters")
            snapshot = self._normalize_snapshot(submission.snapshot)
            next_poll_at = self._next_poll_at(snapshot, now=submitted_at)
            return await self._repository.create(
                task_id=local_task_id,
                user_id=request.user_id,
                thread_id=request.thread_id,
                expected_thread_incarnation=request.thread_incarnation,
                run_id=request.run_id,
                tool_call_id=request.tool_call_id,
                server_name=request.server_name,
                driver_name=driver_name,
                remote_task_id=submission.remote_task_id,
                task_name=request.task_name,
                status=snapshot.status.value,
                result=snapshot.result,
                result_preview=snapshot.result_preview,
                result_truncated=snapshot.result_truncated,
                result_artifact=snapshot.result_artifact,
                error=snapshot.error,
                input_required=snapshot.input_required,
                next_poll_at=next_poll_at,
                driver_data=driver_data,
            )
        except DuplicateMcpRemoteTaskError:
            # This handle already has a durable owner. Cancelling it as
            # compensation would terminate the pre-existing tracked task.
            raise
        except asyncio.CancelledError:
            # Cancellation can race with a successful database commit. If it
            # did, the durable row will converge to cancelled on its next poll;
            # compensating is safer than leaving a live remote task untracked.
            await self._cancel_untracked_task(
                driver=driver,
                task_reference=task_reference,
                driver_name=driver_name,
                reason="caller cancellation during local persistence",
            )
            raise
        except Exception:
            await self._cancel_untracked_task(
                driver=driver,
                task_reference=task_reference,
                driver_name=driver_name,
                reason="local submission finalization failure",
            )
            raise

    async def _cancel_untracked_task(
        self,
        *,
        driver,
        task_reference: TaskReference,
        driver_name: str,
        reason: str,
    ) -> None:
        compensation = asyncio.create_task(
            driver.cancel(task_reference),
            name=f"mcp-submit-compensation-{task_reference.local_task_id}",
        )
        self._compensation_tasks.add(compensation)

        def finalize(task: asyncio.Future[Any]) -> None:
            self._compensation_tasks.discard(task)
            error = _consume_task_error(task)
            if error is None:
                return
            logger.error(
                "Failed to cancel untracked MCP task after %s (task_id=%s, driver=%s, remote_task_id=%s)",
                reason,
                task_reference.local_task_id,
                driver_name,
                task_reference.remote_task_id,
                exc_info=(type(error), error, error.__traceback__),
            )

        compensation.add_done_callback(finalize)
        loop = asyncio.get_running_loop()
        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:
                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:
                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._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:
            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 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))

    async def list_tasks(
        self,
        *,
        thread_id: str,
        user_id: str,
        thread_incarnation: str | None,
        limit: int = 50,
        active_only: bool = False,
    ) -> list[dict[str, Any]]:
        return await self._repository.list_by_thread(
            thread_id,
            user_id=user_id,
            thread_incarnation=thread_incarnation,
            limit=limit,
            active_only=active_only,
        )

    async def cancel_task(
        self,
        *,
        task_id: str,
        thread_id: str,
        user_id: str,
        thread_incarnation: str | None,
        now: datetime | None = None,
    ) -> dict[str, Any] | None:
        return await self._repository.request_cancel(
            task_id,
            user_id=user_id,
            thread_id=thread_id,
            thread_incarnation=thread_incarnation,
            requested_at=now or datetime.now(UTC),
        )

    async def cancel_matching_task(
        self,
        *,
        thread_id: str,
        user_id: str,
        thread_incarnation: str | None,
        task: str | None = None,
    ) -> dict[str, Any]:
        active = await self.list_tasks(
            thread_id=thread_id,
            user_id=user_id,
            thread_incarnation=thread_incarnation,
            active_only=True,
        )
        if task:
            normalized = task.casefold().strip()
            matches = [item for item in active if item["id"] == task or str(item.get("task_name") or "").casefold() == normalized]
        else:
            matches = active
        if not matches:
            raise LookupError("No active background task matches this request")
        if len(matches) > 1:
            names = ", ".join(str(item.get("task_name") or item["id"]) for item in matches[:5])
            raise ValueError(f"More than one active background task matches; specify one task name: {names}")
        result = await self.cancel_task(
            task_id=matches[0]["id"],
            thread_id=thread_id,
            user_id=user_id,
            thread_incarnation=thread_incarnation,
        )
        if result is None:
            raise LookupError("The selected background task no longer exists")
        return result

    async def _run_cancellations(self, *, now: datetime) -> None:
        claim = getattr(self._repository, "claim_cancel_requests", None)
        if claim is None:
            return
        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:
            states, results = await self._run_claimed_batch(
                records,
                operation=self._cancel_one_claimed,
                release=self._release_cancel_after_cancellation,
                action="cancel",
            )
            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_claimed(self, record: dict[str, Any]) -> None:
        driver_name = str(record.get("driver_name") or "")
        driver = self._drivers.get(driver_name)
        try:
            if driver is None:
                raise LookupError(f"No MCP task driver registered as {driver_name!r}")
            snapshot = self._normalize_snapshot(await driver.cancel(TaskReference.from_record(record)))
            if snapshot.status not in (TaskStatus.COMPLETED, TaskStatus.FAILED, TaskStatus.CANCELLED):
                raise McpTaskProtocolError("MCP task cancellation must return a terminal status")
            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,
                result_truncated=snapshot.result_truncated,
                result_artifact=snapshot.result_artifact,
                error=snapshot.error,
                input_required=snapshot.input_required,
                completed_at=datetime.now(UTC),
            )
        except Exception as exc:  # noqa: BLE001 - remote cancellation is retryable
            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)
            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._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:
            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 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__),
                )
                await self._release_notification_failure(record, now=now, error=error)

    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))
        if notification_attempts >= _MAX_NOTIFICATION_ATTEMPTS:
            previous_error = record.get("notification_error") or "delivery failed"
            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,
                now=now,
            )
            return

        if record.get("notification_status") == "dispatched":
            run = await self._get_run(record.get("notification_run_id"), user_id=record["user_id"])
            status = getattr(run, "status", None)
            if run is None:
                run_id = record.get("notification_run_id")
                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)),
                    error=_bound_error(f"Notification run {run_id!r} was not found"),
                    now=now,
                )
            elif status == RunStatus.success:
                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,
                    error=None,
                    now=now,
                )
            elif status in {RunStatus.error, RunStatus.timeout, RunStatus.interrupted}:
                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)),
                    error=_bound_error(getattr(run, "error", None) or f"Notification run ended with {status}"),
                    now=now,
                )
            else:
                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,
                )
            return

        source_run = await self._get_run(record.get("run_id"), user_id=record["user_id"]) if record.get("run_id") else None
        try:
            result = await self._launch_notification(
                thread_id=record["thread_id"],
                assistant_id=getattr(source_run, "assistant_id", None),
                owner_user_id=record["user_id"],
                task_id=task_id,
                dispatch_version=dispatch_version,
                dispatch_attempt=int(record.get("dispatch_attempt") or 0),
                event=dict(record.get("dispatch_event") or {}),
            )
        except PermanentNotificationError as exc:
            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,
                now=now,
            )
            return
        except ConflictError as exc:
            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
            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,
        )

    def _notification_retry_seconds(self, record: dict[str, Any]) -> int:
        failures = max(0, int(record.get("notification_attempt_count") or 0))
        return min(
            self._poll_interval_seconds * (2 ** min(failures, 16)),
            self._max_poll_backoff_seconds,
        )

    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_ordinary_batch_record(
                record,
                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

        try:
            snapshot = self._normalize_snapshot(await driver.get_status(TaskReference.from_record(record)))
        except McpTaskProtocolError as exc:
            logger.error(
                "MCP task status contract failed permanently (task_id=%s, driver=%s): %s",
                record.get("id"),
                driver_name,
                exc,
            )
            await self._apply_snapshot(
                record,
                TaskSnapshot(status=TaskStatus.FAILED, error=_bound_error(str(exc))),
                polled_at=datetime.now(UTC),
            )
            return
        except Exception as exc:  # noqa: BLE001 - driver boundary; retry on the next poll
            polled_at = datetime.now(UTC)
            logger.warning(
                "MCP task status poll failed (task_id=%s, driver=%s); retrying",
                record.get("id"),
                driver_name,
                exc_info=True,
            )
            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)
        await self._apply_snapshot(record, snapshot, polled_at=polled_at)

    async def _apply_snapshot(
        self,
        record: dict,
        snapshot: TaskSnapshot,
        *,
        polled_at: datetime,
    ) -> None:
        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,
            result_truncated=snapshot.result_truncated,
            result_artifact=snapshot.result_artifact,
            error=snapshot.error,
            input_required=snapshot.input_required,
            next_poll_at=self._next_poll_at(snapshot, now=polled_at),
            polled_at=polled_at,
        )
        if not applied:
            logger.info(
                "Discarded MCP task poll result after lease ownership changed or expired (task_id=%s)",
                record.get("id"),
            )

    def _next_poll_at(self, snapshot: TaskSnapshot, *, now: datetime) -> datetime | None:
        if not snapshot.is_pollable:
            return None
        interval = snapshot.poll_after_seconds or self._poll_interval_seconds
        if snapshot.status == TaskStatus.INPUT_REQUIRED:
            interval = max(interval, self._input_required_poll_interval_seconds)
        interval = min(interval, MCP_TASK_POLL_AFTER_MAX_SECONDS)
        return now + timedelta(seconds=interval)

    async def _release_after_error(self, record: dict, *, now: datetime, error: str) -> None:
        consecutive_errors = max(0, int(record.get("consecutive_poll_error_count") or 0))
        retry_seconds = min(
            self._poll_interval_seconds * (2 ** min(consecutive_errors, 16)),
            self._max_poll_backoff_seconds,
        )
        bounded_error = _bound_error(error)
        assert bounded_error is not None
        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,
        )

    def _normalize_snapshot(self, snapshot: TaskSnapshot) -> TaskSnapshot:
        """Bound remote payloads without ever storing truncated JSON."""
        snapshot = replace(snapshot, error=_bound_error(snapshot.error))
        if snapshot.result_artifact is not None:
            encoded_artifact = self._encode_json_payload(
                snapshot.result_artifact,
                field_name="result_artifact",
            )
            if len(encoded_artifact) > MCP_TASK_RESULT_ARTIFACT_MAX_BYTES:
                raise McpTaskProtocolError(f"MCP task result_artifact payload exceeds the {MCP_TASK_RESULT_ARTIFACT_MAX_BYTES}-byte limit")
        if snapshot.input_required is not None:
            encoded_input = self._encode_json_payload(
                snapshot.input_required,
                field_name="input_required",
            )
            if len(encoded_input) > _MAX_INPUT_REQUIRED_BYTES:
                raise McpTaskProtocolError(f"MCP task input_required payload exceeds the {_MAX_INPUT_REQUIRED_BYTES}-byte limit")
        if snapshot.result is None:
            return snapshot
        encoded = self._encode_json_payload(snapshot.result, field_name="result")
        if len(encoded) <= self._max_result_bytes:
            return snapshot

        if isinstance(snapshot.result, str):
            preview_source = snapshot.result
        else:
            preview_source = encoded.decode("utf-8", errors="replace")
        return replace(
            snapshot,
            result=None,
            result_preview=preview_source[: self._result_preview_max_chars],
            result_truncated=True,
        )

    @staticmethod
    def _encode_json_payload(value, *, field_name: str) -> bytes:
        try:
            return json.dumps(
                value,
                ensure_ascii=False,
                separators=(",", ":"),
                allow_nan=False,
            ).encode("utf-8")
        except (TypeError, ValueError) as exc:
            raise McpTaskProtocolError(f"MCP task {field_name} is not valid JSON: {exc}") from exc

    async def start(self) -> None:
        if self._task is not None:
            return
        self._stop.clear()
        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()
        deadline = self._stop_deadline
        assert deadline is not None
        try:
            done, _ = await asyncio.wait(
                {task},
                timeout=max(0.0, deadline - loop.time()),
            )
        except asyncio.CancelledError:
            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():
            try:
                # The first pass runs immediately. Expired leases therefore
                # recover at startup without a separate destructive sweep.
                await self.run_once(now=datetime.now(UTC))
            except Exception:
                logger.exception("MCP task poll failed; retrying next interval")
            try:
                await asyncio.wait_for(
                    self._stop.wait(),
                    timeout=self._poll_interval_seconds,
                )
            except TimeoutError:
                continue