diff --git a/astrbot/core/platform/sources/wecom_ai_bot/wecomai_long_connection.py b/astrbot/core/platform/sources/wecom_ai_bot/wecomai_long_connection.py index 1017dd2300..d6d14c8b5b 100644 --- a/astrbot/core/platform/sources/wecom_ai_bot/wecomai_long_connection.py +++ b/astrbot/core/platform/sources/wecom_ai_bot/wecomai_long_connection.py @@ -34,6 +34,7 @@ def __init__( self._send_lock = asyncio.Lock() self._command_lock = asyncio.Lock() self._response_waiters: dict[str, asyncio.Future[dict[str, Any]]] = {} + self._message_handler_tasks: set[asyncio.Task[None]] = set() @staticmethod def gen_req_id() -> str: @@ -138,7 +139,11 @@ async def _handle_text_message(self, text: str) -> None: cmd = payload.get("cmd") if cmd in {"aibot_msg_callback", "aibot_event_callback"}: - await self.message_handler(payload) + # Keep the receive loop available for command acknowledgements sent by + # the callback handler, such as the configured initial response. + task = asyncio.create_task(self.message_handler(payload)) + self._message_handler_tasks.add(task) + task.add_done_callback(self._on_message_handler_done) return if payload.get("errcode") not in (None, 0): @@ -148,6 +153,17 @@ async def _handle_text_message(self, text: str) -> None: payload.get("errmsg"), ) + def _on_message_handler_done(self, task: asyncio.Task[None]) -> None: + """Release a completed callback task and report its exception.""" + self._message_handler_tasks.discard(task) + if task.cancelled(): + return + if exception := task.exception(): + logger.error( + "[WecomAI][LongConn] 处理回调消息失败", + exc_info=exception, + ) + async def send_command( self, cmd: str, diff --git a/tests/test_wecomai_long_connection.py b/tests/test_wecomai_long_connection.py new file mode 100644 index 0000000000..984a30287b --- /dev/null +++ b/tests/test_wecomai_long_connection.py @@ -0,0 +1,56 @@ +"""Tests for the WeCom AI Bot long-connection client.""" + +import asyncio +import json +from unittest.mock import AsyncMock + +import pytest + +from astrbot.core.platform.sources.wecom_ai_bot.wecomai_long_connection import ( + WecomAIBotLongConnectionClient, +) + + +@pytest.mark.asyncio +async def test_callback_handler_can_receive_command_ack() -> None: + """A callback response must not block the socket's ACK receive path.""" + handler_started = asyncio.Event() + handler_finished = asyncio.Event() + + async def message_handler(_: dict) -> None: + handler_started.set() + sent = await client.send_command( + "aibot_respond_msg", + "response-request", + {"msgtype": "stream"}, + ) + assert sent is True + handler_finished.set() + + client = WecomAIBotLongConnectionClient( + bot_id="bot-id", + secret="secret", + ws_url="wss://example.com", + heartbeat_interval=30, + message_handler=message_handler, + ) + client._ws = AsyncMock(closed=False) + + callback = json.dumps( + { + "cmd": "aibot_msg_callback", + "headers": {"req_id": "callback-request"}, + "body": {}, + } + ) + await asyncio.wait_for(client._handle_text_message(callback), timeout=0.1) + await asyncio.wait_for(handler_started.wait(), timeout=0.1) + + acknowledgement = json.dumps( + {"headers": {"req_id": "response-request"}, "errcode": 0} + ) + await client._handle_text_message(acknowledgement) + + await asyncio.wait_for(handler_finished.wait(), timeout=0.1) + await asyncio.sleep(0) + assert not client._message_handler_tasks