MoviePilot/app/application/agenttask.py

"""Agent 自主定时任务执行的异步应用边界。"""

from __future__ import annotations

import asyncio
from collections.abc import Callable
from dataclasses import dataclass
from typing import Protocol, TypeVar
from uuid import uuid4

from app.application.database import AsyncDatabaseExecutor
from app.runtime.execution import await_task_to_terminal
from app.schemas.exception import DatabaseWorkerOverloadedError


T = TypeVar("T")
SyncTransaction = Callable[[Callable[[object], T]], T]


class AgentTaskRecord(Protocol):
    """执行认领与终态判定所需的任务投影。"""

    id: int
    enabled: bool
    last_run_id: str | None
    last_status: str


class AgentTaskRunRecord(Protocol):
    """执行期间需要脱离数据库会话持有的运行记录字段。"""

    run_id: str
    task_id: int
    trigger_source: str
    name: str
    content: str
    trigger_type: str
    cron_expression: str | None
    run_at: str | None
    user_id: str
    username: str | None
    session_id: str


class AgentTaskRepository(Protocol):
    """AgentTask 执行用例使用的同步短事务仓储合同。"""

    def get(self, task_id: int) -> AgentTaskRecord | None:
        """读取任务当前投影。"""

    def begin_run(
        self,
        task_id: int,
        trigger_source: str = "scheduled",
        *,
        run_id: str | None = None,
    ) -> AgentTaskRunRecord | None:
        """原子认领任务并创建运行快照。"""

    def finish_run_outcome(
        self,
        run_id: str,
        success: bool,
        result: str,
    ) -> AgentTaskFinishRecord:
        """收口运行并返回事务确认的终态事实。"""


class AgentTaskFinishRecord(Protocol):
    """同步仓储返回的结构化运行终态。"""

    run_finalized: bool
    task_projection_updated: bool
    date_task_disabled: bool


AgentTaskRepositoryFactory = Callable[[object], AgentTaskRepository]
AgentTaskScheduleRemover = Callable[[int, int, str], bool]


@dataclass(frozen=True, slots=True)
class AgentTaskRunSnapshot:
    """任务认领成功后可安全跨越数据库会话的执行快照。"""

    run_id: str
    task_id: int
    trigger_source: str
    name: str
    content: str
    trigger_type: str
    cron_expression: str | None
    run_at: str | None
    user_id: str
    username: str | None
    session_id: str


@dataclass(frozen=True, slots=True)
class AgentTaskClaim:
    """任务认领结果；拒绝原因保持现有 Agent 用户提示合同。"""

    run: AgentTaskRunSnapshot | None
    rejection: str | None = None


@dataclass(frozen=True, slots=True)
class AgentTaskFinishOutcome:
    """区分运行收口、任务投影更新与一次任务停用三个事实。"""

    run_finalized: bool
    task_projection_updated: bool
    date_task_disabled: bool


