deer-flow/backend/packages/harness/deerflow/domain/schedule/model/task.py

from __future__ import annotations

import uuid
from dataclasses import dataclass, field, replace
from datetime import UTC, datetime

from deerflow.domain.schedule.exceptions import InvalidContextModeError, TaskNotMutableError
from deerflow.domain.schedule.model.enums import ContextMode, RunStatus, ScheduleType, TaskStatus, TriggerKind
from deerflow.domain.schedule.model.spec import SchedulePolicy, ScheduleSpec

TERMINAL_TASK_STATUSES = frozenset({TaskStatus.COMPLETED, TaskStatus.FAILED, TaskStatus.CANCELLED})
"""The statuses in which a task's current schedule is done.

Lives here rather than in enums.py: it is a business subset of TaskStatus, so
it belongs beside the aggregate that reasons about it.

Two rules read it, and they share one constant rather than drifting into two
same-valued ones:

- `with_schedule` re-arms them to ENABLED once the schedule moves into the
  future — claiming only admits ENABLED rows, so leaving a terminal status
  there would hand back a next_run_at that silently never fires.
- The repository's `protect_terminal` compare-and-set refuses to overwrite
  them, because a fast-failing run's completion hook can land before the launch
  path's own write.

The two coincide because terminal *is* the definition of re-armable. The name
states what the statuses are; what each rule does with them belongs to that
rule.
"""


def _parse_context_mode(value: str | ContextMode) -> ContextMode:
    """Coerce a wire-level context_mode into the enum.

    Reported as InvalidContextModeError rather than letting ValueError
    escape, so the router can map the whole ScheduleError family uniformly
    — an unknown mode was a 422 at router:76-77.
    """
    try:
        return ContextMode(value)
    except ValueError as exc:
        raise InvalidContextModeError(f"Unsupported context_mode: {value}") from exc


