diff --git a/nonebot/drivers/__init__.py b/nonebot/drivers/__init__.py index 4add8a19..f5e87fca 100644 --- a/nonebot/drivers/__init__.py +++ b/nonebot/drivers/__init__.py @@ -430,5 +430,7 @@ class WebSocketSetup: """URL""" headers: Dict[str, str] = field(default_factory=dict) """HTTP headers""" + reconnect: bool = True + """WebSocket 是否重连""" reconnect_interval: float = 3.0 """WebSocket 重连间隔""" diff --git a/nonebot/drivers/aiohttp.py b/nonebot/drivers/aiohttp.py index 40879d62..fbbf673d 100644 --- a/nonebot/drivers/aiohttp.py +++ b/nonebot/drivers/aiohttp.py @@ -407,6 +407,12 @@ class Driver(ForwardDriver): if bot: self._bot_disconnect(bot) bot = None + + if not setup_.reconnect: + logger.info( + f"WebSocket reconnect disabled for bot {setup_.self_id}" + ) + break await asyncio.sleep(setup_.reconnect_interval) except asyncio.CancelledError: diff --git a/nonebot/drivers/fastapi.py b/nonebot/drivers/fastapi.py index 2afb082a..ec579f51 100644 --- a/nonebot/drivers/fastapi.py +++ b/nonebot/drivers/fastapi.py @@ -555,7 +555,7 @@ class FullDriver(ForwardDriver, Driver): asyncio.create_task(bot.handle_message(msg)) except ConnectionClosed: logger.opt(colors=True).error( - "WebSocket connection closed by peer. " + "WebSocket connection closed. " "Try to reconnect..." ) break @@ -568,6 +568,12 @@ class FullDriver(ForwardDriver, Driver): if bot: self._bot_disconnect(bot) bot = None + + if not setup_.reconnect: + logger.info( + f"WebSocket reconnect disabled for bot {setup_.self_id}" + ) + break await asyncio.sleep(setup_.reconnect_interval) except asyncio.CancelledError: