Feature: 为内置驱动器请求方法添加 auto_redirects 参数, 以便控制重定向行为

nonebot/drivers/aiohttp.py

@@ -169,6 +169,7 @@
             headers=setup.headers,
             proxy=setup.proxy or self._proxy,
             timeout=self._get_timeout(setup.timeout),
+            **exclude_unset({"allow_redirects": setup.auto_redirects}),
         ) as response:
             return Response(
                 response.status,
@@ -210,6 +211,7 @@
             headers=setup.headers,
             proxy=setup.proxy or self._proxy,
             timeout=self._get_timeout(setup.timeout),
+            **exclude_unset({"allow_redirects": setup.auto_redirects}),
         ) as response:
             response_headers = response.headers.copy()
             # aiohttp does not guarantee fixed-size chunks; re-chunk to exact size
@@ -335,6 +337,12 @@
                 "the configured ping timeout will be ignored."
             )
 
+        if setup.auto_redirects is not UNSET:
+            logger.warning(
+                "aiohttp driver websocket method does not expose redirects parameter; "
+                "the configured auto_redirects will be ignored."
+            )
+
         async with aiohttp.ClientSession(version=version, trust_env=True) as session:
             async with session.ws_connect(
                 setup.url,
@@ -345,7 +353,7 @@
                 autoping=heartbeat is not None,
                 heartbeat=heartbeat,
             ) as ws:
                 yield WebSocket(request=setup, session=session, websocket=ws)
 
     @override
     def get_session(

nonebot/drivers/httpx.py

@@ -145,6 +145,7 @@ async def request(self, setup: Request) -> Response:
             headers=tuple(setup.headers.items()),
             cookies=setup.cookies.jar,
             timeout=self._get_timeout(setup.timeout),
+            **exclude_unset({"follow_redirects": setup.auto_redirects}),
         )
         return Response(
             response.status_code,
@@ -172,6 +173,7 @@ async def stream_request(
             headers=tuple(setup.headers.items()),
             cookies=setup.cookies.jar,
             timeout=self._get_timeout(setup.timeout),
+            **exclude_unset({"follow_redirects": setup.auto_redirects}),
         ) as response:
             response_headers = response.headers.multi_items()
             async for chunk in response.aiter_bytes(chunk_size=chunk_size):

nonebot/drivers/httpx2.py

@@ -145,6 +145,7 @@ async def request(self, setup: Request) -> Response:
             headers=tuple(setup.headers.items()),
             cookies=setup.cookies.jar,
             timeout=self._get_timeout(setup.timeout),
+            **exclude_unset({"follow_redirects": setup.auto_redirects}),
         )
         return Response(
             response.status_code,
@@ -172,6 +173,7 @@ async def stream_request(
             headers=tuple(setup.headers.items()),
             cookies=setup.cookies.jar,
             timeout=self._get_timeout(setup.timeout),
+            **exclude_unset({"follow_redirects": setup.auto_redirects}),
         ) as response:
             response_headers = response.headers.multi_items()
             async for chunk in response.aiter_bytes(chunk_size=chunk_size):

nonebot/drivers/websockets.py

@@ -111,6 +111,12 @@ async def websocket(self, setup: Request) -> AsyncGenerator["WebSocket", None]:
             }
         )
 