@dataclass(frozen=True)
class ScheduledTask:
    """A user's standing instruction to run one prompt on a schedule.

    Aggregate root of the schedule context. Holds every invariant the old
    router/service pair scattered across three files; nothing here knows about
    HTTP, SQL, or the run runtime.

    The two bookkeeping timestamps have one owner each. `created_at` is the
    aggregate's construction instant -- `create` stamps it from its explicit
    `now` input and the adapter persists it verbatim, never minting its own.
    `updated_at` belongs to the *storage* write path: the CAS methods
    (`record_launch` / `record_completion`) update the row without ever
    holding an aggregate, so only the adapter can stamp "last written" -- the
    `with_*` / `paused` / `resumed` transitions therefore deliberately leave
    `updated_at` alone, and the field's value on a read-back is the adapter's
    stamp, seeded at construction time. Neither timestamp is a rule input;
    the clock the rules see is always an explicit `now=` parameter.
    """

    task_id: str
    user_id: str
    title: str
    prompt: str
    schedule: ScheduleSpec
    context_mode: ContextMode = ContextMode.FRESH_THREAD_PER_RUN
    thread_id: str | None = None
    assistant_id: str | None = "lead_agent"
    status: TaskStatus = TaskStatus.ENABLED
    # The MVP fixes this to "skip"; a single-valued enum would be noise. The
    # comparison is confined to `skips_on_overlap` so a second policy only has
    # to change one place.
    overlap_policy: str = "skip"
    next_run_at: datetime | None = None
    last_run_at: datetime | None = None
    last_run_id: str | None = None
    last_thread_id: str | None = None
    last_error: str | None = None
    run_count: int = 0
    created_at: datetime = field(default_factory=lambda: datetime.now(UTC))
    updated_at: datetime = field(default_factory=lambda: datetime.now(UTC))

    def __post_init__(self) -> None:
        if self.context_mode is ContextMode.REUSE_THREAD and not self.thread_id:
            raise InvalidContextModeError("reuse_thread requires thread_id")

    @classmethod
    def create(
        cls,
        *,
        user_id: str,
        title: str,
        prompt: str,
        schedule: ScheduleSpec,
        context_mode: str | ContextMode,
        thread_id: str | None,
        now: datetime,
        policy: SchedulePolicy,
    ) -> ScheduledTask:
        """Factory: generate identity, normalize context, validate invariants.

        FRESH_THREAD_PER_RUN drops any supplied thread_id (router:164-165 does
        the same on update) — carrying a thread the task will never use would
        make `resolve_execution_thread` ambiguous.

        Thread existence/ownership is NOT checked here: that needs the
        ThreadLookup port and belongs to the service.
        """
        mode = _parse_context_mode(context_mode)
        effective_thread = thread_id if mode is ContextMode.REUSE_THREAD else None
        return cls(
            task_id=f"task-{uuid.uuid4().hex}",
            user_id=user_id,
            title=title,
            prompt=prompt,
            schedule=schedule,
            context_mode=mode,
            thread_id=effective_thread,
            next_run_at=schedule.ensure_launchable(now, policy),
            # The clock is already an explicit rule input here; reading it a
            # second time through the field defaults would give one
            # construction two instants.
            created_at=now,
            updated_at=now,
        )

    @property
    def skips_on_overlap(self) -> bool:
        """Whether an overlapping dispatch is dropped rather than queued."""
        return self.overlap_policy == "skip"

    def resolve_execution_thread(self) -> str:
        """Pick the thread this dispatch executes in (service.py:94-96).

        NOT idempotent: FRESH_THREAD_PER_RUN mints a new uuid4 on every call,
        so a dispatch must call this exactly once and reuse the value for the
        run row, the launch, and the result.

        The empty-thread_id fallback is kept even though __post_init__ makes it
        unreachable for new aggregates — rows predating that invariant can
        still carry REUSE_THREAD with no thread.
        """
        if self.context_mode is ContextMode.FRESH_THREAD_PER_RUN or not self.thread_id:
            return str(uuid.uuid4())
        return self.thread_id

    def ensure_mutable(self) -> None:
        """Reject edits while this dispatch round's lease is held.

        Was `_ensure_task_mutable` in the router (409). Called by every
        state-changing method here rather than by the caller, so the rule
        cannot be forgotten at a new call site. The router historically applied
        it to update/pause/resume but not trigger/delete; that asymmetry is
        intentional and is preserved by keeping those two off this path.

        Raises:
            TaskNotMutableError: the task is currently RUNNING.
        """
        if self.status is TaskStatus.RUNNING:
            raise TaskNotMutableError("Scheduled task is currently running; retry after the active execution finishes")

    def status_after_launch(self, *, trigger: TriggerKind) -> TaskStatus:
        """Status to persist after a successful launch (service.py:152-161).

        Precondition on `self.status`: for a SCHEDULED trigger it is already
        RUNNING (written by claim_due); for a MANUAL trigger it is the
        user-facing status. The PAUSED rule below depends on that.

        Decision order is load-bearing — ONCE is tested first, so manually
        triggering a paused ONCE task yields RUNNING, not PAUSED:

            ONCE                          -> RUNNING (await the completion
                                             hook; declaring COMPLETED at
                                             launch would stick if the run
                                             fails or the process dies)
            MANUAL and self.status PAUSED -> PAUSED  (a manual run must not
                                             silently resume the schedule)
            otherwise                     -> ENABLED
        """
        if self.schedule.schedule_type is ScheduleType.ONCE:
            return TaskStatus.RUNNING
        if trigger is TriggerKind.MANUAL and self.status is TaskStatus.PAUSED:
            return TaskStatus.PAUSED
        return TaskStatus.ENABLED

    def status_after_failure(self, *, trigger: TriggerKind) -> TaskStatus:
        """Status to persist after a failed launch (was _task_status_for_failure).

        Decision order is load-bearing — MANUAL is tested first, so a failed
        manual trigger of a ONCE task keeps its status instead of burning the
        occurrence:

            MANUAL          -> self.status (unchanged; the old dict-based code
                               needed an `or "enabled"` fallback, which the
                               field default makes unnecessary here)
            SCHEDULED, ONCE -> FAILED
            SCHEDULED, CRON -> ENABLED (the next occurrence still stands)
        """
        if trigger is TriggerKind.MANUAL:
            return self.status
        if self.schedule.schedule_type is ScheduleType.ONCE:
            return TaskStatus.FAILED
        return TaskStatus.ENABLED

    def status_after_skip(self) -> TaskStatus:
        """Status to persist after an overlap skip (was _task_status_for_skip).

        No `trigger` parameter: a skip only happens on the SCHEDULED path — a
        manual trigger that overlaps is rejected outright and records no run
        row at all, so a parameter here would imply a branch that cannot exist.

            ONCE -> FAILED   (the single occurrence was lost; COMPLETED would
                              claim an execution that never happened)
            CRON -> ENABLED
        """
        if self.schedule.schedule_type is ScheduleType.ONCE:
            return TaskStatus.FAILED
        return TaskStatus.ENABLED

    def status_after_completion(self, outcome: RunStatus) -> TaskStatus | None:
        """Status to persist when a launched run reaches a terminal state.

        Returns None when the status must not change — every CRON task, whose
        schedule outlives any single run.

            CRON              -> None
            ONCE, SUCCESS     -> COMPLETED
            ONCE, INTERRUPTED -> CANCELLED (a cancel or same-thread takeover is
                                 not an execution failure)
            ONCE, otherwise   -> FAILED

        The occurrence is consumed either way: the run did launch, so re-arming
        would risk duplicate side effects. Note the caller writes `last_error`
        unconditionally, whether or not this returns None (service.py:346).
        """
        if self.schedule.schedule_type is not ScheduleType.ONCE:
            return None
        if outcome is RunStatus.SUCCESS:
            return TaskStatus.COMPLETED
        if outcome is RunStatus.INTERRUPTED:
            return TaskStatus.CANCELLED
        return TaskStatus.FAILED

    def with_schedule(self, schedule: ScheduleSpec, *, now: datetime, policy: SchedulePolicy) -> ScheduledTask:
        """Replace the schedule and recompute next_run_at (router:172-208).

        A terminal task whose schedule just moved into the future must be
        re-armed: claim_due only admits ENABLED rows, so leaving the terminal
        status would return 200 with a next_run_at that silently never fires.

        Raises:
            TaskNotMutableError: via ensure_mutable.
            InvalidScheduleError: via ensure_launchable.
        """
        self.ensure_mutable()
        next_at = schedule.ensure_launchable(now, policy)
        status = TaskStatus.ENABLED if next_at is not None and self.status in TERMINAL_TASK_STATUSES else self.status
        return replace(self, schedule=schedule, next_run_at=next_at, status=status)

    def with_context(self, context_mode: str | ContextMode, thread_id: str | None) -> ScheduledTask:
        """Replace execution context (router:153-165).

        FRESH_THREAD_PER_RUN forces thread_id to None. Thread existence is the
        service's job (ThreadLookup port), not this aggregate's.

        Raises:
            TaskNotMutableError: via ensure_mutable.
            InvalidContextModeError: unknown mode, or REUSE_THREAD with no
                thread (raised by __post_init__ on the replacement).
        """
        self.ensure_mutable()
        mode = _parse_context_mode(context_mode)
        effective_thread = thread_id if mode is ContextMode.REUSE_THREAD else None
        return replace(self, context_mode=mode, thread_id=effective_thread)

    def paused(self) -> ScheduledTask:
        """Stop claiming this task until resumed."""
        self.ensure_mutable()
        return replace(self, status=TaskStatus.PAUSED)

    def resumed(self) -> ScheduledTask:
        """Re-admit this task to claim_due."""
        self.ensure_mutable()
        return replace(self, status=TaskStatus.ENABLED)