diff --git a/astrbot/core/core_lifecycle.py b/astrbot/core/core_lifecycle.py index 9961a8c52b..b79693159a 100644 --- a/astrbot/core/core_lifecycle.py +++ b/astrbot/core/core_lifecycle.py @@ -15,6 +15,8 @@ import time import traceback from asyncio import Queue +from collections.abc import Awaitable +from typing import Any from astrbot.api import logger, sp from astrbot.core import LogBroker, LogManager @@ -316,10 +318,13 @@ def _load(self) -> None: ) diagnostic_tasks = create_event_loop_diagnostic_tasks() - # 把插件中注册的所有协程函数注册到事件总线中并执行 - extra_tasks = [] + # Register every plugin-registered task on the event bus and run it + extra_tasks: list[asyncio.Task] = [] for task in self.star_context._register_tasks: - extra_tasks.append(asyncio.create_task(task, name=task.__name__)) # type: ignore + converted = self._to_task(task) + if converted is not None: + extra_tasks.append(converted) + self.star_context._register_tasks.clear() tasks_ = [ event_bus_task, @@ -337,6 +342,39 @@ def _load(self) -> None: self.start_time = int(time.time()) + @staticmethod + def _to_task(task: Awaitable[Any]) -> asyncio.Task | None: + """Convert a plugin-registered awaitable into an ``asyncio.Task``. + + ``Context.register_task`` accepts ``Awaitable``, so coroutines, + ``asyncio.Task``, ``asyncio.Future`` and any object implementing + ``__await__`` are valid inputs and must all be scheduled. + + ``_task_wrapper`` and ``stop`` rely on ``get_name()`` and ``cancel()`` + of ``asyncio.Task``, so any other valid awaitable is wrapped into a + real Task instead of being dropped silently. + + Args: + task: The awaitable registered through ``Context.register_task``. + + Returns: + A schedulable ``asyncio.Task``, or ``None`` if the input is not + awaitable. + """ + if isinstance(task, asyncio.Task): + return task + if asyncio.iscoroutine(task): + return asyncio.create_task(task, name=task.__name__) + if not isinstance(task, Awaitable): + logger.warning(f"Skipping non-awaitable plugin-registered task: {task!r}") + return None + + async def _await_registered() -> Any: + return await task + + name = getattr(task, "__name__", None) or type(task).__name__ + return asyncio.create_task(_await_registered(), name=name) + async def _task_wrapper(self, task: asyncio.Task) -> None: """异步任务包装器, 用于处理异步任务执行中出现的各种异常. diff --git a/tests/unit/test_core_lifecycle.py b/tests/unit/test_core_lifecycle.py index 058f869fd6..da58eac2e1 100644 --- a/tests/unit/test_core_lifecycle.py +++ b/tests/unit/test_core_lifecycle.py @@ -1000,3 +1000,146 @@ async def test_reload_pipeline_scheduler_raises_for_missing_config( with pytest.raises(ValueError, match="配置文件 .* 不存在"): await lifecycle.reload_pipeline_scheduler("nonexistent") + + +class TestAstrBotCoreLifecycleRegisteredTasks: + @staticmethod + def _make_lifecycle(mock_log_broker, mock_db, registered: list): + lifecycle = AstrBotCoreLifecycle(mock_log_broker, mock_db) + lifecycle.event_bus = MagicMock() + lifecycle.event_bus.dispatch = AsyncMock() + lifecycle.cron_manager = None + lifecycle.temp_dir_cleaner = None + lifecycle.star_context = MagicMock() + lifecycle.star_context._register_tasks = registered + lifecycle.curr_tasks = [] + return lifecycle + + @staticmethod + async def _cleanup(lifecycle: AstrBotCoreLifecycle): + for task in lifecycle.curr_tasks: + task.cancel() + for task in lifecycle.curr_tasks: + try: + await task + except (asyncio.CancelledError, Exception): + pass + + @pytest.mark.asyncio + async def test_load_schedules_coroutine_registered_task( + self, mock_log_broker, mock_db + ): + started = asyncio.Event() + + async def background(): + started.set() + + registered = [background()] + lifecycle = self._make_lifecycle(mock_log_broker, mock_db, registered) + + with patch( + "astrbot.core.core_lifecycle.create_event_loop_diagnostic_tasks", + return_value=[], + ): + lifecycle._load() + + await asyncio.wait_for(started.wait(), timeout=1) + assert registered == [] + assert len(lifecycle.curr_tasks) == 2 + + await self._cleanup(lifecycle) + + @pytest.mark.asyncio + async def test_load_reuses_already_created_task(self, mock_log_broker, mock_db): + started = asyncio.Event() + + async def background(): + started.set() + + existing = asyncio.create_task(background()) + lifecycle = self._make_lifecycle(mock_log_broker, mock_db, [existing]) + + with patch( + "astrbot.core.core_lifecycle.create_event_loop_diagnostic_tasks", + return_value=[], + ): + lifecycle._load() + + await asyncio.wait_for(started.wait(), timeout=1) + assert existing.get_name() in [t.get_name() for t in lifecycle.curr_tasks] + assert lifecycle.star_context._register_tasks == [] + + await self._cleanup(lifecycle) + + class _CustomAwaitable: + """Any object implementing __await__ is a valid Awaitable.""" + + def __await__(self): + async def _inner(): + return "done" + + return _inner().__await__() + + @pytest.mark.asyncio + async def test_load_skips_non_awaitable_registered_task( + self, mock_log_broker, mock_db + ): + lifecycle = self._make_lifecycle(mock_log_broker, mock_db, [object()]) + + with ( + patch( + "astrbot.core.core_lifecycle.create_event_loop_diagnostic_tasks", + return_value=[], + ), + patch("astrbot.core.core_lifecycle.logger") as mock_logger, + ): + lifecycle._load() + + mock_logger.warning.assert_called_once() + assert len(lifecycle.curr_tasks) == 1 + assert lifecycle.star_context._register_tasks == [] + + await self._cleanup(lifecycle) + + @pytest.mark.asyncio + async def test_load_schedules_future_registered_task(self, mock_log_broker, mock_db): + future = asyncio.get_running_loop().create_future() + lifecycle = self._make_lifecycle(mock_log_broker, mock_db, [future]) + + with patch( + "astrbot.core.core_lifecycle.create_event_loop_diagnostic_tasks", + return_value=[], + ): + lifecycle._load() + + # asyncio.Future is a valid Awaitable too and must be scheduled, not dropped + assert len(lifecycle.curr_tasks) == 2 + assert "Future" in [task.get_name() for task in lifecycle.curr_tasks] + assert lifecycle.star_context._register_tasks == [] + + future.set_result("done") + await asyncio.sleep(0) + await self._cleanup(lifecycle) + + @pytest.mark.asyncio + async def test_load_schedules_custom_awaitable_registered_task( + self, mock_log_broker, mock_db + ): + lifecycle = self._make_lifecycle(mock_log_broker, mock_db, [self._CustomAwaitable()]) + + with ( + patch( + "astrbot.core.core_lifecycle.create_event_loop_diagnostic_tasks", + return_value=[], + ), + patch("astrbot.core.core_lifecycle.logger") as mock_logger, + ): + lifecycle._load() + + assert len(lifecycle.curr_tasks) == 2 + assert "_CustomAwaitable" in [t.get_name() for t in lifecycle.curr_tasks] + mock_logger.warning.assert_not_called() + assert lifecycle.star_context._register_tasks == [] + + await asyncio.sleep(0) + await self._cleanup(lifecycle)