From abe8b639aec166201ef547d691b85d98a51b498a Mon Sep 17 00:00:00 2001 From: QingFeng-awa Date: Wed, 5 Aug 2026 01:44:01 +0800 Subject: [PATCH 1/2] feat: allow session_waiter handler to return MessageEventResult Let session_waiter handlers return a MessageEventResult instead of forcing them to call event.send(). The returned result is yielded back into the pipeline by handle_session_control_agent, so it goes through the full ResultDecorateStage (reply-with-mention, reply-with-quote, segmented reply, TTS, text-to-image, etc.) and RespondStage like any normal reply. Backward compatible: handlers that return None keep the original stop_event behavior. --- astrbot/builtin_stars/astrbot/main.py | 8 +- astrbot/core/utils/session_waiter.py | 7 +- tests/unit/test_session_control_agent.py | 122 ++++++++++++++++++++ tests/unit/test_session_waiter.py | 136 +++++++++++++++++++++++ 4 files changed, 267 insertions(+), 6 deletions(-) create mode 100644 tests/unit/test_session_control_agent.py create mode 100644 tests/unit/test_session_waiter.py diff --git a/astrbot/builtin_stars/astrbot/main.py b/astrbot/builtin_stars/astrbot/main.py index 571fa000c3..78af164920 100644 --- a/astrbot/builtin_stars/astrbot/main.py +++ b/astrbot/builtin_stars/astrbot/main.py @@ -9,7 +9,7 @@ from astrbot.api.message_components import Image, Plain from astrbot.api.provider import LLMResponse, ProviderRequest from astrbot.core import logger -from astrbot.core.message.message_event_result import MessageChain +from astrbot.core.message.message_event_result import MessageChain, MessageEventResult from astrbot.core.platform.message_type import MessageType from astrbot.core.utils.session_waiter import ( FILTERS, @@ -42,12 +42,14 @@ def __init__(self, context: star.Context) -> None: logger.error(f"group chat context init failed: {e}") @filter.event_message_type(filter.EventMessageType.ALL, priority=maxsize) - async def handle_session_control_agent(self, event: AstrMessageEvent) -> None: + async def handle_session_control_agent(self, event: AstrMessageEvent): """会话控制代理""" for session_filter in FILTERS: session_id = session_filter.filter(event) if session_id in USER_SESSIONS: - await SessionWaiter.trigger(session_id, event) + result = await SessionWaiter.trigger(session_id, event) + if isinstance(result, MessageEventResult): + yield result event.stop_event() @filter.event_message_type(filter.EventMessageType.ALL, priority=maxsize - 1) diff --git a/astrbot/core/utils/session_waiter.py b/astrbot/core/utils/session_waiter.py index b327a61843..487dccd613 100644 --- a/astrbot/core/utils/session_waiter.py +++ b/astrbot/core/utils/session_waiter.py @@ -151,11 +151,11 @@ def _cleanup(self, error: Exception | None = None) -> None: self.session_controller.stop(error) @classmethod - async def trigger(cls, session_id: str, event: AstrMessageEvent) -> None: + async def trigger(cls, session_id: str, event: AstrMessageEvent) -> Any: """外部输入触发会话处理""" session = USER_SESSIONS.get(session_id) if not session or session.session_controller.future.done(): - return + return None async with session._lock: if not session.session_controller.future.done(): @@ -166,9 +166,10 @@ async def trigger(cls, session_id: str, event: AstrMessageEvent) -> None: try: # TODO: 这里使用 create_task，跟踪 task，防止超时后这里 handler 仍然在执行 assert session.handler is not None - await session.handler(session.session_controller, event) + return await session.handler(session.session_controller, event) except Exception as e: session.session_controller.stop(e) + return None def session_waiter(timeout: int = 30, record_history_chains: bool = False): diff --git a/tests/unit/test_session_control_agent.py b/tests/unit/test_session_control_agent.py new file mode 100644 index 0000000000..e3bf97a3a3 --- /dev/null +++ b/tests/unit/test_session_control_agent.py @@ -0,0 +1,122 @@ +"""Tests for Main.handle_session_control_agent result routing. + +Covers the two branches introduced by the session-waiter return-value change: +- handler returns a MessageEventResult -> yield it, then stop_event +- handler returns None -> no yield, just stop_event +""" + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from astrbot.core.message.message_event_result import MessageEventResult +from astrbot.core.platform.astr_message_event import AstrMessageEvent +from astrbot.core.platform.astrbot_message import AstrBotMessage, MessageMember +from astrbot.core.platform.message_type import MessageType +from astrbot.core.platform.platform_metadata import PlatformMetadata + + +class _ConcreteEvent(AstrMessageEvent): + async def send(self, message): + await super().send(message) + + +def _make_event() -> _ConcreteEvent: + meta = PlatformMetadata(name="test", description="t", id="test_id") + msg = AstrBotMessage() + msg.type = MessageType.FRIEND_MESSAGE + msg.self_id = "bot" + msg.session_id = "s" + msg.message_id = "m" + msg.sender = MessageMember(user_id="u", nickname="U") + msg.message_str = "hi" + return _ConcreteEvent("hi", msg, meta, "s") + + +def _build_main(): + """Construct a Main instance without running its heavy __init__.""" + from astrbot.builtin_stars.astrbot.main import Main + + main = Main.__new__(Main) + return main + + +async def _collect(gen): + """Drain an async generator into a list of yielded values.""" + out = [] + async for item in gen: + out.append(item) + return out + + +@pytest.mark.asyncio +async def test_yields_message_event_result_when_trigger_returns_one(): + """When trigger returns a MessageEventResult, the handler should yield it.""" + main = _build_main() + event = _make_event() + expected = MessageEventResult().message("先见之明") + + fake_filter = MagicMock() + fake_filter.filter = MagicMock(return_value="sid") + with ( + patch( + "astrbot.builtin_stars.astrbot.main.FILTERS", + [fake_filter], + ), + patch( + "astrbot.builtin_stars.astrbot.main.USER_SESSIONS", + {"sid": MagicMock()}, + ), + patch( + "astrbot.builtin_stars.astrbot.main.SessionWaiter.trigger", + new=AsyncMock(return_value=expected), + ), + ): + yielded = await _collect(main.handle_session_control_agent(event)) + + assert yielded == [expected] + assert event.is_stopped() is True + + +@pytest.mark.asyncio +async def test_no_yield_when_trigger_returns_none(): + """When trigger returns None, the handler should not yield and should stop.""" + main = _build_main() + event = _make_event() + + fake_filter = MagicMock() + fake_filter.filter = MagicMock(return_value="sid") + with ( + patch( + "astrbot.builtin_stars.astrbot.main.FILTERS", + [fake_filter], + ), + patch( + "astrbot.builtin_stars.astrbot.main.USER_SESSIONS", + {"sid": MagicMock()}, + ), + patch( + "astrbot.builtin_stars.astrbot.main.SessionWaiter.trigger", + new=AsyncMock(return_value=None), + ), + ): + yielded = await _collect(main.handle_session_control_agent(event)) + + assert yielded == [] + assert event.is_stopped() is True + + +@pytest.mark.asyncio +async def test_no_op_when_no_session_matches(): + """When no session_filter matches, nothing happens and event is not stopped.""" + main = _build_main() + event = _make_event() + + with ( + patch("astrbot.builtin_stars.astrbot.main.FILTERS", []), + patch("astrbot.builtin_stars.astrbot.main.USER_SESSIONS", {}), + ): + yielded = await _collect(main.handle_session_control_agent(event)) + + assert yielded == [] + assert event.is_stopped() is False diff --git a/tests/unit/test_session_waiter.py b/tests/unit/test_session_waiter.py new file mode 100644 index 0000000000..5e3a7d029e --- /dev/null +++ b/tests/unit/test_session_waiter.py @@ -0,0 +1,136 @@ +"""Tests for session_waiter.trigger return value propagation.""" + +import asyncio +from unittest.mock import AsyncMock + +import pytest + +from astrbot.core.message.message_event_result import MessageEventResult +from astrbot.core.platform.astr_message_event import AstrMessageEvent +from astrbot.core.platform.astrbot_message import AstrBotMessage, MessageMember +from astrbot.core.platform.message_type import MessageType +from astrbot.core.platform.platform_metadata import PlatformMetadata +from astrbot.core.utils.session_waiter import ( + FILTERS, + USER_SESSIONS, + SessionController, + SessionFilter, + SessionWaiter, + session_waiter, +) + +SESSION_ID = "fixed_session_id" + + +class _ConcreteEvent(AstrMessageEvent): + async def send(self, message): + await super().send(message) + + +class _IdFilter(SessionFilter): + def filter(self, event: AstrMessageEvent) -> str: + return SESSION_ID + + +def _make_event() -> _ConcreteEvent: + meta = PlatformMetadata(name="test", description="t", id="test_id") + msg = AstrBotMessage() + msg.type = MessageType.FRIEND_MESSAGE + msg.self_id = "bot" + msg.session_id = "s" + msg.message_id = "m" + msg.sender = MessageMember(user_id="u", nickname="U") + msg.message_str = "hi" + return _ConcreteEvent("hi", msg, meta, "s") + + +@pytest.fixture +def waiter_handler(): + """Return a mutable holder so each test can install its own handler. + + Also cleans up the module-global USER_SESSIONS / FILTERS so tests stay + isolated even if the waiter task is interrupted. + """ + + holder = {} + + @session_waiter(timeout=5, record_history_chains=False) + async def waiter(controller: SessionController, event: AstrMessageEvent): + handler = holder["handler"] + if handler is None: + return None + return await handler(controller, event) + + holder["waiter"] = waiter + holder["handler"] = None + yield holder + # Teardown: stop any still-pending session and clear globals. + for session in list(USER_SESSIONS.values()): + session.session_controller.stop() + USER_SESSIONS.clear() + FILTERS.clear() + + +async def _start_waiter(holder) -> asyncio.Task: + """Register the waiter (which blocks on the controller future) in background.""" + event = _make_event() + task = asyncio.create_task(holder["waiter"](event, session_filter=_IdFilter())) + # Yield control so the waiter registers itself in USER_SESSIONS. + await asyncio.sleep(0) + return task + + +@pytest.mark.asyncio +async def test_trigger_returns_message_event_result(waiter_handler): + """trigger should propagate the handler's returned MessageEventResult.""" + expected = MessageEventResult().message("先见之明") + waiter_handler["handler"] = AsyncMock(return_value=expected) + + task = await _start_waiter(waiter_handler) + result = await SessionWaiter.trigger(SESSION_ID, _make_event()) + + assert result is expected + # register_wait is still pending; stop it cleanly. + for session in list(USER_SESSIONS.values()): + session.session_controller.stop() + await task + + +@pytest.mark.asyncio +async def test_trigger_returns_none_when_handler_returns_none(waiter_handler): + """trigger should return None when the handler returns None.""" + waiter_handler["handler"] = AsyncMock(return_value=None) + + task = await _start_waiter(waiter_handler) + result = await SessionWaiter.trigger(SESSION_ID, _make_event()) + + assert result is None + for session in list(USER_SESSIONS.values()): + session.session_controller.stop() + await task + + +@pytest.mark.asyncio +async def test_trigger_returns_none_when_no_session(): + """trigger should return None when there is no registered session.""" + event = _make_event() + result = await SessionWaiter.trigger("nonexistent_session", event) + assert result is None + + +@pytest.mark.asyncio +async def test_trigger_returns_none_when_handler_raises(waiter_handler): + """trigger should return None (and stop the controller) when the handler raises. + + The still-pending register_wait task will re-raise the exception after we + stop the controller, so we swallow it explicitly. + """ + waiter_handler["handler"] = AsyncMock(side_effect=RuntimeError("boom")) + + task = await _start_waiter(waiter_handler) + result = await SessionWaiter.trigger(SESSION_ID, _make_event()) + + assert result is None + # The controller was stopped with the exception; register_wait re-raises it. + with pytest.raises(RuntimeError, match="boom"): + await task From b31ffb5744d53881115d25730894dd1c3c70448c Mon Sep 17 00:00:00 2001 From: QingFeng-awa Date: Wed, 5 Aug 2026 03:10:17 +0800 Subject: [PATCH 2/2] docs(session-control): update session-control guide for return MessageEventResult --- docs/en/dev/star/guides/session-control.md | 29 ++++++++------------- docs/zh/dev/star/guides/session-control.md | 30 +++++++--------------- 2 files changed, 19 insertions(+), 40 deletions(-) diff --git a/docs/en/dev/star/guides/session-control.md b/docs/en/dev/star/guides/session-control.md index e08bae7ae1..b718077857 100644 --- a/docs/en/dev/star/guides/session-control.md +++ b/docs/en/dev/star/guides/session-control.md @@ -20,15 +20,7 @@ Import: ```py import astrbot.api.message_components as Comp -from astrbot.core.utils.session_waiter import ( - session_waiter, - SessionController, -) -``` - -Code within the handler can be written as follows: - -```python +from astrbot.api.utils import SessionController, session_waiter from astrbot.api.event import filter, AstrMessageEvent @filter.command("idiom-chain") @@ -40,27 +32,22 @@ async def handle_empty_mention(self, event: AstrMessageEvent): # How to use the session controller @session_waiter(timeout=60, record_history_chains=False) # Register a session controller with a 60-second timeout, without recording message history async def empty_mention_waiter(controller: SessionController, event: AstrMessageEvent): - idiom = event.message_str # The idiom sent by the user, e.g., "one horse takes the lead" + idiom = event.message_str # The idiom sent by the user, e.g., "一马当先" if idiom == "exit": # If the user wants to exit the idiom chain game by typing "exit" - await event.send(event.plain_result("Exited the idiom chain game~")) controller.stop() # Stop the session controller, which will end immediately. - return + return event.plain_result("Exited the idiom chain game~") if len(idiom) != 4: # If the user's input is not a 4-character idiom - await event.send(event.plain_result("The idiom must be four characters~")) # Send a reply, cannot use yield - return # Exit the current method without executing subsequent logic, but the session is not interrupted; subsequent user input will still enter the current session + return event.plain_result("The idiom must be four characters~") # ... - message_result = event.make_result() - message_result.chain = [Comp.Plain("Foresight")] # import astrbot.api.message_components as Comp - await event.send(message_result) # Send a reply, cannot use yield - controller.keep(timeout=60, reset_timeout=True) # Reset timeout to 60s. If not reset, it will continue the previous timeout countdown. # controller.stop() # Stop the session controller, which will end immediately. # If history chains are recorded, you can retrieve them via controller.get_history_chains() + return event.plain_result("先见之明") # Send a reply try: await empty_mention_waiter(event) @@ -74,6 +61,10 @@ async def handle_empty_mention(self, event: AstrMessageEvent): logger.error("handle_empty_mention error: " + str(e)) ``` +> [!TIP] +> Inside the session controller handler, you **cannot use `yield`**, but you can `return` a `MessageEventResult`. The returned result is automatically sent through the framework's message decoration pipeline, behaving the same as a regular command reply.\ +> If you call `event.send()` directly, the message bypasses the decoration pipeline and is sent directly. + Once the session controller is activated, messages subsequently sent by that sender will first be processed by the `empty_mention_waiter` function you defined above, until the session controller is stopped or times out. ## SessionController @@ -92,7 +83,7 @@ By default, the AstrBot session controller uses `sender_id` (the sender's ID) as ```py import astrbot.api.message_components as Comp -from astrbot.core.utils.session_waiter import ( +from astrbot.api.utils import ( session_waiter, SessionFilter, SessionController, diff --git a/docs/zh/dev/star/guides/session-control.md b/docs/zh/dev/star/guides/session-control.md index beaea69c61..b4359fdb6e 100644 --- a/docs/zh/dev/star/guides/session-control.md +++ b/docs/zh/dev/star/guides/session-control.md @@ -16,19 +16,9 @@ AstrBot 提供了开箱即用的会话控制功能： -导入： - ```py import astrbot.api.message_components as Comp -from astrbot.core.utils.session_waiter import ( - session_waiter, - SessionController, -) -``` - -handler 内的代码可以如下： - -```python +from astrbot.api.utils import SessionController, session_waiter from astrbot.api.event import filter, AstrMessageEvent @filter.command("成语接龙") @@ -43,24 +33,18 @@ async def handle_empty_mention(self, event: AstrMessageEvent): idiom = event.message_str # 用户发来的成语，假设是 "一马当先" if idiom == "退出": # 假设用户想主动退出成语接龙，输入了 "退出" - await event.send(event.plain_result("已退出成语接龙~")) controller.stop() # 停止会话控制器，会立即结束。 - return + return event.plain_result("已退出成语接龙~") if len(idiom) != 4: # 假设用户输入的不是4字成语 - await event.send(event.plain_result("成语必须是四个字的呢~")) # 发送回复，不能使用 yield - return - # 退出当前方法，不执行后续逻辑，但此会话并未中断，后续的用户输入仍然会进入当前会话 + return event.plain_result("成语必须是四个字的呢~") # 返回后不执行后续逻辑，但此会话并未中断，后续的用户输入仍然会进入当前会话 # ... - message_result = event.make_result() - message_result.chain = [Comp.Plain("先见之明")] # import astrbot.api.message_components as Comp - await event.send(message_result) # 发送回复，不能使用 yield - controller.keep(timeout=60, reset_timeout=True) # 重置超时时间为 60s，如果不重置，则会继续之前的超时时间计时。 # controller.stop() # 停止会话控制器，会立即结束。 # 如果记录了历史消息链，可以通过 controller.get_history_chains() 获取历史消息链 + return event.plain_result("先见之明") # 发送回复 try: await empty_mention_waiter(event) @@ -74,6 +58,10 @@ async def handle_empty_mention(self, event: AstrMessageEvent): logger.error("handle_empty_mention error: " + str(e)) ``` +> [!TIP] +> 会话控制器处理函数内**不能使用 `yield`**，但可以 `return` 一个 `MessageEventResult`。返回的结果会被框架自动送入消息装饰流程后再发送，与普通指令回复的行为一致。\ +> 如果直接调用 `event.send()`，消息会绕过装饰流程直接发送。 + 当激活会话控制器后，该发送人之后发送的消息会首先经过上面你定义的 `empty_mention_waiter` 函数处理，直到会话控制器被停止或者超时。 ## SessionController @@ -92,7 +80,7 @@ async def handle_empty_mention(self, event: AstrMessageEvent): ```py import astrbot.api.message_components as Comp -from astrbot.core.utils.session_waiter import ( +from astrbot.api.utils import ( session_waiter, SessionFilter, SessionController,