AgentLab/backend/app/scheduler/loop.py

"""调度循环：每一拍认领到点的定时任务、发起运行，再顺便查已发起的运行要不要推送通知。

认领靠条件 UPDATE：UPDATE schedules SET next_run_at = 下一次 WHERE id = … AND next_run_at = 旧值。两拍撞在
一起（或者以后多开了一个进程），只有一拍改得动这一行，另一拍拿到 0 行就知道这一次已经有人处理了。触发记录
(任务, fire_at, kind) 唯一是第二道保险。

错过的触发（服务停机、进程卡住）：只看最近应触发的那一次——在宽限内补一次（kind=misfire），超出就记一条
skipped_misfire；更早的积压只记一条跳过，不补。下一次一律推进到现在之后，不会醒来就把积压的一口气全发起。
"""
from __future__ import annotations

import asyncio
import contextlib
import logging
from datetime import datetime

from sqlalchemy import select, update

from app.db.base import SessionLocal, utcnow
from app.db.models import Schedule, ScheduleFire
from app.scheduler import cron, service

logger = logging.getLogger(__name__)

#: 一拍的间隔（秒）。测试里 start(interval=…) 调小
TICK_SECONDS = 15.0
#: 晚了这么多秒以内都算准点（kind=scheduled），和补发宽限无关：拍与拍之间的十几秒不该算错过
ON_TIME_S = 60
#: stop 等手上那一拍的发起收尾最多等多久，超过就取消。关停时 main.py 先停调度、再让 run_manager 平稳挂起在执行
#: 的运行，docker stop 默认只给 10 秒：这里要留出大半
STOP_GRACE_S = 3.0
#: stop 等正在推送的通知多久。通知本来就是至多一次，接收端挂住时等满超时加重试（几十秒）换不来什么，宽限一过就取消
NOTIFY_GRACE_S = 1.5


async def fire_due(now: datetime | None = None) -> list[ScheduleFire]:
    """处理所有到点的任务（enabled 且 next_run_at <= now），返回这一拍记下的触发。一个任务出错不影响别的。"""
    now = now or utcnow()
    async with SessionLocal() as session:
        due = (await session.execute(
            select(Schedule.id, Schedule.next_run_at)
            .where(Schedule.enabled.is_(True), Schedule.next_run_at.is_not(None), Schedule.next_run_at <= now)
            .order_by(Schedule.next_run_at)
        )).all()
    out: list[ScheduleFire] = []
    for schedule_id, claimed_at in due:
        try:
            out.extend(await _due_one(schedule_id, claimed_at, now))
        except Exception:  # noqa: BLE001 - 一个任务出错不能拖住这一拍的其他任务，也不能让循环停下
            logger.exception("定时任务 %s 这一拍处理失败", schedule_id)
    return out