+        if setup.auto_redirects is not UNSET:
+            logger.warning(
+                "websockets driver websocket method does not expose "
+                "redirects parameter; the configured auto_redirects will be ignored."
+            )
+
         connection = connect(
             str(setup.url),
             additional_headers={**setup.headers, **setup.cookies.as_header(setup)},

nonebot/internal/driver/model.py

@@ -79,6 +79,7 @@ def __init__(
         timeout: TimeoutTypes | UnsetType = UNSET,
         proxy: str | None = None,
         ping_interval: PingIntervalTypes | UnsetType = UNSET,
+        auto_redirects: bool | UnsetType = UNSET,
     ):
         # method
         self.method: str = (
@@ -94,6 +95,8 @@ def __init__(
         self.proxy: str | None = proxy
         # ping interval
         self.ping_interval: PingIntervalTypes | UnsetType = ping_interval
+        # auto redirects
+        self.auto_redirects: bool | UnsetType = auto_redirects
 
         # url
         if isinstance(url, tuple):

tests/fake_server.py

@@ -79,6 +79,10 @@ def http_echo(request: Request) -> Response:
     )
 
 
+def http_redirect(request: Request) -> Response:
+    return Response("redirecting", status=302, headers={"Location": "/"})
+
+
 def websocket_echo(request: Request) -> Response:
     stream = request.environ["werkzeug.socket"]
 
@@ -137,5 +141,7 @@ def websocket_echo(request: Request) -> Response:
 def request_handler(request: Request) -> Response:
     if request.headers.get("Connection") == "Upgrade":
         return websocket_echo(request)
+    elif request.path == "/redirect":
+        return http_redirect(request)
     else:
         return http_echo(request)

tests/test_driver.py

@@ -372,8 +372,8 @@ async def test_http_client(driver: Driver, server_url: URL):
     )
     chunks = []
     async for resp in driver.stream_request(request, chunk_size=4):
-        assert response.status_code == 200
-        assert response.content
+        assert resp.status_code == 200
+        assert resp.content
         chunks.append(resp.content)
     assert all(len(chunk) == 4 for chunk in chunks[:-1])
     data = json.loads(b"".join(chunks))
@@ -388,6 +388,62 @@ async def test_http_client(driver: Driver, server_url: URL):
     await anyio.sleep(1)
 
 
+@pytest.mark.anyio
+@pytest.mark.parametrize(
+    "driver",
+    [
+        pytest.param("nonebot.drivers.httpx:Driver", id="httpx"),
+        pytest.param("nonebot.drivers.httpx2:Driver", id="httpx2"),
+        pytest.param("nonebot.drivers.aiohttp:Driver", id="aiohttp"),
+    ],
+    indirect=True,
+)
+async def test_http_client_redirects(driver: Driver, server_url: URL):
+    assert isinstance(driver, HTTPClientMixin)
+
+    redirect_url = server_url.join(URL("/redirect"))
+
+    # follow redirects by default
+    request = Request("GET", redirect_url)
+    assert request.auto_redirects is UNSET
+    response = await driver.request(request)
+    assert response.status_code == 200
+    assert response.content
+    data = json.loads(response.content)
+    assert data["method"] == "GET"
+    assert URL(data["url"]).path == "/"
+
+    # disable auto redirects
+    request = Request("GET", redirect_url, auto_redirects=False)
+    assert request.auto_redirects is False
+    response = await driver.request(request)
+    assert response.status_code == 302
+    assert response.headers["Location"] == "/"
+    assert response.content == b"redirecting"
+
+    # stream request following redirects
+    request = Request("GET", redirect_url)
+    chunks = []
+    async for resp in driver.stream_request(request, chunk_size=4):
+        assert resp.status_code == 200
+        assert resp.content
+        chunks.append(resp.content)
+    data = json.loads(b"".join(chunks))
+    assert data["method"] == "GET"
+    assert URL(data["url"]).path == "/"
+
+    # stream request without auto redirects
+    request = Request("GET", redirect_url, auto_redirects=False)
+    chunks = []
+    async for resp in driver.stream_request(request, chunk_size=4):
+        assert resp.status_code == 302
+        assert resp.content
+        chunks.append(resp.content)
+    assert b"".join(chunks) == b"redirecting"
+
+    await anyio.sleep(1)
+
+
 @pytest.mark.anyio
 @pytest.mark.parametrize(
     "driver",
@@ -600,6 +656,64 @@ async def test_http_client_session(driver: Driver, server_url: URL):
     await anyio.sleep(1)
 
 
+@pytest.mark.anyio
+@pytest.mark.parametrize(
+    "driver",
+    [
+        pytest.param("nonebot.drivers.httpx:Driver", id="httpx"),
+        pytest.param("nonebot.drivers.httpx2:Driver", id="httpx2"),
+        pytest.param("nonebot.drivers.aiohttp:Driver", id="aiohttp"),
+    ],
+    indirect=True,
+)
+async def test_http_client_session_redirects(driver: Driver, server_url: URL):
+    assert isinstance(driver, HTTPClientMixin)
+
+    redirect_url = server_url.join(URL("/redirect"))
+
+    session = driver.get_session()
+    async with session:
+        # follow redirects by default
+        request = Request("GET", redirect_url)
+        assert request.auto_redirects is UNSET
+        response = await session.request(request)
+        assert response.status_code == 200
+        assert response.content
+        data = json.loads(response.content)
+        assert data["method"] == "GET"
+        assert URL(data["url"]).path == "/"
+
+        # disable auto redirects
+        request = Request("GET", redirect_url, auto_redirects=False)
+        assert request.auto_redirects is False
+        response = await session.request(request)
+        assert response.status_code == 302
+        assert response.headers["Location"] == "/"
+        assert response.content == b"redirecting"
+
+        # stream request following redirects
+        request = Request("GET", redirect_url)
+        chunks = []
+        async for resp in session.stream_request(request, chunk_size=4):
+            assert resp.status_code == 200
+            assert resp.content
+            chunks.append(resp.content)
+        data = json.loads(b"".join(chunks))
+        assert data["method"] == "GET"
+        assert URL(data["url"]).path == "/"
+
+        # stream request without auto redirects
+        request = Request("GET", redirect_url, auto_redirects=False)
+        chunks = []
+        async for resp in session.stream_request(request, chunk_size=4):
+            assert resp.status_code == 302
+            assert resp.content
+            chunks.append(resp.content)
+        assert b"".join(chunks) == b"redirecting"
+
+    await anyio.sleep(1)
+
+
 @pytest.mark.anyio
 async def test_aiohttp_stream_request_skip_empty_chunk() -> None:
     class _FakeContent: