✨ Feature: WS 支持 ping interval/timeout 配置 (#3964)

Co-authored-by: Ju4tCode <42488585+yanyongyu@users.noreply.github.com>

nonebot/drivers/aiohttp.py

@@ -46,6 +46,7 @@ from nonebot.internal.driver import (
     Timeout,
     TimeoutTypes,
 )
+from nonebot.log import logger
 from nonebot.utils import UNSET, UnsetType, exclude_unset
 
 try:
@@ -324,6 +325,16 @@ class Mixin(HTTPClientMixin, WebSocketClientMixin):
                 )
             )
 
+        heartbeat = None
+        if setup.ping_interval is not UNSET:
+            heartbeat = setup.ping_interval
+
+        if isinstance(setup.timeout, Timeout) and setup.timeout.ping is not UNSET:
+            logger.warning(
+                "aiohttp driver does not expose a separate ping timeout; "
+                "the configured ping timeout will be ignored."
+            )
+
         async with aiohttp.ClientSession(version=version, trust_env=True) as session:
             async with session.ws_connect(
                 setup.url,
@@ -331,6 +342,8 @@ class Mixin(HTTPClientMixin, WebSocketClientMixin):
                 timeout=timeout,
                 headers=setup.headers,
                 proxy=setup.proxy,
+                autoping=heartbeat is not None,
+                heartbeat=heartbeat,
             ) as ws:
                 yield WebSocket(request=setup, session=session, websocket=ws)
 

nonebot/drivers/websockets.py

@@ -36,7 +36,7 @@ from nonebot.drivers import WebSocket as BaseWebSocket
 from nonebot.drivers.none import Driver as NoneDriver
 from nonebot.exception import WebSocketClosed
 from nonebot.log import LoguruHandler
-from nonebot.utils import UNSET, exclude_unset
+from nonebot.utils import UNSET, UnsetType, exclude_unset
 
 try:
     from websockets import ClientConnection, ConnectionClosed, connect
@@ -77,14 +77,17 @@ class Mixin(WebSocketClientMixin):
     @override
     @asynccontextmanager
     async def websocket(self, setup: Request) -> AsyncGenerator["WebSocket", None]:
-        timeout_kwargs: dict[str, float | None] = {}
+        timeout_kwargs: dict[str, float | None | UnsetType] = {}
         if isinstance(setup.timeout, Timeout):
             open_timeout = (
                 setup.timeout.connect or setup.timeout.read or setup.timeout.total
             )