class AgentTaskExecutionService:
    """通过有界数据库 worker 认领并收口一次 AgentTask 执行。"""

    def __init__(
        self,
        *,
        repository: AgentTaskRepositoryFactory,
        async_executor: AsyncDatabaseExecutor,
        sync_transaction: SyncTransaction,
    ) -> None:
        """保存同步仓储、事务和异步执行器。"""
        self._repository = repository
        self._async_executor = async_executor
        self._sync_transaction = sync_transaction

    @staticmethod
    def _snapshot(run: AgentTaskRunRecord) -> AgentTaskRunSnapshot:
        """在事务内复制运行字段，避免 ORM 对象越过会话边界。"""
        return AgentTaskRunSnapshot(
            run_id=run.run_id,
            task_id=run.task_id,
            trigger_source=run.trigger_source,
            name=run.name,
            content=run.content,
            trigger_type=run.trigger_type,
            cron_expression=run.cron_expression,
            run_at=run.run_at,
            user_id=run.user_id,
            username=run.username,
            session_id=run.session_id,
        )

    async def claim(
        self,
        task_id: int,
        trigger_source: str = "scheduled",
        *,
        scheduler_generation: int | None = None,
        remove_schedule: AgentTaskScheduleRemover | None = None,
    ) -> AgentTaskClaim:
        """认领一次执行；取消发生在提交后时先补偿收口再传播取消。"""

        run_id = uuid4().hex

        def transaction(session: object) -> AgentTaskClaim:
            repository = self._repository(session)
            run = repository.begin_run(
                task_id=task_id,
                trigger_source=trigger_source,
                run_id=run_id,
            )
            if not run:
                task = repository.get(task_id)
                return AgentTaskClaim(
                    run=None,
                    rejection=(
                        "Agent 定时任务不存在或已停用"
                        if not task or not task.enabled
                        else "Agent 定时任务当前不可执行"
                    ),
                )
            return AgentTaskClaim(run=self._snapshot(run))

        async def claim_to_terminal() -> AgentTaskClaim:
            """容量瞬时耗尽时保留本轮调度，直到认领取得 admission。"""
            while True:
                try:
                    return await self._async_executor.run(
                        lambda: self._sync_transaction(transaction)
                    )
                except DatabaseWorkerOverloadedError:
                    await asyncio.sleep(0.01)

        claim_task = asyncio.create_task(claim_to_terminal())
        try:
            return await claim_task
        except asyncio.CancelledError as cancellation:
            finalize_task = asyncio.create_task(self._finalize(
                run_id=run_id,
                task_id=task_id,
                success=False,
                result="Agent 定时任务已取消",
                scheduler_generation=scheduler_generation,
                remove_schedule=remove_schedule,
            ))
            await await_task_to_terminal(finalize_task)
            raise cancellation

    async def finalize(
        self,
        run: AgentTaskRunSnapshot,
        *,
        success: bool,
        result: str,
        scheduler_generation: int | None = None,
        remove_schedule: AgentTaskScheduleRemover | None = None,
    ) -> AgentTaskFinishOutcome:
        """等待终态事务完成，并仅清理仍属于该 generation 的一次任务。"""

        return await self._finalize(
            run_id=run.run_id,
            task_id=run.task_id,
            success=success,
            result=result,
            scheduler_generation=scheduler_generation,
            remove_schedule=remove_schedule,
        )

    async def _finalize(
        self,
        *,
        run_id: str,
        task_id: int,
        success: bool,
        result: str,
        scheduler_generation: int | None,
        remove_schedule: AgentTaskScheduleRemover | None,
    ) -> AgentTaskFinishOutcome:
        """按稳定运行 ID 收口，供正常路径和取消补偿共享。"""

        def transaction(session: object) -> AgentTaskFinishOutcome:
            repository = self._repository(session)
            outcome = repository.finish_run_outcome(
                run_id=run_id,
                success=success,
                result=result,
            )
            return AgentTaskFinishOutcome(
                run_finalized=outcome.run_finalized,
                task_projection_updated=outcome.task_projection_updated,
                date_task_disabled=outcome.date_task_disabled,
            )

        async def finish_to_terminal() -> AgentTaskFinishOutcome:
            """容量瞬时耗尽时保留 owner，直到收口取得 admission。"""
            while True:
                try:
                    return await self._async_executor.run(
                        lambda: self._sync_transaction(transaction)
                    )
                except DatabaseWorkerOverloadedError:
                    await asyncio.sleep(0.01)

        finish_task = asyncio.create_task(finish_to_terminal())
        cancellation: asyncio.CancelledError | None = None
        try:
            outcome = await asyncio.shield(finish_task)
        except asyncio.CancelledError as error:
            cancellation = error
            outcome = await await_task_to_terminal(finish_task)

        if (
            outcome.date_task_disabled
            and scheduler_generation is not None
            and remove_schedule is not None
        ):
            remove_schedule(task_id, scheduler_generation, run_id)
        if cancellation is not None:
            raise cancellation
        return outcome


_service: AgentTaskExecutionService | None = None


def configure_agent_task_execution(service: AgentTaskExecutionService) -> None:
    """由启动组合根登记 AgentTask 执行服务。"""
    global _service
    _service = service


def get_agent_task_execution_service() -> AgentTaskExecutionService:
    """返回已登记的 AgentTask 执行服务。"""
    if _service is None:
        raise RuntimeError("AgentTask 执行服务尚未配置")
    return _service