async def _due_one(schedule_id: str, claimed_at: datetime, now: datetime) -> list[ScheduleFire]:
    async with SessionLocal() as session:
        row = await session.get(Schedule, schedule_id)
        if row is None or not row.enabled or row.next_run_at != claimed_at:
            return []
        cron_expr, tz, grace = row.cron, row.timezone, max(int(row.misfire_grace_s or 0), 0)
    try:
        spec, zone = cron.parse(cron_expr), cron.zone(tz)
        # 最近应触发的那一次。库里的 next_run_at 理应就是一个触发时刻，万一不是（手改过库）就按它算
        latest = max(cron.last_due(spec, zone, now) or claimed_at, claimed_at)
        upcoming = cron.next_fire(spec, zone, now)
    except (cron.CronError, cron.TimezoneError) as e:
        # 存进库里之后才坏掉的表达式（手改过库、时区库升级删了时区）：停止排期，在触发历史里照实说
        if not await _claim(schedule_id, claimed_at, None):
            return []
        logger.warning("定时任务 %s 的 cron 或时区无法解析，已停止排期：%s", schedule_id, e)
        fire = await service._record(schedule_id, claimed_at, "scheduled", "failed_to_start",
                                     f"cron 表达式或时区无法解析，已停止排期：{e}", {})
        return [fire] if fire else []
    if not await _claim(schedule_id, claimed_at, upcoming):
        return []

    out: list[ScheduleFire | None] = []
    lateness = (now - latest).total_seconds()
    if lateness <= max(ON_TIME_S, grace):
        if latest > claimed_at:
            # 最近这一次照常（或补一次），更早的积压只记一条，不补
            last_missed = cron.prev_fire(spec, zone, latest) or claimed_at
            reason = f"服务未运行，错过了 {_span(claimed_at, last_missed, tz)} 的触发；积压的触发不补，只发起最近一次"
            out.append(await service.record_skip(schedule_id, claimed_at, reason))
        kind = "scheduled" if lateness <= ON_TIME_S else "misfire"
        out.append(await service.fire_once(schedule_id, latest, kind))
    else:
        reason = (f"服务未运行，错过了 {_span(claimed_at, latest, tz)} 的触发，"
                  f"超出补发宽限（{service.duration_text(grace)}），没有发起")
        out.append(await service.record_skip(schedule_id, latest, reason))
    return [fire for fire in out if fire is not None]


async def _claim(schedule_id: str, claimed_at: datetime, upcoming: datetime | None) -> bool:
    """把 next_run_at 从 claimed_at 推进到 upcoming。改动了一行才算认领到。

    updated_at 照旧：它说的是任务最近一次被人修改，推进排期不算。version 也不动——乐观锁防的是两个人同时改设置，
    调度推进 next_run_at 不该让正在编辑的人保存时撞上 409。"""
    async with SessionLocal() as session:
        result = await session.execute(
            update(Schedule)
            .where(Schedule.id == schedule_id, Schedule.enabled.is_(True), Schedule.next_run_at == claimed_at)
            .values(next_run_at=upcoming, updated_at=Schedule.updated_at)
            .execution_options(synchronize_session=False)
        )
        await session.commit()
        return result.rowcount == 1


def _span(first: datetime, last: datetime, tz: str) -> str:
    a, b = service.local_text(first, tz), service.local_text(last, tz)
    return a if a == b else f"{a} 至 {b}"


async def tick(now: datetime | None = None) -> None:
    """完整的一拍：发起到点的，处理流程的时限和核实，再推送该推送的通知。测试和手动调用用；循环里分开跑（见 _Runner._loop）。"""
    from app.scheduler import webhook

    await fire_due(now)
    await process_due(now)
    await webhook.notify_due(now)


async def process_due(now: datetime | None = None) -> None:
    """流程的一拍（app/process）：待办的提醒、升级、超时驳回（永不自动批准）；结果未知的写入按退避去查；
    上传了却一直没提交的文件过了时限清掉（app/files，每小时一次）；访问密钥的请求记录批量写库、过了保留期的清掉
    （auth/keys.py，每小时清一次：每个请求同步写一行会和运行抢 SQLite 的写锁，所以攒到这一拍写）；每天到点发待办汇总、
    访问密钥到期前 7 天和 1 天各提醒一次（notify/digest.py，04 细案 B3：按日期认领、按审计去重，推送放在后台，不拖这一拍）。
    各自出错只记日志：这一拍出错不该让调度循环停下，下一拍照常再查。"""
    from app.auth import keys
    from app.files import store as files_store
    from app.notify import digest
    from app.process import actions, tasks

    for name, step in (("待办的时限处理", tasks.tick), ("写入结果的查询", actions.tick),
                       ("没有提交的上传文件清理", files_store.cleanup_tick), ("访问密钥的请求记录", keys.tick),
                       ("每天的待办汇总", digest.digest_due), ("访问密钥快到期的提醒", digest.keys_due)):
        try:
            done = await step(now)
            if isinstance(done, list) and done:
                logger.info("%s：这一拍处理了 %d 条", name, len(done))
        except Exception:  # noqa: BLE001
            logger.exception("%s这一拍出错，下一拍照常进行", name)


