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

from __future__ import annotations

import asyncio
import logging
import socket
import uuid
from datetime import UTC, datetime, timedelta
from typing import Any, Literal

from fastapi import HTTPException

from deerflow.persistence.scheduled_task_runs import ActiveScheduledRunConflict, ScheduledTaskAdmissionRejected
from deerflow.runtime import ConflictError, RunRecord
from deerflow.scheduler.schedules import next_run_at
from deerflow.trace_context import ensure_trace_context
from deerflow.utils.thread_id import validate_thread_id

logger = logging.getLogger(__name__)

# Shared so the active-row fast path and the atomic-admission conflict path
# return byte-identical outcomes for the same active-occurrence condition.
_ACTIVE_RUN_CONFLICT_ERROR = "task already has an active run"
_RESTART_RECOVERY_ERROR = "interrupted: gateway restarted before the run reached a terminal state"
_LEASE_RECOVERY_ERROR = "interrupted: the owning gateway stopped renewing its run lease"
_QUEUE_TIMEOUT_ERROR = "scheduled task queue wait timeout exceeded"


class ScheduledTaskService:
    def __init__(
        self,
        *,
        task_repo,
        task_run_repo,
        launch_run,
        poll_interval_seconds: int,
        lease_seconds: int,
        max_concurrent_runs: int,
        queue_timeout_seconds: int = 3600,
        multi_instance: bool = False,
        run_lease_grace_seconds: int = 10,
    ) -> None:
        self._task_repo = task_repo
        self._task_run_repo = task_run_repo
        self._launch_run = launch_run
        self._poll_interval_seconds = poll_interval_seconds
        self._lease_seconds = lease_seconds
        self._max_concurrent_runs = max_concurrent_runs
        self._queue_timeout_seconds = queue_timeout_seconds
        self._multi_instance = multi_instance
        self._run_lease_grace_seconds = run_lease_grace_seconds
        self._lease_owner = f"{socket.gethostname()}:{uuid.uuid4().hex}"
        self._task: asyncio.Task | None = None
        self._stop = asyncio.Event()
        self._skip_next_lease_reconciliation = False

    async def run_once(self, *, now: datetime) -> None:
        if self._multi_instance:
            if self._skip_next_lease_reconciliation:
                self._skip_next_lease_reconciliation = False
            else:
                await self._reconcile_active_state(now=now)
        else:
            await self._task_run_repo.recover_expired_launch_claims(
                error=_LEASE_RECOVERY_ERROR,
                now=now,
            )
        await self._expire_waiting_runs(now=now)
        await self._drain_queue(now=now)
        # Admission and execution capacity are separate. Due occurrences are
        # persisted even when all execution slots are busy; claim_queued_run()
        # applies the global launch budget under the database lock.
        claimed = await self._task_repo.claim_due_tasks(
            now=now,
            lease_owner=self._lease_owner,
            lease_seconds=self._lease_seconds,
            limit=self._max_concurrent_runs,
        )
        for task in claimed:
            await self.dispatch_task(task, now=now, trigger="scheduled")

    @staticmethod
    def _is_overlap_conflict(exc: Exception) -> bool:
        if isinstance(exc, ConflictError):
            return True
        return isinstance(exc, HTTPException) and exc.status_code == 409

    @staticmethod
    def _task_status_for_failure(task: dict[str, Any], *, trigger: str) -> str:
        if trigger == "manual":
            # A failed manual trigger must not consume the task's scheduled
            # future: a `once` task with run_at still ahead would otherwise be
            # flipped to "failed" and never claimed again.
            return task.get("status") or "enabled"
        if task["schedule_type"] == "once":
            return "failed"
        return "enabled"

    @staticmethod
    def _task_status_for_launch(task: dict[str, Any], *, trigger: str) -> str:
        # The task-level status to write once _launch_run has produced a live
        # run. A `once` task stays "running" until handle_run_completion
        # observes the real terminal outcome; declaring "completed" at launch
        # would stick if the run fails or the process dies (startup
        # reconciliation is cancel_stuck_once_tasks).
        if task["schedule_type"] == "once":
            return "running"
        if trigger == "manual" and task.get("status") == "paused":
            return "paused"
        return "enabled"

    async def dispatch_task(
        self,
        task: dict[str, Any],
        *,
        now: datetime,
        trigger: str,
    ) -> dict[str, Any]:
        expected_lease_owner = self._lease_owner if trigger == "scheduled" else None
        execution_thread_id = task.get("thread_id")
        if task.get("context_mode") == "fresh_thread_per_run" or execution_thread_id is None:
            execution_thread_id = str(uuid.uuid4())
        try:
            validate_thread_id(execution_thread_id)
        except ValueError as exc:
            # Rows persisted before the thread-id contract was centralized may
            # hold IDs that were valid then (dots, unlimited length) but fail
            # the canonical pattern now. Route through the normal failure
            # bookkeeping instead of raising: an uncaught ValueError here would
            # surface as HTTP 500 on manual trigger and, in the poller, abort
            # the rest of the claimed batch every cycle while the task itself
            # is never marked with last_error.
            task_status = self._task_status_for_failure(task, trigger=trigger)
            await self._task_repo.update_after_launch(
                task["id"],
                status=task_status,
                next_run_at=next_run_at(
                    task["schedule_type"],
                    task["schedule_spec"],
                    task["timezone"],
                    now=now,
                ),
                last_run_at=now,
                last_run_id=None,
                last_thread_id=execution_thread_id,
                last_error=str(exc),
                increment_run_count=False,
                expected_lease_owner=expected_lease_owner,
            )
            return {
                "outcome": "failed",
                "task_run_id": None,
                "run_id": None,
                "thread_id": execution_thread_id,
                "error": str(exc),
            }
        active = await self._task_run_repo.get_active_run(task["id"])
        if active is not None:
            if trigger == "scheduled":
                await self._release_admission_lease(task, trigger=trigger)
            return self._existing_active_result(active, execution_thread_id, trigger=trigger)

        task_run_id = f"task-run-{uuid.uuid4().hex}"
        try:
            await self._task_run_repo.create(
                run_record_id=task_run_id,
                task_id=task["id"],
                thread_id=execution_thread_id,
                scheduled_for=now,
                trigger=trigger,
                status="queued",
                coordinate_with_task=True,
                expected_task_user_id=task.get("user_id"),
                expected_task_status=task.get("status") if trigger == "manual" else None,
                expected_task_updated_at=task.get("updated_at") if trigger == "manual" else None,
                expected_task_lease_owner=self._lease_owner if trigger == "scheduled" else None,
                release_task_lease_status="enabled" if trigger == "scheduled" else None,
            )
        except ActiveScheduledRunConflict:
            active = await self._task_run_repo.get_active_run(task["id"])
            if trigger == "scheduled":
                await self._release_admission_lease(task, trigger=trigger)
            if active is None:
                return self._active_run_conflict_result(execution_thread_id)
            return self._existing_active_result(active, execution_thread_id, trigger=trigger)
        except ScheduledTaskAdmissionRejected as exc:
            if exc.reason == "not_found":
                return {
                    "outcome": "not_found",
                    "task_run_id": None,
                    "run_id": None,
                    "thread_id": execution_thread_id,
                    "error": "scheduled task no longer exists",
                }
            return {
                "outcome": "conflict",
                "task_run_id": None,
                "run_id": None,
                "thread_id": execution_thread_id,
                "error": "scheduled task changed before trigger admission",
            }

        # Scheduled admission inserted the queue row and released its parent
        # lease in one transaction. Manual admission verified that this task
        # snapshot was still current under the same parent lock.
        queued = {
            "id": task_run_id,
            "task_id": task["id"],
            "thread_id": execution_thread_id,
            "trigger": trigger,
        }
        return await self._attempt_queued_run(task, queued, now=now)

    async def _release_admission_lease(self, task: dict[str, Any], *, trigger: str) -> None:
        status = "enabled" if trigger == "scheduled" else (task.get("status") or "enabled")
        await self._task_repo.release_dispatch_lease(
            task["id"],
            expected_lease_owner=self._lease_owner if trigger == "scheduled" else None,
            status=status,
        )

    async def _attempt_queued_run(
        self,
        task: dict[str, Any],
        queued: dict[str, Any],
        *,
        now: datetime,
    ) -> dict[str, Any]:
        """Turn one queued occurrence into a live run under its own trace scope.

        The poller is a non-HTTP entry point, so no ``TraceMiddleware`` has
        bound anything: each occurrence opens its own scope rather than
        sharing one id across a whole poll cycle. A manual trigger arrives
        inside a Gateway request and keeps that request's trace instead, so
        the launched run stays correlated with the call that asked for it.
        """
        with ensure_trace_context():
            return await self._launch_queued_occurrence(task, queued, now=now)

    async def _launch_queued_occurrence(
        self,
        task: dict[str, Any],
        queued: dict[str, Any],
        *,
        now: datetime,
    ) -> dict[str, Any]:
        task_run_id = queued["id"]
        execution_thread_id = queued["thread_id"]
        trigger = queued["trigger"]
        claimed = await self._task_run_repo.claim_queued_run(
            task_run_id,
            lease_owner=self._lease_owner,
            now=now,
            lease_seconds=self._lease_seconds,
            global_max_concurrent_runs=self._max_concurrent_runs,
        )
        if claimed is None:
            return self._queued_result(task_run_id, execution_thread_id)

        # Track whether _launch_run has produced a live run. A bookkeeping
        # failure after launch must retain the non-terminal slot so a later
        # poll cannot start the same occurrence twice.
        launched_run_id: str | None = None
        launched_thread_id: str | None = None
        launch_succeeded = False
        try:
            result = await self._launch_run(
                thread_id=execution_thread_id,
                assistant_id=task.get("assistant_id"),
                prompt=task["prompt"],
                owner_user_id=task.get("user_id"),
                metadata={
                    "scheduled_task_id": task["id"],
                    "scheduled_task_run_id": task_run_id,
                    "scheduled_trigger": trigger,
                },
            )
            launch_succeeded = True
            launched_run_id = result["run_id"]
            launched_thread_id = result["thread_id"]
            next_at = next_run_at(
                task["schedule_type"],
                task["schedule_spec"],
                task["timezone"],
                now=now,
            )
            task_status = self._task_status_for_launch(task, trigger=trigger)
            await self._record_launched_run(
                task_run_id=task_run_id,
                task_id=task["id"],
                run_id=launched_run_id,
                started_at=now,
            )
            await self._task_repo.update_after_launch(
                task["id"],
                status=task_status,
                next_run_at=next_at,
                last_run_at=now,
                last_run_id=launched_run_id,
                last_thread_id=launched_thread_id,
                last_error=None,
                increment_run_count=True,
                task_run_id=task_run_id,
                # Same race as the run-row write above: a fast-failing run's
                # completion hook may have already finalized a `once` task.
                protect_terminal=True,
            )
            return {
                "outcome": "launched",
                "task_run_id": task_run_id,
                "run_id": launched_run_id,
                "thread_id": launched_thread_id,
                "error": None,
            }
        except Exception as exc:
            if not launch_succeeded and self._is_overlap_conflict(exc):
                await self._task_run_repo.requeue_claimed_run(
                    task_run_id,
                    lease_owner=self._lease_owner,
                    error=str(exc),
                )
                return self._queued_result(task_run_id, execution_thread_id, error=str(exc))

            next_at = next_run_at(
                task["schedule_type"],
                task["schedule_spec"],
                task["timezone"],
                now=now,
            )

            if launch_succeeded:
                # _launch_run succeeded, so a run is live even though
                # post-launch bookkeeping raised. Keep the task-run row
                # "running" so it keeps holding the task's single active slot
                # (preventing a duplicate launch on the next dispatch) and
                # persist the run_id on the parent task for recovery /
                # reconciliation / cancellation. These writes are best-effort:
                # if the DB is still down the row stays "queued" -- still
                # active, still holding the slot -- so we log and still report
                # the run as launched so callers know a run is in flight.
                task_status = self._task_status_for_launch(task, trigger=trigger)
                try:
                    await self._record_launched_run(
                        task_run_id=task_run_id,
                        task_id=task["id"],
                        run_id=launched_run_id,
                        started_at=now,
                    )
                except Exception:
                    logger.exception(
                        "Scheduled task-run %s: post-launch bookkeeping failed; run %s is still live (task %s)",
                        task_run_id,
                        launched_run_id,
                        task["id"],
                    )
                try:
                    await self._task_repo.update_after_launch(
                        task["id"],
                        status=task_status,
                        next_run_at=next_at,
                        last_run_at=now,
                        last_run_id=launched_run_id,
                        last_thread_id=launched_thread_id,
                        # The bookkeeping exception is an infrastructure-level
                        # transient, not a run-level failure: the run launched
                        # and is still in flight. Clear last_error like the
                        # success path so the task list does not show an error
                        # on a task whose run is actively running; the real
                        # terminal outcome is written by handle_run_completion.
                        # The transient itself is logged above.
                        last_error=None,
                        increment_run_count=True,
                        task_run_id=task_run_id,
                        protect_terminal=True,
                    )
                except Exception:
                    logger.exception(
                        "Scheduled task %s: post-launch update failed; run %s is still live",
                        task["id"],
                        launched_run_id,
                    )
                return {
                    "outcome": "launched",
                    "task_run_id": task_run_id,
                    "run_id": launched_run_id,
                    "thread_id": launched_thread_id,
                    "error": str(exc),
                }

            # _launch_run itself failed (or a step before it did): no live run
            # was created, so it is safe to release the active slot.
            finalized = await self._task_run_repo.fail_launching_run(
                task_run_id,
                task_id=task["id"],
                lease_owner=self._lease_owner,
                error=str(exc),
                now=now,
            )
            if not finalized:
                logger.warning(
                    "Scheduled task-run %s lost its launch claim before failure bookkeeping; leaving recovery-owned state unchanged",
                    task_run_id,
                )
                return self._queued_result(task_run_id, execution_thread_id, error=str(exc))
            return {
                "outcome": "failed",
                "task_run_id": task_run_id,
                "run_id": None,
                "thread_id": execution_thread_id,
                "error": str(exc),
            }

    async def _record_launched_run(
        self,
        *,
        task_run_id: str,
        task_id: str,
        run_id: str,
        started_at: datetime,
    ) -> None:
        updated = await self._task_run_repo.update_status(
            task_run_id,
            status="running",
            run_id=run_id,
            started_at=started_at,
            protect_terminal=True,
            expected_lease_owner=self._lease_owner,
        )
        if updated:
            return
        reconciled = await self._task_run_repo.reconcile_launched_run(
            task_run_id,
            task_id=task_id,
            run_id=run_id,
            started_at=started_at,
        )
        if not reconciled:
            logger.error(
                "Scheduled task-run %s launched durable run %s but could not restore its occurrence association",
                task_run_id,
                run_id,
            )

    def _active_run_conflict_result(self, thread_id: str) -> dict[str, Any]:
        """Manual-trigger response when the task already has an active run.

        Nothing was scheduled to happen, so no run-history row is recorded; the
        router maps this to a 409.
        """
        return {
            "outcome": "conflict",
            "task_run_id": None,
            "run_id": None,
            "thread_id": thread_id,
            "error": _ACTIVE_RUN_CONFLICT_ERROR,
        }

    def _existing_active_result(
        self,
        active: dict[str, Any],
        thread_id: str,
        *,
        trigger: str,
    ) -> dict[str, Any]:
        if active["status"] == "queued":
            return self._queued_result(active["id"], active["thread_id"])
        return self._active_run_conflict_result(thread_id)

    @staticmethod
    def _queued_result(
        task_run_id: str,
        thread_id: str,
        *,
        error: str | None = None,
    ) -> dict[str, Any]:
        return {
            "outcome": "queued",
            "task_run_id": task_run_id,
            "run_id": None,
            "thread_id": thread_id,
            "error": error,
        }

    async def _drain_queue(self, *, now: datetime) -> None:
        queued_rows = await self._task_run_repo.list_queued_runs(limit=max(16, self._max_concurrent_runs * 4))
        for queued in queued_rows:
            await self._task_repo.release_queued_admission_lease(queued["task_id"])
            task = await self._task_repo.get_internal(queued["task_id"])
            if task is None:
                await self._task_run_repo.update_status(
                    queued["id"],
                    status="interrupted",
                    error="scheduled task was deleted while queued",
                    finished_at=now,
                )
                continue
            # Pausing suppresses automatic occurrences, but a manual trigger is
            # an explicit request and has always been allowed to run without
            # resuming the schedule. A later pause still cancels an already
            # queued manual row atomically in pause_with_queue_cancellation().
            if task.get("status") == "paused" and queued["trigger"] != "manual":
                await self._task_run_repo.update_status(
                    queued["id"],
                    status="interrupted",
                    error="scheduled task was paused while queued",
                    finished_at=now,
                )
                continue
            await self._attempt_queued_run(task, queued, now=now)

    async def _expire_waiting_runs(self, *, now: datetime) -> None:
        await self._task_run_repo.expire_queued_runs(
            created_before=now - timedelta(seconds=self._queue_timeout_seconds),
            error=_QUEUE_TIMEOUT_ERROR,
            now=now,
        )

    async def handle_run_completion(self, record: RunRecord) -> None:
        metadata = record.metadata or {}
        task_id = metadata.get("scheduled_task_id")
        task_run_id = metadata.get("scheduled_task_run_id")
        user_id = record.user_id
        if not isinstance(task_id, str) or not isinstance(task_run_id, str) or not user_id:
            return

        terminal_status: Literal["success", "failed", "interrupted"] | None
        if record.status.value == "success":
            terminal_status = "success"
            error = None
        elif record.status.value == "interrupted":
            # Distinct from "failed": an interrupt (user cancel, same-thread
            # takeover) carries no error and is not an execution failure.
            terminal_status = "interrupted"
            error = record.error or "run was interrupted before completion"
        elif record.status.value in {"error", "timeout"}:
            terminal_status = "failed"
            error = record.error
        else:
            terminal_status = None
            error = record.error
        if terminal_status is None:
            return

        await self._task_repo.complete_run(
            task_id,
            user_id=user_id,
            task_run_id=task_run_id,
            run_id=record.run_id,
            status=terminal_status,
            error=error,
            finished_at=datetime.now(UTC),
        )

    async def start(self) -> None:
        if self._task is not None:
            return
        restart_error = _RESTART_RECOVERY_ERROR
        if self._multi_instance:
            await self._reconcile_active_state(now=datetime.now(UTC))
            self._skip_next_lease_reconciliation = True
        else:
            # This destructive sweep is safe only while Gateway lifespan awaits
            # start(): no request or poll admission can create a run owned by
            # this process yet. Complete occurrence -> parent recovery before
            # returning; moving either pass into run_once() can interrupt live
            # work or race manual admission.
            try:
                stale = await self._task_run_repo.mark_stale_active_runs(error=restart_error)
                if stale:
                    logger.warning("Marked %d stale scheduled task run(s) as interrupted after restart", stale)
            except Exception:
                logger.exception("Failed to sweep stale scheduled task runs at startup")
                raise
            try:
                # The run rows above are only half the story: a launched `once`
                # task is parked in "running" until the (now dead) completion hook
                # would have finalized it.
                stuck = await self._task_repo.cancel_stuck_once_tasks(error=restart_error)
                if stuck:
                    logger.warning("Reconciled %d stuck once task(s) after restart", stuck)
            except Exception:
                logger.exception("Failed to reconcile stuck once tasks at startup")
                raise
        self._stop.clear()
        self._task = asyncio.create_task(self._run_loop())

    async def _reconcile_active_state(self, *, now: datetime) -> None:
        error = _LEASE_RECOVERY_ERROR
        try:
            stale = await self._task_run_repo.reconcile_active_runs(
                error=error,
                now=now,
                lease_grace_seconds=self._run_lease_grace_seconds,
            )
            if stale:
                logger.warning("Marked %d stale scheduled task run(s) as interrupted after lease reconciliation", stale)
        except Exception:
            logger.exception("Failed to reconcile scheduled task runs with leases")
        try:
            stuck = await self._task_repo.reconcile_stuck_once_tasks(
                error=error,
                now=now,
                lease_grace_seconds=self._run_lease_grace_seconds,
            )
            if stuck:
                logger.warning("Reconciled %d stuck once task(s) after lease reconciliation", stuck)
        except Exception:
            logger.exception("Failed to reconcile once tasks with leases")

    async def stop(self) -> None:
        if self._task is None:
            return
        self._stop.set()
        await self._task
        self._task = None

    async def _run_loop(self) -> None:
        while not self._stop.is_set():
            try:
                await self.run_once(now=datetime.now(UTC))
            except Exception:
                # A transient DB error (e.g. SQLite "database is locked") must
                # not kill the poller task for the rest of the process life.
                logger.exception("Scheduled task poll failed; retrying next interval")
            try:
                await asyncio.wait_for(
                    self._stop.wait(),
                    timeout=self._poll_interval_seconds,
                )
            except TimeoutError:
                continue