diff --git a/astrbot/core/utils/session_waiter.py b/astrbot/core/utils/session_waiter.py index b327a61843..636557b243 100644 --- a/astrbot/core/utils/session_waiter.py +++ b/astrbot/core/utils/session_waiter.py @@ -121,6 +121,9 @@ def __init__( self._lock = asyncio.Lock() """需要保证一个 session 同时只有一个 trigger""" + self._handler_task: asyncio.Task | None = None + """当前正在运行的 handler 任务""" + async def register_wait( self, handler: Callable[[SessionController, AstrMessageEvent], Awaitable[Any]], @@ -136,18 +139,29 @@ async def register_wait( try: return await self.session_controller.future except Exception as e: - self._cleanup(e) + await self._cleanup(e) raise e finally: - self._cleanup() + await self._cleanup() + + async def _cleanup(self, error: Exception | None = None) -> None: + """清理会话。 - def _cleanup(self, error: Exception | None = None) -> None: - """清理会话""" + 这是一个内部私有 async 方法,调用方必须 ``await``。保持 async 是因为取消 + ``_handler_task`` 需要在 ``self._lock`` 保护下进行以避免与 ``trigger`` 竞态; + 不应移除 async 或额外增加同步兼容层(过度工程)。 + """ USER_SESSIONS.pop(self.session_id, None) try: FILTERS.remove(self.session_filter) except ValueError: pass + # 在锁内取消任务:trigger 在锁内只做创建/取消任务的快速同步操作、不在锁内 + # await handler,且 trigger 锁外 await 期间已释放锁;因此 _cleanup 总能及时 + # 获得锁,两者不会互相阻塞等待对方持有的锁(不会死锁)。 + async with self._lock: + if self._handler_task and not self._handler_task.done(): + self._handler_task.cancel() self.session_controller.stop(error) @classmethod @@ -157,6 +171,7 @@ async def trigger(cls, session_id: str, event: AstrMessageEvent) -> None: if not session or session.session_controller.future.done(): return + task_to_await = None async with session._lock: if not session.session_controller.future.done(): if session.record_history_chains: @@ -164,11 +179,44 @@ async def trigger(cls, session_id: str, event: AstrMessageEvent) -> None: [copy.deepcopy(comp) for comp in event.get_messages()], ) try: - # TODO: 这里使用 create_task,跟踪 task,防止超时后这里 handler 仍然在执行 + # 使用 create_task 跟踪任务,防止超时后 handler 仍然在执行 assert session.handler is not None - await session.handler(session.session_controller, event) + + async def _run_handler(): + # 内层异常处理:把 handler 的普通异常转为 stop(e) 写入 future, + # 使 register_wait 能感知;CancelledError 则重新抛出交由外层处理。 + try: + await session.handler(session.session_controller, event) + except asyncio.CancelledError: + raise # 重新抛出,让外层的 except asyncio.CancelledError 处理 + except Exception as e: + session.session_controller.stop(e) + + # 取消上一个可能还在运行的任务 + if session._handler_task and not session._handler_task.done(): + session._handler_task.cancel() + + session._handler_task = asyncio.create_task(_run_handler()) + task_to_await = session._handler_task except Exception as e: session.session_controller.stop(e) + return + + # 在锁外等待任务完成,避免死锁:trigger 在锁内只创建任务、不在锁内 await + # handler,锁已在进入此处前释放;而 _cleanup 仅在锁内取消任务,因此两者不会 + # 互相阻塞等待对方持有的锁。 + if task_to_await: + try: + await task_to_await + except asyncio.CancelledError: + # 外层只静默处理任务被取消的情况:取消是预期行为(如被 _cleanup 或 + # 后续 trigger 取消),不应上报为异常,也不应写入 future。 + pass + finally: + # 清理已完成任务的引用,避免已完成任务残留在 _handler_task 上。 + # 仅当 _handler_task 仍是当前任务时才置 None,防止并发 trigger 已替换它。 + if session._handler_task is task_to_await: + session._handler_task = None def session_waiter(timeout: int = 30, record_history_chains: bool = False): diff --git a/tests/unit/test_session_waiter.py b/tests/unit/test_session_waiter.py new file mode 100644 index 0000000000..50b0ce0925 --- /dev/null +++ b/tests/unit/test_session_waiter.py @@ -0,0 +1,326 @@ +"""Tests for astrbot.core.utils.session_waiter. + +These tests lock in the expected behavior of the session-waiter task-tracking +mechanism introduced in PR #9458: a handler task that is still running after a +timeout (or after being superseded by a second trigger) must be cancelled, and +the resulting ``asyncio.CancelledError`` must NOT be funneled into +``SessionController.stop(error)`` as if it were an ordinary handler exception. + +The tests encode the *expected* behavior. They must not be weakened to satisfy +the current implementation. +""" + +import asyncio +from contextlib import suppress +from unittest.mock import MagicMock + +import pytest + +from astrbot.core.utils.session_waiter import ( + FILTERS, + USER_SESSIONS, + DefaultSessionFilter, + SessionWaiter, +) + + +def make_event(unified_msg_origin: str = "test_umo") -> MagicMock: + """Build a minimal AstrMessageEvent mock for session_waiter tests.""" + event = MagicMock() + event.unified_msg_origin = unified_msg_origin + event.get_messages.return_value = [] + return event + + +async def _wait_for_session(session_id: str, timeout: float = 5.0) -> None: + """Wait until ``session_id`` appears in USER_SESSIONS (register_wait started).""" + loop = asyncio.get_running_loop() + deadline = loop.time() + timeout + while session_id not in USER_SESSIONS: + if loop.time() >= deadline: + raise AssertionError( + f"Session {session_id!r} was not registered within {timeout}s" + ) + await asyncio.sleep(0.01) + + +@pytest.fixture(autouse=True) +def reset_session_globals(): + """Clear module-level USER_SESSIONS / FILTERS before and after each test.""" + USER_SESSIONS.clear() + FILTERS.clear() + yield + USER_SESSIONS.clear() + FILTERS.clear() + + +class TestSessionWaiterBasicFlow: + """基础流程:trigger 触发 handler 执行;handler 调 stop() 后 register_wait 返回。""" + + @pytest.mark.asyncio + async def test_trigger_runs_handler_and_register_wait_returns(self): + handler_executed = asyncio.Event() + + async def handler(controller, event): # noqa: ARG001 + handler_executed.set() + controller.stop() + + waiter = SessionWaiter(DefaultSessionFilter(), "session_basic", False) + register_task = asyncio.create_task(waiter.register_wait(handler, timeout=30)) + await _wait_for_session("session_basic") + + event = make_event("session_basic") + await SessionWaiter.trigger("session_basic", event) + + assert handler_executed.is_set(), "handler should have been executed by trigger" + + result = await asyncio.wait_for(register_task, timeout=5) + assert register_task.done() + assert not register_task.cancelled() + assert result is None + + +class TestSessionWaiterDoubleTrigger: + """重复 trigger 取消上一个仍在运行的 handler。""" + + @pytest.mark.asyncio + async def test_second_trigger_cancels_previous_handler(self): + first_handler_started = asyncio.Event() + first_handler_cancelled = asyncio.Event() + call_count = 0 + + async def handler(controller, event): # noqa: ARG001 + nonlocal call_count + call_count += 1 + if call_count == 1: + first_handler_started.set() + try: + await asyncio.sleep(5) + except asyncio.CancelledError: + first_handler_cancelled.set() + raise + else: + controller.stop() + + waiter = SessionWaiter(DefaultSessionFilter(), "session_double", False) + register_task = asyncio.create_task(waiter.register_wait(handler, timeout=30)) + await _wait_for_session("session_double") + + event1 = make_event("session_double") + trigger1_task = asyncio.create_task( + SessionWaiter.trigger("session_double", event1) + ) + await asyncio.wait_for(first_handler_started.wait(), timeout=5) + + # Previous handler is still running -> second trigger must cancel it. + event2 = make_event("session_double") + await SessionWaiter.trigger("session_double", event2) + + await asyncio.wait_for(first_handler_cancelled.wait(), timeout=5) + assert first_handler_cancelled.is_set(), ( + "previous handler task should have been cancelled by the second trigger" + ) + + # trigger1 was awaiting the cancelled task; it should swallow the + # CancelledError and return cleanly. + await asyncio.wait_for(trigger1_task, timeout=5) + assert trigger1_task.done() + assert not trigger1_task.cancelled() + + # Second handler called stop() -> register_wait returns normally. + await asyncio.wait_for(register_task, timeout=5) + assert register_task.done() + assert not register_task.cancelled() + + +class TestSessionWaiterCleanupCancelsHandler: + """超时/清理取消运行中的 handler。""" + + @pytest.mark.asyncio + async def test_cleanup_cancels_running_handler(self): + handler_started = asyncio.Event() + handler_cancelled = asyncio.Event() + + async def handler(controller, event): # noqa: ARG001 + handler_started.set() + try: + await asyncio.sleep(5) + except asyncio.CancelledError: + handler_cancelled.set() + raise + + waiter = SessionWaiter(DefaultSessionFilter(), "session_cleanup", False) + register_task = asyncio.create_task(waiter.register_wait(handler, timeout=30)) + await _wait_for_session("session_cleanup") + + event = make_event("session_cleanup") + trigger_task = asyncio.create_task( + SessionWaiter.trigger("session_cleanup", event) + ) + await asyncio.wait_for(handler_started.wait(), timeout=5) + + # Handler is running. _cleanup() is what register_wait calls on timeout + # / finally; invoking it directly simulates that path. + await waiter._cleanup() + + await asyncio.wait_for(handler_cancelled.wait(), timeout=5) + assert handler_cancelled.is_set(), ( + "running handler task should be cancelled by _cleanup" + ) + + # trigger was awaiting the now-cancelled handler task; it returns. + await asyncio.wait_for(trigger_task, timeout=5) + assert trigger_task.done() + assert not trigger_task.cancelled() + + # _cleanup set the future result -> register_wait completes. + await asyncio.wait_for(register_task, timeout=5) + assert register_task.done() + assert not register_task.cancelled() + + +class TestSessionWaiterCancelledErrorHandling: + """CancelledError 不被当作普通异常处理(PR 核心修复点)。 + + 当 handler 被取消时,CancelledError 不得通过 stop(e) 被当作普通异常设置到 + future 上。验证:stop() 从未被以 CancelledError 为参数调用;且若 future 已完成, + 其 exception() 绝不是 CancelledError。 + """ + + def _install_stop_spy(self, waiter: SessionWaiter) -> list: + """Record every error passed to SessionController.stop().""" + original_stop = waiter.session_controller.stop + stop_calls: list = [] + + def stop_spy(error=None): + stop_calls.append(error) + return original_stop(error) + + waiter.session_controller.stop = stop_spy + return stop_calls + + @pytest.mark.asyncio + async def test_stop_not_called_with_cancelled_error_on_cleanup(self): + handler_started = asyncio.Event() + handler_cancelled = asyncio.Event() + + async def handler(controller, event): # noqa: ARG001 + handler_started.set() + try: + await asyncio.sleep(5) + except asyncio.CancelledError: + handler_cancelled.set() + raise + + waiter = SessionWaiter(DefaultSessionFilter(), "session_cancel_cleanup", False) + stop_calls = self._install_stop_spy(waiter) + + register_task = asyncio.create_task(waiter.register_wait(handler, timeout=30)) + await _wait_for_session("session_cancel_cleanup") + + event = make_event("session_cancel_cleanup") + trigger_task = asyncio.create_task( + SessionWaiter.trigger("session_cancel_cleanup", event) + ) + await asyncio.wait_for(handler_started.wait(), timeout=5) + + await waiter._cleanup() + await asyncio.wait_for(handler_cancelled.wait(), timeout=5) + + # Core: stop() must never receive a CancelledError. + for err in stop_calls: + assert not isinstance(err, asyncio.CancelledError), ( + f"stop() must not be called with CancelledError, got {err!r}" + ) + + # Core: a completed future must not carry a CancelledError. + future = waiter.session_controller.future + if future.done(): + exc = future.exception() + assert not isinstance(exc, asyncio.CancelledError), ( + f"future exception must not be CancelledError, got {exc!r}" + ) + + await asyncio.wait_for(trigger_task, timeout=5) + await asyncio.wait_for(register_task, timeout=5) + + @pytest.mark.asyncio + async def test_stop_not_called_with_cancelled_error_on_double_trigger(self): + first_handler_started = asyncio.Event() + first_handler_cancelled = asyncio.Event() + call_count = 0 + + async def handler(controller, event): # noqa: ARG001 + nonlocal call_count + call_count += 1 + if call_count == 1: + first_handler_started.set() + try: + await asyncio.sleep(5) + except asyncio.CancelledError: + first_handler_cancelled.set() + raise + else: + controller.stop() + + waiter = SessionWaiter(DefaultSessionFilter(), "session_cancel_dt", False) + stop_calls = self._install_stop_spy(waiter) + + register_task = asyncio.create_task(waiter.register_wait(handler, timeout=30)) + await _wait_for_session("session_cancel_dt") + + event1 = make_event("session_cancel_dt") + trigger1_task = asyncio.create_task( + SessionWaiter.trigger("session_cancel_dt", event1) + ) + await asyncio.wait_for(first_handler_started.wait(), timeout=5) + + event2 = make_event("session_cancel_dt") + await SessionWaiter.trigger("session_cancel_dt", event2) + await asyncio.wait_for(first_handler_cancelled.wait(), timeout=5) + + for err in stop_calls: + assert not isinstance(err, asyncio.CancelledError), ( + f"stop() must not be called with CancelledError, got {err!r}" + ) + + future = waiter.session_controller.future + if future.done(): + exc = future.exception() + assert not isinstance(exc, asyncio.CancelledError), ( + f"future exception must not be CancelledError, got {exc!r}" + ) + + await asyncio.wait_for(trigger1_task, timeout=5) + await asyncio.wait_for(register_task, timeout=5) + + +class TestSessionWaiterHandlerException: + """handler 普通异常向上传播。""" + + @pytest.mark.asyncio + async def test_handler_exception_propagates_to_register_wait(self): + handler_executed = asyncio.Event() + + async def handler(controller, event): # noqa: ARG001 + handler_executed.set() + raise ValueError("handler error") + + waiter = SessionWaiter(DefaultSessionFilter(), "session_exc", False) + register_task = asyncio.create_task(waiter.register_wait(handler, timeout=30)) + await _wait_for_session("session_exc") + + event = make_event("session_exc") + await SessionWaiter.trigger("session_exc", event) + + assert handler_executed.is_set() + + # register_wait re-raises the ValueError surfaced via stop(e). + with pytest.raises(ValueError, match="handler error"): + await asyncio.wait_for(register_task, timeout=5) + + # The future itself carries the ValueError. + assert waiter.session_controller.future.done() + exc = waiter.session_controller.future.exception() + assert isinstance(exc, ValueError) + assert "handler error" in str(exc)