Fix: 修正 http/websocket client timeout

nonebot/drivers/__init__.py

@@ -9,6 +9,7 @@
     description: nonebot.drivers 模块
 """
 
+from nonebot.internal.driver import DEFAULT_TIMEOUT as DEFAULT_TIMEOUT
 from nonebot.internal.driver import URL as URL
 from nonebot.internal.driver import ASGIMixin as ASGIMixin
 from nonebot.internal.driver import Cookies as Cookies
@@ -31,6 +32,7 @@
 from nonebot.internal.driver import combine_driver as combine_driver
 
 __autodoc__ = {
+    "DEFAULT_TIMEOUT": True,
     "URL": True,
     "Cookies": True,
     "Request": True,

nonebot/drivers/aiohttp.py

@@ -38,13 +38,15 @@
 from nonebot.drivers.none import Driver as NoneDriver
 from nonebot.exception import WebSocketClosed
 from nonebot.internal.driver import (
+    DEFAULT_TIMEOUT,
     Cookies,
     CookieTypes,
     HeaderTypes,
     QueryTypes,
     Timeout,
     TimeoutTypes,
 )
+from nonebot.utils import UNSET, UnsetType, exclude_unset
 
 try:
     import aiohttp
@@ -63,7 +65,7 @@ def __init__(
         headers: HeaderTypes = None,
         cookies: CookieTypes = None,
         version: str | HTTPVersion = HTTPVersion.H11,
-        timeout: TimeoutTypes = None,
+        timeout: TimeoutTypes | UnsetType = UNSET,
         proxy: str | None = None,
     ):
         self._client: aiohttp.ClientSession | None = None
@@ -85,15 +87,32 @@ def __init__(
         else:
             raise RuntimeError(f"Unsupported HTTP version: {version}")
 
+        _timeout = None
         if isinstance(timeout, Timeout):
-            self._timeout = aiohttp.ClientTimeout(
-                total=timeout.total,
-                connect=timeout.connect,
-                sock_read=timeout.read,
+            timeout_kwargs: dict[str, float | None] = exclude_unset(
+                {
+                    "total": timeout.total,
+                    "connect": timeout.connect,
+                    "sock_read": timeout.read,
+                }
+            )
+            if timeout_kwargs:
+                _timeout = aiohttp.ClientTimeout(**timeout_kwargs)  # type: ignore
+        elif timeout is not UNSET:
+            _timeout = aiohttp.ClientTimeout(connect=timeout, sock_read=timeout)
+
+        if _timeout is None:
+            _timeout = aiohttp.ClientTimeout(
+                **exclude_unset(
+                    {
+                        "total": DEFAULT_TIMEOUT.total,
+                        "connect": DEFAULT_TIMEOUT.connect,
+                        "sock_read": DEFAULT_TIMEOUT.read,
+                    }
+                )
             )
-        else:
-            self._timeout = aiohttp.ClientTimeout(timeout)
 
+        self._timeout = _timeout
         self._proxy = proxy
 
     @property
@@ -102,6 +121,25 @@ def client(self) -> aiohttp.ClientSession:
             raise RuntimeError("Session is not initialized")
         return self._client
 
+    def _get_timeout(self, timeout: TimeoutTypes | UnsetType) -> aiohttp.ClientTimeout:
+        _timeout = None
+        if isinstance(timeout, Timeout):
+            timeout_kwargs: dict[str, float | None] = exclude_unset(
+                {
+                    "total": timeout.total,
+                    "connect": timeout.connect,
+                    "sock_read": timeout.read,
+                }
+            )
+            if timeout_kwargs:
+                _timeout = aiohttp.ClientTimeout(**timeout_kwargs)  # type: ignore
+        elif timeout is not UNSET:
+            _timeout = aiohttp.ClientTimeout(connect=timeout, sock_read=timeout)
+
+        if _timeout is None:
+            return self._timeout
+        return _timeout
+
     @override
     async def request(self, setup: Request) -> Response:
         if self._params:
@@ -121,15 +159,6 @@ async def request(self, setup: Request) -> Response:
             if cookie.value is not None
         )
 
-        if isinstance(setup.timeout, Timeout):
-            timeout = aiohttp.ClientTimeout(
-                total=setup.timeout.total,
-                connect=setup.timeout.connect,
-                sock_read=setup.timeout.read,
-            )
-        else:
-            timeout = aiohttp.ClientTimeout(setup.timeout)
-
         async with await self.client.request(
             setup.method,
             url,
@@ -138,7 +167,7 @@ async def request(self, setup: Request) -> Response:
             cookies=cookies,
             headers=setup.headers,
             proxy=setup.proxy or self._proxy,
-            timeout=timeout,
+            timeout=self._get_timeout(setup.timeout),
         ) as response:
             return Response(
                 response.status,
@@ -171,15 +200,6 @@ async def stream_request(
             if cookie.value is not None
         )
 
-        if isinstance(setup.timeout, Timeout):
-            timeout = aiohttp.ClientTimeout(
-                total=setup.timeout.total,
-                connect=setup.timeout.connect,
-                sock_read=setup.timeout.read,
-            )
-        else:
-            timeout = aiohttp.ClientTimeout(setup.timeout)
-
         async with self.client.request(
             setup.method,
             url,
@@ -188,7 +208,7 @@ async def stream_request(
             cookies=cookies,
             headers=setup.headers,
             proxy=setup.proxy or self._proxy,
-            timeout=timeout,
+            timeout=self._get_timeout(setup.timeout),
         ) as response:
             response_headers = response.headers.copy()
             # aiohttp does not guarantee fixed-size chunks; re-chunk to exact size
@@ -270,13 +290,39 @@ async def websocket(self, setup: Request) -> AsyncGenerator["WebSocket", None]:
         else:
             raise RuntimeError(f"Unsupported HTTP version: {setup.version}")
 
+        timeout = None
         if isinstance(setup.timeout, Timeout):
+            timeout_kwargs: dict[str, float | None] = exclude_unset(
+                {
+                    "ws_receive": setup.timeout.read,
+                    "ws_close": (
+                        setup.timeout.total
+                        if setup.timeout.close is UNSET
+                        else setup.timeout.close
+                    ),
+                }
+            )
+            if timeout_kwargs:
+                timeout = aiohttp.ClientWSTimeout(**timeout_kwargs)
+        elif setup.timeout is not UNSET:
             timeout = aiohttp.ClientWSTimeout(
-                ws_receive=setup.timeout.read,  # type: ignore
-                ws_close=setup.timeout.total,  # type: ignore
+                ws_receive=setup.timeout,  # type: ignore
+                ws_close=setup.timeout,  # type: ignore
+            )
+
+        if timeout is None:
+            timeout = aiohttp.ClientWSTimeout(
+                **exclude_unset(
+                    {
+                        "ws_receive": DEFAULT_TIMEOUT.read,
+                        "ws_close": (
+                            DEFAULT_TIMEOUT.total
+                            if DEFAULT_TIMEOUT.close is UNSET
+                            else DEFAULT_TIMEOUT.close
+                        ),
+                    }
+                )
             )
-        else:
-            timeout = aiohttp.ClientWSTimeout(ws_close=setup.timeout or 10.0)  # type: ignore
 
         async with aiohttp.ClientSession(version=version, trust_env=True) as session:
             async with session.ws_connect(
@@ -295,7 +341,7 @@ def get_session(
         headers: HeaderTypes = None,
         cookies: CookieTypes = None,
         version: str | HTTPVersion = HTTPVersion.H11,
-        timeout: TimeoutTypes = None,
+        timeout: TimeoutTypes | UnsetType = UNSET,
         proxy: str | None = None,
     ) -> Session:
         return Session(

nonebot/drivers/httpx.py

@@ -34,13 +34,15 @@
 )
 from nonebot.drivers.none import Driver as NoneDriver
 from nonebot.internal.driver import (
+    DEFAULT_TIMEOUT,
     Cookies,
     CookieTypes,
     HeaderTypes,
     QueryTypes,
     Timeout,
     TimeoutTypes,
 )
+from nonebot.utils import UNSET, UnsetType, exclude_unset
 
 try:
     import httpx
@@ -59,7 +61,7 @@ def __init__(
         headers: HeaderTypes = None,
         cookies: CookieTypes = None,
         version: str | HTTPVersion = HTTPVersion.H11,
-        timeout: TimeoutTypes = None,
+        timeout: TimeoutTypes | UnsetType = UNSET,
         proxy: str | None = None,
     ):
         self._client: httpx.AsyncClient | None = None
@@ -73,15 +75,34 @@ def __init__(
         self._cookies = Cookies(cookies)
         self._version = HTTPVersion(version)
 
+        _timeout = None
         if isinstance(timeout, Timeout):
-            self._timeout = httpx.Timeout(
-                timeout=timeout.total,
-                connect=timeout.connect,
-                read=timeout.read,
+            avg_timeout = timeout.total and timeout.total / 4
+            timeout_kwargs: dict[str, float | None] = exclude_unset(
+                {
+                    "timeout": avg_timeout,
+                    "connect": timeout.connect,
+                    "read": timeout.read,
+                }
+            )
+            if timeout_kwargs:
+                _timeout = httpx.Timeout(**timeout_kwargs)
+        elif timeout is not UNSET:
+            _timeout = httpx.Timeout(timeout)
+
+        if _timeout is None:
+            avg_timeout = DEFAULT_TIMEOUT.total and DEFAULT_TIMEOUT.total / 4
+            _timeout = httpx.Timeout(
+                **exclude_unset(
+                    {
+                        "timeout": avg_timeout,
+                        "connect": DEFAULT_TIMEOUT.connect,
+                        "read": DEFAULT_TIMEOUT.read,
+                    }
+                )
             )
-        else:
-            self._timeout = httpx.Timeout(timeout)
 
+        self._timeout = _timeout
         self._proxy = proxy
 
     @property
@@ -90,17 +111,28 @@ def client(self) -> httpx.AsyncClient:
             raise RuntimeError("Session is not initialized")
         return self._client
 
-    @override
-    async def request(self, setup: Request) -> Response:
-        if isinstance(setup.timeout, Timeout):
-            timeout = httpx.Timeout(
-                timeout=setup.timeout.total,
-                connect=setup.timeout.connect,
-                read=setup.timeout.read,
+    def _get_timeout(self, timeout: TimeoutTypes | UnsetType) -> httpx.Timeout:
+        _timeout = None
+        if isinstance(timeout, Timeout):
+            avg_timeout = timeout.total and timeout.total / 4
+            timeout_kwargs: dict[str, float | None] = exclude_unset(
+                {
+                    "timeout": avg_timeout,
+                    "connect": timeout.connect,
+                    "read": timeout.read,
+                }
             )
-        else:
-            timeout = httpx.Timeout(setup.timeout)
+            if timeout_kwargs:
+                _timeout = httpx.Timeout(**timeout_kwargs)
+        elif timeout is not UNSET:
+            _timeout = httpx.Timeout(timeout)
 
+        if _timeout is None:
+            return self._timeout
+        return _timeout
+
+    @override
+    async def request(self, setup: Request) -> Response:
         response = await self.client.request(
             setup.method,
             str(setup.url),
@@ -112,7 +144,7 @@ async def request(self, setup: Request) -> Response:
             params=setup.url.raw_query_string,
             headers=tuple(setup.headers.items()),
             cookies=setup.cookies.jar,
-            timeout=timeout,
+            timeout=self._get_timeout(setup.timeout),
         )
         return Response(
             response.status_code,
@@ -128,15 +160,6 @@ async def stream_request(
         *,
         chunk_size: int = 1024,
     ) -> AsyncGenerator[Response, None]:
-        if isinstance(setup.timeout, Timeout):
-            timeout = httpx.Timeout(
-                timeout=setup.timeout.total,
-                connect=setup.timeout.connect,
-                read=setup.timeout.read,
-            )
-        else:
-            timeout = httpx.Timeout(setup.timeout)
-
         async with self.client.stream(
             setup.method,
             str(setup.url),
@@ -148,7 +171,7 @@ async def stream_request(
             params=setup.url.raw_query_string,
             headers=tuple(setup.headers.items()),
             cookies=setup.cookies.jar,
-            timeout=timeout,
+            timeout=self._get_timeout(setup.timeout),
         ) as response:
             response_headers = response.headers.multi_items()
             async for chunk in response.aiter_bytes(chunk_size=chunk_size):