diff --git a/src/langbot/libs/wecom_ai_bot_api/ws_client.py b/src/langbot/libs/wecom_ai_bot_api/ws_client.py index 23cbb8d47..06f58552f 100644 --- a/src/langbot/libs/wecom_ai_bot_api/ws_client.py +++ b/src/langbot/libs/wecom_ai_bot_api/ws_client.py @@ -887,6 +887,12 @@ class WecomBotWsClient: try: await self._send_frame(frame) except Exception: + # A failed ping must also interrupt the listener. Otherwise + # the heartbeat task exits while the receive loop can remain + # blocked on an apparently open socket indefinitely. + await self.logger.warning('Heartbeat send failed; closing WebSocket to reconnect') + if self._ws and not self._ws.closed: + await self._ws.close() break except asyncio.CancelledError: pass diff --git a/tests/unit_tests/platform/test_wecombot_ws_heartbeat_failure.py b/tests/unit_tests/platform/test_wecombot_ws_heartbeat_failure.py new file mode 100644 index 000000000..874d75e63 --- /dev/null +++ b/tests/unit_tests/platform/test_wecombot_ws_heartbeat_failure.py @@ -0,0 +1,56 @@ +"""A failed outbound heartbeat must unblock the WebSocket receive loop.""" + +import asyncio +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest + +from langbot.libs.wecom_ai_bot_api.ws_client import WecomBotWsClient + + +@pytest.mark.asyncio +async def test_heartbeat_send_failure_closes_socket_for_reconnect(): + logger = SimpleNamespace(warning=AsyncMock()) + client = WecomBotWsClient( + 'bot-id', 'secret', logger, heartbeat_interval=0, + ) + socket = SimpleNamespace( + closed=False, + send_str=AsyncMock(side_effect=ConnectionResetError('socket write failed')), + close=AsyncMock(), + ) + client._ws = socket + client._running = True + + await asyncio.wait_for(client._heartbeat_loop(), timeout=1.0) + + socket.send_str.assert_awaited_once() + socket.close.assert_awaited_once() + logger.warning.assert_awaited_once() + assert client._missed_pong_count == 1 + + +@pytest.mark.asyncio +async def test_successful_heartbeat_does_not_close_socket(): + logger = SimpleNamespace(warning=AsyncMock()) + client = WecomBotWsClient( + 'bot-id', 'secret', logger, heartbeat_interval=0, + ) + + async def sent_ping(_frame): + client._running = False + + socket = SimpleNamespace( + closed=False, + send_str=AsyncMock(side_effect=sent_ping), + close=AsyncMock(), + ) + client._ws = socket + client._running = True + + await asyncio.wait_for(client._heartbeat_loop(), timeout=1.0) + + socket.send_str.assert_awaited_once() + socket.close.assert_not_awaited() + logger.warning.assert_not_awaited()