from __future__ import annotations import hashlib import json from datetime import UTC, datetime, timedelta from typing import Any from sqlalchemy import or_, select, update from sqlalchemy.exc import IntegrityError from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker from deerflow.mcp.tasks import ATTENTION_TASK_STATUSES, POLLABLE_TASK_STATUSES, TERMINAL_TASK_STATUSES from deerflow.persistence.mcp_tasks.model import McpTaskRow from deerflow.persistence.thread_meta.model import ThreadMetaRow from deerflow.utils.time import coerce_iso _POLLABLE_STATUS_VALUES = tuple(status.value for status in POLLABLE_TASK_STATUSES) _ATTENTION_STATUS_VALUES = frozenset(status.value for status in ATTENTION_TASK_STATUSES) _TERMINAL_STATUS_VALUES = frozenset(status.value for status in TERMINAL_TASK_STATUSES) _TIMESTAMP_FIELDS = ( "next_poll_at", "last_polled_at", "lease_expires_at", "notification_lease_expires_at", "next_notification_at", "cancel_requested_at", "next_cancel_at", "completed_at", "created_at", "updated_at", ) _INFLIGHT_NOTIFICATION_STATUSES = frozenset({"claimed", "dispatched", "retry"}) 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 return { "task_id": row.id, "task_name": row.task_name, "status": row.status, "result": row.result, "result_preview": row.result_preview, "result_truncated": bool(row.result_truncated), "result_artifact": row.result_artifact, "error": row.error, "input_required": row.input_required, "tracking_degraded": tracking_degraded, "last_poll_error": row.last_poll_error if tracking_degraded else None, } def _event_fingerprint(event: dict[str, Any]) -> str: encoded = json.dumps(event, ensure_ascii=False, sort_keys=True, separators=(",", ":"), default=str).encode("utf-8") return hashlib.sha256(encoded).hexdigest() def _record_event_if_changed(row: McpTaskRow, *, tracking_degraded: bool, now: datetime) -> bool: event = _notification_event(row, tracking_degraded=tracking_degraded) if event is None: return False fingerprint = _event_fingerprint(event) if fingerprint == row.event_fingerprint: return False row.event_fingerprint = fingerprint row.event_version = int(row.event_version or 0) + 1 if row.notification_status not in _INFLIGHT_NOTIFICATION_STATUSES: row.notification_status = "pending" row.next_notification_at = now row.notification_error = None row.notification_attempt_count = 0 row.dispatch_version = None row.dispatch_attempt = 0 row.dispatch_event = None row.notification_run_id = None return True class DuplicateMcpRemoteTaskError(RuntimeError): """The current user already tracks this server's remote task handle.""" def _is_remote_task_unique_conflict(exc: IntegrityError) -> bool: original = exc.orig diagnostic = getattr(original, "diag", None) if getattr(diagnostic, "constraint_name", None) == "uq_mcp_tasks_user_server_remote": return True message = str(original) return "uq_mcp_tasks_user_server_remote" in message or "mcp_tasks.user_id, mcp_tasks.server_name, mcp_tasks.remote_task_id" in message class McpTaskRepository: """Durable source of truth for long-running MCP task lifecycle state.""" def __init__(self, session_factory: async_sessionmaker[AsyncSession]) -> None: self._sf = session_factory @staticmethod def _row_to_dict(row: McpTaskRow) -> dict[str, Any]: data = row.to_dict() data.pop("thread_incarnation", None) for key in _TIMESTAMP_FIELDS: if data.get(key) is not None: data[key] = coerce_iso(data[key]) return data async def create( self, *, task_id: str, user_id: str, thread_id: str, run_id: str | None, tool_call_id: str | None, server_name: str, driver_name: str, remote_task_id: str, task_name: str, status: str, result: Any | None, result_preview: str | None, result_truncated: bool, result_artifact: dict[str, str] | None, error: str | None, input_required: dict[str, Any] | None, next_poll_at: datetime | None, driver_data: dict[str, Any] | None = None, ) -> dict[str, Any]: now = datetime.now(UTC) row = McpTaskRow( id=task_id, user_id=user_id, thread_id=thread_id, run_id=run_id, tool_call_id=tool_call_id, server_name=server_name, driver_name=driver_name, remote_task_id=remote_task_id, task_name=task_name, status=status, result=result, result_preview=result_preview, result_truncated=result_truncated, result_artifact=result_artifact, error=error, input_required=input_required, driver_data=dict(driver_data or {}), notification_status="none", next_poll_at=next_poll_at, completed_at=now if status in _TERMINAL_STATUS_VALUES else None, created_at=now, updated_at=now, ) _record_event_if_changed(row, tracking_degraded=False, now=now) async with self._sf() as session: matching_thread = select(ThreadMetaRow.incarnation).where( ThreadMetaRow.thread_id == thread_id, or_(ThreadMetaRow.user_id == user_id, ThreadMetaRow.user_id.is_(None)), ) if session.get_bind().dialect.name == "sqlite": # Keep lookup and write in one SQLite statement. A preliminary # read would leave a delete/recreate window before the INSERT. row.thread_incarnation = matching_thread.scalar_subquery() else: # FOR SHARE also conflicts with the FOR NO KEY UPDATE lock taken # by an older writer's plain owner UPDATE. KEY SHARE would not, # leaving a mixed-version ownership race before this INSERT. row.thread_incarnation = (await session.execute(matching_thread.with_for_update(read=True))).scalar_one_or_none() session.add(row) try: await session.commit() except IntegrityError as exc: await session.rollback() if _is_remote_task_unique_conflict(exc): raise DuplicateMcpRemoteTaskError(f"Remote MCP task {remote_task_id!r} is already tracked for server {server_name!r} by this user") from exc raise await session.refresh(row) return self._row_to_dict(row) async def get(self, task_id: str, *, user_id: str) -> dict[str, Any] | None: async with self._sf() as session: row = await session.get(McpTaskRow, task_id) if row is None or row.user_id != user_id: return None return self._row_to_dict(row) async def list_by_thread( self, thread_id: str, *, user_id: str, limit: int = 50, active_only: bool = False, ) -> list[dict[str, Any]]: stmt = select(McpTaskRow).where( McpTaskRow.thread_id == thread_id, McpTaskRow.user_id == user_id, ) if active_only: stmt = stmt.where(McpTaskRow.status.in_(_POLLABLE_STATUS_VALUES)) stmt = stmt.order_by(McpTaskRow.created_at.desc(), McpTaskRow.id.desc()).limit(limit) async with self._sf() as session: result = await session.execute(stmt) return [self._row_to_dict(row) for row in result.scalars()] async def claim_due_tasks( self, *, now: datetime, lease_owner: str, lease_seconds: int, limit: int, ) -> list[dict[str, Any]]: lease_expires_at = now + timedelta(seconds=lease_seconds) stmt = ( select(McpTaskRow) .where( McpTaskRow.status.in_(_POLLABLE_STATUS_VALUES), McpTaskRow.cancel_requested_at.is_(None), McpTaskRow.next_poll_at.is_not(None), McpTaskRow.next_poll_at <= now, or_( McpTaskRow.lease_expires_at.is_(None), McpTaskRow.lease_expires_at < now, ), ) .order_by(McpTaskRow.next_poll_at.asc(), McpTaskRow.id.asc()) .limit(limit) .with_for_update(skip_locked=True) ) async with self._sf() as session: result = await session.execute(stmt) rows = list(result.scalars()) for row in rows: row.lease_owner = lease_owner row.lease_expires_at = lease_expires_at row.poll_attempt_count += 1 row.updated_at = now await session.commit() return [self._row_to_dict(row) for row in rows] async def apply_snapshot( self, task_id: str, *, lease_owner: str, status: str, result: Any | None, result_preview: str | None, result_truncated: bool, result_artifact: dict[str, str] | None, error: str | None, input_required: dict[str, Any] | None, next_poll_at: datetime | None, polled_at: datetime, ) -> bool: async with self._sf() as session: stmt = ( select(McpTaskRow) .where( McpTaskRow.id == task_id, McpTaskRow.lease_owner == lease_owner, McpTaskRow.lease_expires_at >= polled_at, McpTaskRow.status.not_in(_TERMINAL_STATUS_VALUES), McpTaskRow.cancel_requested_at.is_(None), ) .with_for_update() ) row = (await session.execute(stmt)).scalar_one_or_none() if row is None: 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) await session.commit() return True async def release_claim( self, task_id: str, *, lease_owner: 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, ) await session.commit() return True async def request_cancel( self, task_id: str, *, user_id: str, thread_id: str, requested_at: datetime, ) -> dict[str, Any] | None: """Persist a user-scoped cancellation request without exposing the remote id.""" async with self._sf() as session: stmt = ( select(McpTaskRow) .where( McpTaskRow.id == task_id, McpTaskRow.user_id == user_id, McpTaskRow.thread_id == thread_id, ) .with_for_update() ) row = (await session.execute(stmt)).scalar_one_or_none() if row is None: return None if row.status not in _TERMINAL_STATUS_VALUES and row.cancel_requested_at is None: row.cancel_requested_at = requested_at if row.next_cancel_at is None: row.next_cancel_at = requested_at # A cancel request fences any in-flight poll result, so its # poll lease can be released immediately for the cancellation # worker. A repeated request must preserve an existing cancel # lease so it cannot trigger a concurrent remote cancellation. row.lease_owner = None row.lease_expires_at = None row.updated_at = requested_at await session.commit() return self._row_to_dict(row) async def claim_cancel_requests( self, *, now: datetime, lease_owner: str, lease_seconds: int, limit: int, task_id: str | None = None, ) -> list[dict[str, Any]]: stmt = select(McpTaskRow).where( McpTaskRow.cancel_requested_at.is_not(None), McpTaskRow.status.not_in(_TERMINAL_STATUS_VALUES), McpTaskRow.next_cancel_at.is_not(None), McpTaskRow.next_cancel_at <= now, or_(McpTaskRow.lease_expires_at.is_(None), McpTaskRow.lease_expires_at < now), ) if task_id is not None: stmt = stmt.where(McpTaskRow.id == task_id) stmt = stmt.order_by(McpTaskRow.next_cancel_at.asc(), McpTaskRow.id.asc()).limit(limit).with_for_update(skip_locked=True) async with self._sf() as session: rows = list((await session.execute(stmt)).scalars()) expires_at = now + timedelta(seconds=lease_seconds) for row in rows: row.lease_owner = lease_owner row.lease_expires_at = expires_at row.cancel_attempt_count = int(row.cancel_attempt_count or 0) + 1 row.updated_at = now await session.commit() return [self._row_to_dict(row) for row in rows] async def apply_cancel_snapshot( self, task_id: str, *, lease_owner: str, status: str, result: Any | None, result_preview: str | None, result_truncated: bool, result_artifact: dict[str, str] | None, error: str | None, input_required: dict[str, Any] | None, completed_at: datetime, ) -> bool: 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) .where( McpTaskRow.id == task_id, McpTaskRow.lease_owner == lease_owner, McpTaskRow.lease_expires_at >= completed_at, McpTaskRow.status.not_in(_TERMINAL_STATUS_VALUES), ) .with_for_update() ) row = (await session.execute(stmt)).scalar_one_or_none() if row is None: 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) await session.commit() return True async def release_cancel_claim( self, task_id: str, *, lease_owner: str, next_cancel_at: datetime, error: str, ) -> bool: stmt = ( update(McpTaskRow) .where(McpTaskRow.id == task_id, McpTaskRow.lease_owner == lease_owner) .values( next_cancel_at=next_cancel_at, last_cancel_error=error, lease_owner=None, lease_expires_at=None, updated_at=datetime.now(UTC), ) ) async with self._sf() as session: result = await session.execute(stmt) await session.commit() return bool(result.rowcount) async def claim_notification_work( self, *, now: datetime, lease_owner: str, lease_seconds: int, limit: int, tracking_degraded_after_errors: int, ) -> list[dict[str, Any]]: statuses = ("pending", "claimed", "retry", "dispatched") stmt = ( select(McpTaskRow) .where( McpTaskRow.event_version > McpTaskRow.notified_version, McpTaskRow.notification_status.in_(statuses), or_(McpTaskRow.next_notification_at.is_(None), McpTaskRow.next_notification_at <= now), or_(McpTaskRow.notification_lease_expires_at.is_(None), McpTaskRow.notification_lease_expires_at < now), ) .order_by(McpTaskRow.next_notification_at.asc(), McpTaskRow.id.asc()) .limit(limit) .with_for_update(skip_locked=True) ) async with self._sf() as session: rows = list((await session.execute(stmt)).scalars()) expires_at = now + timedelta(seconds=lease_seconds) for row in rows: row.notification_lease_owner = lease_owner row.notification_lease_expires_at = expires_at 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: row.dispatch_attempt = 0 row.notification_attempt_count = 0 row.dispatch_version = row.event_version row.dispatch_event = _notification_event( row, tracking_degraded=int(row.consecutive_poll_error_count or 0) >= tracking_degraded_after_errors, ) row.notification_run_id = None row.notification_status = "claimed" row.updated_at = now await session.commit() return [self._row_to_dict(row) for row in rows] async def mark_notification_dispatched( self, task_id: str, *, lease_owner: str, dispatch_version: int, run_id: str, now: datetime, ) -> bool: stmt = ( update(McpTaskRow) .where( McpTaskRow.id == task_id, McpTaskRow.notification_lease_owner == lease_owner, McpTaskRow.notification_lease_expires_at >= now, McpTaskRow.dispatch_version == dispatch_version, McpTaskRow.notification_status.in_(("claimed", "retry")), ) .values( notification_status="dispatched", notification_run_id=run_id, notification_error=None, next_notification_at=now, notification_lease_owner=None, notification_lease_expires_at=None, updated_at=now, ) ) async with self._sf() as session: result = await session.execute(stmt) await session.commit() return bool(result.rowcount) async def release_notification_claim( self, task_id: str, *, lease_owner: str, next_notification_at: datetime, error: str, replace_with_latest: bool, count_failure: bool = False, ) -> bool: values: dict[str, Any] = { "notification_status": "pending" if replace_with_latest else "retry", "notification_error": error, "next_notification_at": next_notification_at, "notification_lease_owner": None, "notification_lease_expires_at": None, "updated_at": datetime.now(UTC), } if replace_with_latest: values.update( dispatch_event=None, notification_run_id=None, ) 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) async with self._sf() as session: result = await session.execute(stmt) await session.commit() return bool(result.rowcount) async def finish_notification_run( self, task_id: str, *, lease_owner: str, dispatch_version: int, delivered: bool, next_notification_at: datetime | None, error: str | None, now: datetime, ) -> bool: async with self._sf() as session: stmt = ( select(McpTaskRow) .where( McpTaskRow.id == task_id, McpTaskRow.notification_lease_owner == lease_owner, McpTaskRow.notification_lease_expires_at >= now, McpTaskRow.dispatch_version == dispatch_version, McpTaskRow.notification_status == "dispatched", ) .with_for_update() ) 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 async def release_notification_lease( self, task_id: str, *, lease_owner: str, next_notification_at: datetime, error: str, count_failure: bool = False, ) -> bool: """Release unexpected notification work without changing its phase.""" values: dict[str, Any] = { "notification_error": error, "next_notification_at": next_notification_at, "notification_lease_owner": None, "notification_lease_expires_at": None, "updated_at": datetime.now(UTC), } 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) ) async with self._sf() as session: result = await session.execute(stmt) await session.commit() return bool(result.rowcount) async def dead_letter_notification( self, task_id: str, *, lease_owner: str, dispatch_version: int, error: str, count_failure: bool, now: datetime, ) -> bool: """Stop one failed snapshot, preserving any newer event for delivery.""" base_filters = ( McpTaskRow.id == task_id, McpTaskRow.notification_lease_owner == lease_owner, McpTaskRow.notification_lease_expires_at >= now, McpTaskRow.dispatch_version == dispatch_version, McpTaskRow.notification_status.in_(("claimed", "retry", "dispatched")), ) dead_letter_values: dict[str, Any] = { "notification_status": "dead_letter", "notification_error": error, "next_notification_at": None, "notification_lease_owner": None, "notification_lease_expires_at": None, "dispatch_version": None, "dispatch_attempt": 0, "dispatch_event": None, "notification_run_id": None, "updated_at": now, } if count_failure: dead_letter_values["notification_attempt_count"] = McpTaskRow.notification_attempt_count + 1 async with self._sf() as session: dead_lettered = await session.execute(update(McpTaskRow).where(*base_filters, McpTaskRow.event_version <= dispatch_version).values(**dead_letter_values)) if dead_lettered.rowcount: await session.commit() return True replaced_by_latest = await session.execute( update(McpTaskRow) .where(*base_filters, McpTaskRow.event_version > dispatch_version) .values( notification_status="pending", notification_error=None, notification_attempt_count=0, next_notification_at=now, notification_lease_owner=None, notification_lease_expires_at=None, dispatch_version=None, dispatch_attempt=0, dispatch_event=None, notification_run_id=None, updated_at=now, ) ) await session.commit() return bool(replaced_by_latest.rowcount) async def defer_dispatched_notification( self, task_id: str, *, lease_owner: str, dispatch_version: int, next_notification_at: datetime, now: datetime, ) -> bool: """Release a notification lease while its Agent run is still active.""" stmt = ( update(McpTaskRow) .where( McpTaskRow.id == task_id, McpTaskRow.notification_lease_owner == lease_owner, McpTaskRow.notification_lease_expires_at >= now, McpTaskRow.dispatch_version == dispatch_version, McpTaskRow.notification_status == "dispatched", ) .values( next_notification_at=next_notification_at, notification_lease_owner=None, notification_lease_expires_at=None, updated_at=now, ) ) async with self._sf() as session: result = await session.execute(stmt) await session.commit() return bool(result.rowcount)