""" 服务层 (Service Manager) 定义 SchedulerManager 类作为定时任务服务的公共 API 入口。 它负责编排业务逻辑，并调用 Repository 和 Adapter 层来完成具体工作。 """ from __future__ import annotations from collections.abc import Awaitable, Callable, Coroutine from datetime import datetime import inspect from typing import Any, ClassVar import uuid from arclet.alconna import Alconna, Option import nonebot from nonebot.adapters import Bot from pydantic import BaseModel from zhenxun.configs.config import Config from zhenxun.models.scheduled_job import ScheduledJob from zhenxun.services.log import logger from zhenxun.utils.pydantic_compat import model_dump, model_validate from .engine import APSchedulerAdapter from .repository import ScheduleRepository from .targeting import ( ScheduleTargeter, ) from .types import ( BaseTrigger, EphemeralJobDeclaration, ExecutionOptions, ExecutionPolicy, ScheduleContext, ScheduledJobDeclaration, ) class SchedulerManager: ALL_GROUPS: ClassVar[str] = "__ALL_GROUPS__" _registered_tasks: ClassVar[ dict[ str, dict[str, Callable | type[BaseModel] | int | list[Option] | Alconna | None], ] ] = {} _declared_tasks: ClassVar[list[ScheduledJobDeclaration]] = [] _ephemeral_declared_tasks: ClassVar[list[EphemeralJobDeclaration]] = [] _running_tasks: ClassVar[set] = set() _target_resolvers: ClassVar[ dict[str, Callable[[str, Bot], Awaitable[list[str | None]]]] ] = {} def __init__(self): self._register_builtin_resolvers() def _register_builtin_resolvers(self): """在管理器初始化时注册所有内置的目标解析器。""" from .targeting import ( _resolve_all_groups, _resolve_global_or_user, _resolve_group, _resolve_tag, _resolve_user, ) if "GROUP" in self._target_resolvers: return self.register_target_resolver("GROUP", _resolve_group) self.register_target_resolver("TAG", _resolve_tag) self.register_target_resolver("ALL_GROUPS", _resolve_all_groups) self.register_target_resolver("GLOBAL", _resolve_global_or_user) self.register_target_resolver("USER", _resolve_user) logger.debug("已注册所有内置的定时任务目标解析器。") def register_target_resolver( self, target_type: str, resolver_func: Callable[[str, Bot], Awaitable[list[str | None]]], ): """ 注册一个新的目标类型解析器。 """ if target_type in self._target_resolvers: logger.warning(f"目标解析器 '{target_type}' 已存在，将被覆盖。") self._target_resolvers[target_type.upper()] = resolver_func logger.info(f"已注册新的定时任务目标解析器: '{target_type}'") def target(self, **filters: Any) -> ScheduleTargeter: """ 创建目标选择器以执行批量操作 参数: **filters: 过滤条件，支持plugin_name、group_id、bot_id等字段。 返回: ScheduleTargeter: 目标选择器对象，可用于批量操作。 """ return ScheduleTargeter(self, **filters) def job( self, trigger: BaseTrigger, group_id: str | None = None, bot_id: str | None = None, default_params: BaseModel | None = None, policy: ExecutionPolicy | None = None, default_jitter: int | None = None, default_spread: int | None = None, default_interval: int | None = None, ): """ 声明式定时任务的统一装饰器。 """ def decorator(func: Callable[..., Coroutine]) -> Callable[..., Coroutine]: try: plugin = nonebot.get_plugin_by_module_name(func.__module__) if not plugin: raise ValueError(f"函数 {func.__name__} 不在任何已加载的插件中。") plugin_name = plugin.name params_model = None from .types import ScheduleContext for param in inspect.signature(func).parameters.values(): if ( isinstance(param.annotation, type) and issubclass(param.annotation, BaseModel) and param.annotation is not ScheduleContext ): params_model = param.annotation break if plugin_name in self._registered_tasks: logger.warning(f"插件 '{plugin_name}' 的定时任务已被重复注册。") self._registered_tasks[plugin_name] = { "func": func, "model": params_model, "default_jitter": default_jitter, "default_spread": default_spread, "default_interval": default_interval, } job_kwargs = model_dump(default_params) if default_params else {} if policy: job_kwargs["execution_policy"] = model_dump(policy) task_declaration = ScheduledJobDeclaration( plugin_name=plugin_name, group_id=group_id, bot_id=bot_id, trigger=trigger, job_kwargs=job_kwargs, ) self._declared_tasks.append(task_declaration) logger.debug( f"发现声明式定时任务 '{plugin_name}'，将在启动时进行注册。" ) except Exception as e: logger.error(f"注册声明式定时任务失败: {func.__name__}, 错误: {e}") return func return decorator def runtime_job(self, trigger: BaseTrigger): """ 声明一个临时的、非持久化的定时任务。 """ def decorator(func: Callable[..., Coroutine]) -> Callable[..., Coroutine]: try: plugin = nonebot.get_plugin_by_module_name(func.__module__) if not plugin: raise ValueError(f"函数 {func.__name__} 不在任何已加载的插件中。") plugin_name = plugin.name self._registered_tasks[f"ephemeral::{plugin_name}::{func.__name__}"] = { "func": func, "model": None, } declaration = EphemeralJobDeclaration( plugin_name=plugin_name, func=func, trigger=trigger, ) self._ephemeral_declared_tasks.append(declaration) logger.debug( f"发现临时定时任务 '{plugin_name}:{func.__name__}'，将在启动时调度" ) except Exception as e: logger.error(f"注册临时定时任务失败: {func.__name__}, 错误: {e}") return func return decorator def register( self, plugin_name: str, params_model: type[BaseModel] | None = None, cli_parser: Alconna | None = None, default_permission: int = 5, default_jitter: int | None = None, default_spread: int | None = None, default_interval: int | None = None, ) -> Callable: """ 注册可调度的任务函数 """ def decorator(func: Callable[..., Coroutine]) -> Callable[..., Coroutine]: if plugin_name in self._registered_tasks: logger.warning(f"插件 '{plugin_name}' 的定时任务已被重复注册。") self._registered_tasks[plugin_name] = { "func": func, "model": params_model, "cli_parser": cli_parser, "default_permission": default_permission, "default_jitter": default_jitter, "default_spread": default_spread, "default_interval": default_interval, } model_name = params_model.__name__ if params_model else "无" logger.debug( f"插件 '{plugin_name}' 的定时任务已注册，参数模型: {model_name}" ) return func return decorator def get_registered_plugins(self) -> list[str]: """ 获取已注册插件列表 """ return list(self._registered_tasks.keys()) async def run_at(self, func: Callable[..., Coroutine], trigger: BaseTrigger) -> str: """ 在未来的某个时间点，运行一个一次性的临时任务。 """ job_id = f"ephemeral_runtime_{uuid.uuid4()}" context = ScheduleContext( schedule_id=0, plugin_name=f"runtime::{func.__module__}", bot_id=None, group_id=None, job_kwargs={}, ) APSchedulerAdapter.add_ephemeral_job( job_id=job_id, func=func, trigger_type=trigger.trigger_type, trigger_config=model_dump(trigger, exclude={"trigger_type"}), context=context, ) logger.info(f"已动态调度一个临时任务 (ID: {job_id})，将在 {trigger} 触发。") return job_id async def schedule_once( self, func: Callable[..., Coroutine], trigger: BaseTrigger, *, user_id: str | None = None, group_id: str | None = None, bot_id: str | None = None, job_kwargs: dict | None = None, name: str | None = None, created_by: str | None = None, required_permission: int = 5, ) -> "ScheduledJob | None": """ 编程式API，用于动态调度一个持久化的、一次性的任务。 """ if user_id and group_id: raise ValueError("user_id 和 group_id 不能同时提供。") temp_plugin_name = f"runtime_one_off__{func.__module__}.{func.__name__}__{uuid.uuid4().hex[:8]}" # noqa: E501 self._registered_tasks[temp_plugin_name] = {"func": func, "model": None} logger.debug(f"为一次性任务动态注册临时插件: '{temp_plugin_name}'") target_type = "USER" if user_id else ("GROUP" if group_id else "GLOBAL") target_identifier = user_id or group_id or "" return await self.add_schedule( plugin_name=temp_plugin_name, target_type=target_type, target_identifier=target_identifier, trigger_type=trigger.trigger_type, trigger_config=model_dump(trigger, exclude={"trigger_type"}), job_kwargs=job_kwargs, bot_id=bot_id, name=name, created_by=created_by, required_permission=required_permission, is_one_off=True, ) async def add_daily_task( self, plugin_name: str, group_id: str | None, hour: int, minute: int, second: int = 0, job_kwargs: dict | None = None, bot_id: str | None = None, ) -> "ScheduledJob | None": """ 添加每日定时任务 """ trigger_config = { "hour": hour, "minute": minute, "second": second, "timezone": Config.get_config("SchedulerManager", "SCHEDULER_TIMEZONE"), } return await self.add_schedule( plugin_name, target_type="GROUP" if group_id else "GLOBAL", target_identifier=group_id or "", trigger_type="cron", trigger_config=trigger_config, job_kwargs=job_kwargs, bot_id=bot_id, ) async def add_interval_task( self, plugin_name: str, group_id: str | None, *, weeks: int = 0, days: int = 0, hours: int = 0, minutes: int = 0, seconds: int = 0, start_date: str | datetime | None = None, job_kwargs: dict | None = None, bot_id: str | None = None, ) -> "ScheduledJob | None": """ 添加间隔性定时任务 """ trigger_config = { "weeks": weeks, "days": days, "hours": hours, "minutes": minutes, "seconds": seconds, "start_date": start_date, } trigger_config = {k: v for k, v in trigger_config.items() if v} return await self.add_schedule( plugin_name, target_type="GROUP" if group_id else "GLOBAL", target_identifier=group_id or "", trigger_type="interval", trigger_config=trigger_config, job_kwargs=job_kwargs, bot_id=bot_id, ) def _validate_and_prepare_kwargs( self, plugin_name: str, job_kwargs: dict | None ) -> tuple[bool, str | dict]: """验证并准备任务参数，应用默认值""" from pydantic import ValidationError task_meta = self._registered_tasks.get(plugin_name) if not task_meta: return False, f"插件 '{plugin_name}' 未注册。" params_model = task_meta.get("model") job_kwargs = job_kwargs if job_kwargs is not None else {} if not params_model: if job_kwargs: logger.warning( f"插件 '{plugin_name}' 未定义参数模型，但收到了参数: {job_kwargs}" ) return True, job_kwargs if not (isinstance(params_model, type) and issubclass(params_model, BaseModel)): logger.error(f"插件 '{plugin_name}' 的参数模型不是有效的 BaseModel 类") return False, f"插件 '{plugin_name}' 的参数模型配置错误" try: validated_model = model_validate(params_model, job_kwargs) return True, model_dump(validated_model) except ValidationError as e: errors = [f" - {err['loc'][0]}: {err['msg']}" for err in e.errors()] error_str = "\n".join(errors) msg = f"插件 '{plugin_name}' 的任务参数验证失败:\n{error_str}" return False, msg async def add_schedule( self, plugin_name: str, target_type: str, target_identifier: str, trigger_type: str, trigger_config: dict, job_kwargs: dict | None = None, bot_id: str | None = None, *, name: str | None = None, created_by: str | None = None, required_permission: int = 5, source: str = "USER", is_one_off: bool = False, execution_options: dict | None = None, ) -> "ScheduledJob | None": """ 添加定时任务（通用方法） 参数: plugin_name: 插件名称。 target_type: 目标类型 (GROUP, USER, TAG, ALL_GROUPS, GLOBAL)。 target_identifier: 目标标识符。 trigger_type: 触发器类型 (cron, interval, date)。 trigger_config: 触发器配置字典。 job_kwargs: 传递给任务函数的额外参数。 bot_id: Bot ID约束。 name: 任务别名。 created_by: 创建者ID。 required_permission: 管理此任务所需的权限。 source: 任务来源 (USER, PLUGIN_DEFAULT)。 is_one_off: 是否为一次性任务。 execution_options: 任务执行的额外选项 (例如: jitter, spread)。 返回: ScheduledJob | None: 创建的任务信息，失败时返回None。 """ if plugin_name not in self._registered_tasks: logger.error(f"插件 '{plugin_name}' 没有注册可用的定时任务。") return None is_valid, result = self._validate_and_prepare_kwargs(plugin_name, job_kwargs) if not is_valid: logger.error(f"任务参数校验失败: {result}") return None options_dict = execution_options or {} validated_options = ExecutionOptions(**options_dict) search_kwargs = { "plugin_name": plugin_name, "target_type": target_type, "target_identifier": target_identifier, } if bot_id: search_kwargs["bot_id"] = bot_id defaults = { "name": name, "trigger_type": trigger_type, "trigger_config": trigger_config, "job_kwargs": result, "is_enabled": True, "created_by": created_by, "required_permission": required_permission, "source": source, "is_one_off": is_one_off, "execution_options": model_dump( validated_options, exclude_none=True ), } defaults = {k: v for k, v in defaults.items() if v is not None} schedule, created = await ScheduleRepository.update_or_create( defaults, **search_kwargs ) APSchedulerAdapter.add_or_reschedule_job(schedule) action_str = "创建" if created else "更新" logger.info( f"已成功{action_str}任务 '{name or plugin_name}' (ID: {schedule.id})" ) return schedule async def get_schedules( self, page: int | None = None, page_size: int | None = None, **filters: Any ) -> tuple[list[ScheduledJob], int]: """ 根据条件获取定时任务列表 """ cleaned_filters = {k: v for k, v in filters.items() if v is not None} return await ScheduleRepository.query_schedules( page=page, page_size=page_size, **cleaned_filters ) async def get_schedules_status_bulk( self, schedule_ids: list[int] ) -> list[dict[str, Any]]: """ 批量获取多个定时任务的详细状态信息 """ if not schedule_ids: return [] schedules = await ScheduleRepository.filter(id__in=schedule_ids).all() schedule_map = {s.id: s for s in schedules} statuses = [] for schedule_id in schedule_ids: if schedule := schedule_map.get(schedule_id): status_from_scheduler = APSchedulerAdapter.get_job_status(schedule.id) status_dict = { field: getattr(schedule, field) for field in schedule._meta.fields_map } status_dict.update(status_from_scheduler) status_dict["is_enabled"] = ( "运行中" if schedule_id in self._running_tasks else ("启用" if schedule.is_enabled else "暂停") ) statuses.append(status_dict) return statuses async def update_schedule( self, schedule_id: int, trigger_type: str | None = None, trigger_config: dict | None = None, job_kwargs: dict | None = None, ) -> tuple[bool, str]: """ 更新定时任务配置 """ schedule = await ScheduleRepository.get_by_id(schedule_id) if not schedule: return False, f"未找到 ID 为 {schedule_id} 的任务。" updated_fields = [] if trigger_config is not None: schedule.trigger_config = trigger_config updated_fields.append("trigger_config") if trigger_type is not None and schedule.trigger_type != trigger_type: schedule.trigger_type = trigger_type updated_fields.append("trigger_type") if job_kwargs is not None: existing_kwargs = ( schedule.job_kwargs.copy() if isinstance(schedule.job_kwargs, dict) else {} ) existing_kwargs.update(job_kwargs) is_valid, result = self._validate_and_prepare_kwargs( schedule.plugin_name, existing_kwargs ) if not is_valid: return False, str(result) assert isinstance(result, dict), "验证成功时 result 应该是字典类型" schedule.job_kwargs = result updated_fields.append("job_kwargs") if not updated_fields: return True, "没有任何需要更新的配置。" await ScheduleRepository.save(schedule, update_fields=updated_fields) APSchedulerAdapter.add_or_reschedule_job(schedule) return True, f"成功更新了任务 ID: {schedule_id} 的配置。" async def get_schedule_status(self, schedule_id: int) -> dict | None: """ 获取定时任务的详细状态信息 """ schedule = await ScheduleRepository.get_by_id(schedule_id) if not schedule: return None status_from_scheduler = APSchedulerAdapter.get_job_status(schedule.id) status_text = ( "运行中" if schedule_id in self._running_tasks else ("启用" if schedule.is_enabled else "暂停") ) return { "id": schedule.id, "bot_id": schedule.bot_id, "plugin_name": schedule.plugin_name, "target_type": schedule.target_type, "target_identifier": schedule.target_identifier, "is_enabled": status_text, "trigger_type": schedule.trigger_type, "trigger_config": schedule.trigger_config, "job_kwargs": schedule.job_kwargs, **status_from_scheduler, } async def pause_schedule(self, schedule_id: int) -> tuple[bool, str]: """ 暂停指定的定时任务 """ schedule = await ScheduleRepository.get_by_id(schedule_id) if not schedule or not schedule.is_enabled: return False, "任务不存在或已暂停。" schedule.is_enabled = False await ScheduleRepository.save(schedule, update_fields=["is_enabled"]) APSchedulerAdapter.pause_job(schedule_id) return True, f"已暂停任务 (ID: {schedule.id})。" async def resume_schedule(self, schedule_id: int) -> tuple[bool, str]: """ 恢复指定的定时任务 """ schedule = await ScheduleRepository.get_by_id(schedule_id) if not schedule or schedule.is_enabled: return False, "任务不存在或已启用。" schedule.is_enabled = True await ScheduleRepository.save(schedule, update_fields=["is_enabled"]) APSchedulerAdapter.resume_job(schedule_id) return True, f"已恢复任务 (ID: {schedule.id})。" async def trigger_now(self, schedule_id: int) -> tuple[bool, str]: """ 立即手动触发指定的定时任务 """ from .engine import _execute_job schedule = await ScheduleRepository.get_by_id(schedule_id) if not schedule: return False, f"未找到 ID 为 {schedule_id} 的定时任务。" if schedule.plugin_name not in self._registered_tasks: return False, f"插件 '{schedule.plugin_name}' 没有注册可用的定时任务。" try: await _execute_job(schedule.id, force=True) return True, f"已手动触发任务 (ID: {schedule.id})。" except Exception as e: logger.error(f"手动触发任务失败: {e}") return False, f"手动触发任务失败: {e}" async def get_schedule_by_id(self, schedule_id: int) -> "ScheduledJob | None": """ 通过ID获取任务对象的公共方法。 参数: schedule_id: 任务ID。 返回: ScheduledJob | None: 任务对象，不存在时返回None。 """ return await ScheduleRepository.get_by_id(schedule_id) scheduler_manager = SchedulerManager()