async def _notify(webhook) -> None:  # noqa: ANN001 - 传进来的是 app.scheduler.webhook 模块
    """后台跑一遍通知扫描。出错只记日志：这个任务没人 await，异常留在里面只会变成一句「从未取回的异常」。"""
    try:
        await webhook.notify_due()
    except Exception:  # noqa: BLE001
        logger.exception("定时任务的通知推送这一拍出错，下一拍照常进行")


class _Runner:
    """调度循环的任务。照 sandbox/manager.py 回收闲置会话的写法：一个 asyncio 任务，出错记日志接着下一拍。"""

    def __init__(self) -> None:
        self.task: asyncio.Task[None] | None = None
        #: 通知推送另起任务：Webhook 超时加重试最长要几十秒，不能拖着下一拍的发起。同一时刻只跑一份
        self.notify: asyncio.Task[None] | None = None
        self._stopping: asyncio.Event | None = None

    def running(self) -> bool:
        return self.task is not None and not self.task.done()

    async def start(self, interval: float = TICK_SECONDS) -> None:
        if self.running():
            return
        self._stopping = asyncio.Event()
        self.task = asyncio.create_task(self._loop(interval, self._stopping), name="agentlab-scheduler")
        logger.info("定时任务调度已启动，每 %s 秒检查一次", interval)

    async def _loop(self, interval: float, stopping: asyncio.Event) -> None:
        from app.scheduler import webhook

        while not stopping.is_set():
            try:
                await fire_due()
                await process_due()
                # stop 撞在这一拍的发起当中时，收尾不再新起通知扫描：stop 只等它动手时手上的那一份
                if not stopping.is_set() and (self.notify is None or self.notify.done()):
                    self.notify = asyncio.create_task(_notify(webhook), name="agentlab-scheduler-notify")
            except asyncio.CancelledError:
                raise
            except Exception:  # noqa: BLE001 - 这一拍出错不该让循环停下
                logger.exception("定时任务调度这一拍出错，下一拍照常进行")
            with contextlib.suppress(TimeoutError):
                await asyncio.wait_for(stopping.wait(), timeout=interval)

    async def stop(self) -> None:
        """让循环收完手上这一拍再停、正在推送的通知推完；各等一个宽限（STOP_GRACE_S、NOTIFY_GRACE_S），两件事并行
        等，过了就取消。没 start 过直接返回。

        先置停止标志、再取两个任务：标志置上之后循环不会再起新的通知扫描，取到的就是最后一份。"""
        if self._stopping is not None:
            self._stopping.set()
        tick, notify = self.task, self.notify
        await asyncio.gather(_settle(tick, STOP_GRACE_S), _settle(notify, NOTIFY_GRACE_S))
        if self.notify is not None and self.notify is not notify:
            # 兜底：万一还是落下了一份（以后改循环时漏掉上面的条件），不让它在关停期间接着推
            await _settle(self.notify, 0)
        self.task = self.notify = None
        # 关停前把还攒在进程里的访问密钥请求记录写进库（没起调度循环时也写：请求记录不该因为关停丢掉最后一拍）
        from app.auth import keys

        try:
            await keys.flush()
        except Exception:  # noqa: BLE001 - 写不进去只记日志，不拦关停
            logger.exception("关停前写访问密钥的请求记录失败")
        if tick is not None or notify is not None:
            logger.info("定时任务调度已停止")


async def _settle(task: asyncio.Task[None] | None, grace: float) -> None:
    """等 task 结束，最多等 grace 秒，过了就取消。task 自己的异常不再追究：停的时候只管它停下。"""
    if task is None or task.done():
        return
    try:
        await asyncio.wait_for(asyncio.shield(task), timeout=grace)
    except TimeoutError:
        task.cancel()
        with contextlib.suppress(asyncio.CancelledError, Exception):
            await task
    except Exception:  # noqa: BLE001
        pass


runner = _Runner()