MoviePilot/app/db/models/agenttask.py

from typing import Optional

from sqlalchemy import Boolean, Index, Integer, String, Text, select, update
from sqlalchemy.orm import Mapped, Session, mapped_column

from app.db.base import Base, execute_dml, get_id_column
from app.db.decorators import run_legacy_sync_query


def _get_for_user_statement(
        model: type["AgentTask"],
        task_id: int,
        user_id: Optional[str] = None,
):
    """构造按任务 ID 与可选用户归属收窄的查询语句。"""
    statement = select(model).where(model.id == task_id)
    if user_id is not None:
        statement = statement.where(model.user_id == user_id)
    return statement


def _list_for_user_statement(
        model: type["AgentTask"],
        user_id: Optional[str] = None,
        enabled: Optional[bool] = None,
):
    """构造按用户、启用状态和创建时间排序的任务列表语句。"""
    statement = select(model)
    if user_id is not None:
        statement = statement.where(model.user_id == user_id)
    if enabled is not None:
        statement = statement.where(model.enabled.is_(enabled))
    return statement.order_by(model.created_at.desc(), model.id.desc())


class AgentTask(Base):
    """
    Agent 自主定时任务表。
    """

    id = get_id_column()
    # 任务名称
    name: Mapped[str] = mapped_column(String, nullable=False)
    # 交给 Agent 执行的完整任务内容
    content: Mapped[str] = mapped_column(Text, nullable=False)
    # 触发类型：date-单次触发，cron-周期触发
    trigger_type: Mapped[str] = mapped_column(String, nullable=False)
    # 标准五段 cron 表达式
    cron_expression: Mapped[Optional[str]] = mapped_column(String)
    # 单次触发时间，使用带时区的 ISO 8601 格式
    run_at: Mapped[Optional[str]] = mapped_column(String)
    # 是否继续接受调度
    enabled: Mapped[bool] = mapped_column(Boolean, nullable=False, default=True)
    # 创建任务的用户与会话上下文
    user_id: Mapped[str] = mapped_column(String, nullable=False)
    username: Mapped[Optional[str]] = mapped_column(String)
    session_id: Mapped[str] = mapped_column(String, nullable=False)
    channel: Mapped[Optional[str]] = mapped_column(String)
    source: Mapped[Optional[str]] = mapped_column(String)
    original_chat_id: Mapped[Optional[str]] = mapped_column(String)
    # 最近一次执行状态与结果
    last_status: Mapped[str] = mapped_column(String, nullable=False, default="waiting")
    last_run_at: Mapped[Optional[str]] = mapped_column(String)
    last_result: Mapped[Optional[str]] = mapped_column(Text)
    # 最新一次真实执行的公开 ID，用于保护 last_* 投影不被旧运行覆盖
    last_run_id: Mapped[Optional[str]] = mapped_column(String)
    # 已收口执行次数；进程中断的未完成尝试不计入
    run_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
    created_at: Mapped[str] = mapped_column(String, nullable=False)
    updated_at: Mapped[str] = mapped_column(String, nullable=False)

    __table_args__ = (
        Index("ix_agenttask_enabled", "enabled"),
        Index("ix_agenttask_user_created", "user_id", "created_at", "id"),
    )

    @classmethod
    def add_task(cls, db: Session, **kwargs: object) -> int:
        """
        新增 Agent 定时任务并返回任务 ID。
        """
        task = cls(**kwargs)
        db.add(task)
        db.flush()
        return task.id

    @classmethod
    def get_for_user(
            cls,
            db: Session | int | None = None,
            task_id: int | None = None,
            user_id: Optional[str] = None,
    ) -> Optional["AgentTask"]:
        """
        按任务 ID 和可选用户 ID 查询，并保留无 Session 的旧插件调用方式。
        """
        if task_id is None and isinstance(db, int):
            task_id, db = db, None
        if task_id is None:
            raise TypeError("task_id is required")

        def query(session: Session) -> Optional["AgentTask"]:
            """在给定会话中读取单个 Agent 任务。"""
            return session.execute(
                _get_for_user_statement(cls, task_id=task_id, user_id=user_id)
            ).scalars().first()

        if isinstance(db, Session):
            return query(db)
        return run_legacy_sync_query(query)

    @classmethod
    def list_for_user(
            cls,
            db: Session | None = None,
            user_id: Optional[str] = None,
            enabled: Optional[bool] = None,
    ) -> list["AgentTask"]:
        """
        按用户和启用状态查询，并保留无 Session 的旧插件调用方式。
        """
        def query(session: Session) -> list["AgentTask"]:
            """在给定会话中读取 Agent 任务列表。"""
            return list(session.execute(
                _list_for_user_statement(cls, user_id=user_id, enabled=enabled)
            ).scalars().all())

        if isinstance(db, Session):
            return query(db)
        return run_legacy_sync_query(query)

    @classmethod
    def update_task(
            cls,
            db: Session,
            task_id: int,
            payload: dict,
            user_id: Optional[str] = None,
    ) -> bool:
        """
        仅在任务未运行时按任务 ID 和可选用户 ID 更新配置。

        运行状态与配置必须在同一条条件更新中判定，避免执行认领后被并发配置写入
        覆盖回可再次执行的状态。
        """
        statement = update(cls).where(
            cls.id == task_id,
            cls.last_status != "running",
        )
        if user_id is not None:
            statement = statement.where(cls.user_id == user_id)
        return bool(execute_dml(db, statement.values(payload)))