From f963bfbd9b9e4288fd1795b991530aa05814878c Mon Sep 17 00:00:00 2001 From: Qixuan112 <3573568193@qq.com> Date: Thu, 30 Jul 2026 14:49:54 +0800 Subject: [PATCH 1/4] fix(session_waiter): track handler task to prevent execution after timeout - Add _handler_task attribute to SessionWaiter to track running handler tasks - Use asyncio.create_task to wrap handler execution for proper task management - Cancel any running handler task during cleanup to prevent orphaned execution - Cancel previous handler task if a new trigger arrives before completion - Handle asyncio.CancelledError explicitly to avoid error propagation This resolves the TODO at line 167 and ensures that handler tasks are properly managed and cancelled when sessions timeout or cleanup, preventing potential resource leaks and unexpected behavior from handlers that continue running after their session has ended. --- astrbot/core/utils/session_waiter.py | 25 +++++++++++++++++++++++-- 1 file changed, 23 insertions(+), 2 deletions(-) diff --git a/astrbot/core/utils/session_waiter.py b/astrbot/core/utils/session_waiter.py index b327a61843..4c8764b563 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]], @@ -148,6 +151,9 @@ def _cleanup(self, error: Exception | None = None) -> None: FILTERS.remove(self.session_filter) except ValueError: pass + # 取消正在运行的 handler 任务 + if self._handler_task and not self._handler_task.done(): + self._handler_task.cancel() self.session_controller.stop(error) @classmethod @@ -164,9 +170,24 @@ 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(): + try: + await session.handler(session.session_controller, event) + 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()) + await session._handler_task + except asyncio.CancelledError: + # 任务被取消时不需要处理 + pass except Exception as e: session.session_controller.stop(e) From b3ff355528511db05e3960d5cd9ab1da488ec07c Mon Sep 17 00:00:00 2001 From: Qixuan112 <3573568193@qq.com> Date: Thu, 30 Jul 2026 15:40:40 +0800 Subject: [PATCH 2/4] fix: re-raise CancelledError in _run_handler to avoid treating cancellations as errors --- astrbot/core/utils/session_waiter.py | 30 ++++++++++++++++++---------- 1 file changed, 20 insertions(+), 10 deletions(-) diff --git a/astrbot/core/utils/session_waiter.py b/astrbot/core/utils/session_waiter.py index 4c8764b563..5dab5ae5fd 100644 --- a/astrbot/core/utils/session_waiter.py +++ b/astrbot/core/utils/session_waiter.py @@ -139,21 +139,22 @@ 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() - def _cleanup(self, error: Exception | None = None) -> None: + async def _cleanup(self, error: Exception | None = None) -> None: """清理会话""" USER_SESSIONS.pop(self.session_id, None) try: FILTERS.remove(self.session_filter) except ValueError: pass - # 取消正在运行的 handler 任务 - if self._handler_task and not self._handler_task.done(): - self._handler_task.cancel() + # 使用锁保护任务取消操作,防止与 trigger 方法竞态 + async with self._lock: + if self._handler_task and not self._handler_task.done(): + self._handler_task.cancel() self.session_controller.stop(error) @classmethod @@ -163,6 +164,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: @@ -176,6 +178,8 @@ async def trigger(cls, session_id: str, event: AstrMessageEvent) -> None: async def _run_handler(): try: await session.handler(session.session_controller, event) + except asyncio.CancelledError: + raise # 重新抛出,让外层的 except asyncio.CancelledError 处理 except Exception as e: session.session_controller.stop(e) @@ -184,12 +188,18 @@ async def _run_handler(): session._handler_task.cancel() session._handler_task = asyncio.create_task(_run_handler()) - await session._handler_task - except asyncio.CancelledError: - # 任务被取消时不需要处理 - pass + task_to_await = session._handler_task except Exception as e: session.session_controller.stop(e) + return + + # 在锁外等待任务完成,避免死锁 + if task_to_await: + try: + await task_to_await + except asyncio.CancelledError: + # 任务被取消时不需要处理 + pass def session_waiter(timeout: int = 30, record_history_chains: bool = False): From 7379572d4349c9f402f0df4763c29bdc35bf0a68 Mon Sep 17 00:00:00 2001 From: Qixuan112 <3573568193@qq.com> Date: Fri, 31 Jul 2026 20:32:26 +0800 Subject: [PATCH 3/4] test(session_waiter): add unit tests for handler task cancellation Add locked unit tests covering the task-tracking mechanism introduced in #9458, addressing review feedback that cancellation behavior was untested: - Basic trigger flow: handler runs and register_wait returns - Double trigger cancels the previous still-running handler task - _cleanup cancels a running handler task (timeout path) - CancelledError is NOT funneled into SessionController.stop() as an ordinary exception (cleanup and double-trigger paths) - Handler exceptions propagate to the future via stop(e) These tests encode the expected behavior and serve as the acceptance contract for future changes. --- tests/unit/test_session_waiter.py | 326 ++++++++++++++++++++++++++++++ 1 file changed, 326 insertions(+) create mode 100644 tests/unit/test_session_waiter.py 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) From 30aa3abe8249073a09a5836e7b81a305eb76d051 Mon Sep 17 00:00:00 2001 From: Qixuan112 <3573568193@qq.com> Date: Fri, 31 Jul 2026 20:32:27 +0800 Subject: [PATCH 4/4] fix(session_waiter): address review feedback on task tracking Respond to @xiaoyuyu6420 and sourcery review comments on #9458: - Document _cleanup as an internal async method (it must acquire _lock to safely cancel _handler_task); confirmed no external callers exist - Add comments explaining why awaiting the handler task outside the lock cannot deadlock with _cleanup - Add comments explaining the two-layer exception handling: inner layer converts handler exceptions to stop(e), outer layer silently handles task cancellation - Clear _handler_task reference after the task completes (guarded by an identity check to avoid clobbering a concurrent trigger's task) No control-flow changes; behavior is unchanged and all unit tests pass. --- astrbot/core/utils/session_waiter.py | 25 +++++++++++++++++++++---- 1 file changed, 21 insertions(+), 4 deletions(-) diff --git a/astrbot/core/utils/session_waiter.py b/astrbot/core/utils/session_waiter.py index 5dab5ae5fd..636557b243 100644 --- a/astrbot/core/utils/session_waiter.py +++ b/astrbot/core/utils/session_waiter.py @@ -145,13 +145,20 @@ async def register_wait( await self._cleanup() async 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 方法竞态 + # 在锁内取消任务:trigger 在锁内只做创建/取消任务的快速同步操作、不在锁内 + # await handler,且 trigger 锁外 await 期间已释放锁;因此 _cleanup 总能及时 + # 获得锁,两者不会互相阻塞等待对方持有的锁(不会死锁)。 async with self._lock: if self._handler_task and not self._handler_task.done(): self._handler_task.cancel() @@ -176,6 +183,8 @@ async def trigger(cls, session_id: str, event: AstrMessageEvent) -> None: assert session.handler is not None async def _run_handler(): + # 内层异常处理:把 handler 的普通异常转为 stop(e) 写入 future, + # 使 register_wait 能感知;CancelledError 则重新抛出交由外层处理。 try: await session.handler(session.session_controller, event) except asyncio.CancelledError: @@ -193,13 +202,21 @@ async def _run_handler(): 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):