-            timeout_kwargs = exclude_unset(
-                {"open_timeout": open_timeout, "close_timeout": setup.timeout.close}
-            )
+            timeout_kwargs = {
+                "open_timeout": open_timeout,
+                "close_timeout": setup.timeout.close,
+                "ping_timeout": setup.timeout.ping,
+            }
+
         elif setup.timeout is not UNSET:
             timeout_kwargs = {
                 "open_timeout": setup.timeout,
@@ -95,18 +98,24 @@ class Mixin(WebSocketClientMixin):
             open_timeout = (
                 DEFAULT_TIMEOUT.connect or DEFAULT_TIMEOUT.read or DEFAULT_TIMEOUT.total
             )
-            timeout_kwargs = exclude_unset(
-                {
-                    "open_timeout": open_timeout,
-                    "close_timeout": DEFAULT_TIMEOUT.close,
-                }
-            )
+            timeout_kwargs = {
+                "open_timeout": open_timeout,
+                "close_timeout": DEFAULT_TIMEOUT.close,
+                "ping_timeout": DEFAULT_TIMEOUT.ping,
+            }
+
+        kwargs = exclude_unset(
+            {
+                **timeout_kwargs,
+                "ping_interval": setup.ping_interval,
+            }
+        )
 
         connection = connect(
             str(setup.url),
             additional_headers={**setup.headers, **setup.cookies.as_header(setup)},
             proxy=setup.proxy if setup.proxy is not None else True,
-            **timeout_kwargs,  # type: ignore
+            **kwargs,  # type: ignore
         )
         async with connection as ws:
             yield WebSocket(request=setup, websocket=ws)

nonebot/internal/driver/model.py

@@ -20,9 +20,10 @@ class Timeout:
     connect: float | None | UnsetType = UNSET
     read: float | None | UnsetType = UNSET
     close: float | None | UnsetType = UNSET
+    ping: float | None | UnsetType = UNSET
 
 
-DEFAULT_TIMEOUT = Timeout(total=None, connect=5.0, read=30.0, close=10.0)
+DEFAULT_TIMEOUT = Timeout(total=None, connect=5.0, read=30.0, close=10.0, ping=20.0)
 
 
 RawURL: TypeAlias = tuple[bytes, bytes, int | None, bytes]
@@ -52,6 +53,7 @@ FileTypes: TypeAlias = (
 )
 FilesTypes: TypeAlias = dict[str, FileTypes] | list[tuple[str, FileTypes]] | None
 TimeoutTypes: TypeAlias = float | Timeout | None
+PingIntervalTypes: TypeAlias = float | None
 
 
 class HTTPVersion(Enum):
@@ -76,6 +78,7 @@ class Request:
         version: str | HTTPVersion = HTTPVersion.H11,
         timeout: TimeoutTypes | UnsetType = UNSET,
         proxy: str | None = None,
+        ping_interval: PingIntervalTypes | UnsetType = UNSET,
     ):
         # method
         self.method: str = (
@@ -89,6 +92,8 @@ class Request:
         self.timeout: TimeoutTypes | UnsetType = timeout
         # proxy
         self.proxy: str | None = proxy
+        # ping interval
+        self.ping_interval: PingIntervalTypes | UnsetType = ping_interval
 
         # url
         if isinstance(url, tuple):

tests/test_driver.py

@@ -878,6 +878,92 @@ async def test_websocket_client_timeout(driver: Driver, server_url: URL):
     await anyio.sleep(1)
 
 
+@pytest.mark.anyio
+@pytest.mark.parametrize(
+    "driver",
+    [
+        pytest.param("nonebot.drivers.websockets:Driver", id="websockets"),
+        pytest.param("nonebot.drivers.aiohttp:Driver", id="aiohttp"),
+    ],
+    indirect=True,
+)
+async def test_websocket_client_ping_timeout(driver: Driver, server_url: URL):
+    """WebSocket connections work with different ping_timeout settings."""
+    assert isinstance(driver, WebSocketClientMixin)
+
+    ws_url = server_url.with_scheme("ws")
+
+    # ping timeout not set (UNSET), falls back to DEFAULT_TIMEOUT.ping
+    request = Request("GET", ws_url, timeout=Timeout())
+    async with driver.websocket(request) as ws:
+        await ws.send("quit")
+        with pytest.raises(WebSocketClosed):
+            await ws.receive()
+
+    await anyio.sleep(1)
+
+    # ping timeout explicitly set to None (disable ping timeout)
+    request = Request("GET", ws_url, timeout=Timeout(ping=None))
+    async with driver.websocket(request) as ws:
+        await ws.send("quit")
+        with pytest.raises(WebSocketClosed):
+            await ws.receive()
+
+    await anyio.sleep(1)
+
+    # ping timeout set to a float value
+    request = Request("GET", ws_url, timeout=Timeout(ping=20.0))
+    async with driver.websocket(request) as ws:
+        await ws.send("quit")
+        with pytest.raises(WebSocketClosed):
+            await ws.receive()
+
+    await anyio.sleep(1)
+
+
+@pytest.mark.anyio
+@pytest.mark.parametrize(
+    "driver",
+    [
+        pytest.param("nonebot.drivers.websockets:Driver", id="websockets"),
+        pytest.param("nonebot.drivers.aiohttp:Driver", id="aiohttp"),
+    ],
+    indirect=True,
+)
+async def test_websocket_client_ping_interval(driver: Driver, server_url: URL):
+    """WebSocket connections work with different ping_interval settings."""
+    assert isinstance(driver, WebSocketClientMixin)
+
+    ws_url = server_url.with_scheme("ws")
+
+    # ping_interval not set (UNSET), default behavior
+    request = Request("GET", ws_url)
+    async with driver.websocket(request) as ws:
+        await ws.send("quit")
+        with pytest.raises(WebSocketClosed):
+            await ws.receive()
+
+    await anyio.sleep(1)
+
+    # ping_interval explicitly set to None (disable ping)
+    request = Request("GET", ws_url, ping_interval=None)
+    async with driver.websocket(request) as ws:
+        await ws.send("quit")
+        with pytest.raises(WebSocketClosed):
+            await ws.receive()
+
+    await anyio.sleep(1)
+
+    # ping_interval set to a float value
+    request = Request("GET", ws_url, ping_interval=20.0)
+    async with driver.websocket(request) as ws:
+        await ws.send("quit")
+        with pytest.raises(WebSocketClosed):
+            await ws.receive()
+
+    await anyio.sleep(1)
+
+
 @pytest.mark.parametrize(
     ("driver", "driver_type"),
     [