🐛 Fix: 修复 aiohttp 流式响应分块大小不固定问题 (#3919)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
keep_running
2026-03-27 11:28:05 +08:00
committed by GitHub
parent 4f1c590ee9
commit eea5257394
2 changed files with 57 additions and 1 deletions

View File

@@ -191,11 +191,26 @@ class Session(HTTPClientSession):
timeout=timeout, timeout=timeout,
) as response: ) as response:
response_headers = response.headers.copy() response_headers = response.headers.copy()
# aiohttp does not guarantee fixed-size chunks; re-chunk to exact size
buffer = bytearray()
async for chunk in response.content.iter_chunked(chunk_size): async for chunk in response.content.iter_chunked(chunk_size):
if not chunk:
continue
buffer.extend(chunk)
while len(buffer) >= chunk_size:
out = bytes(buffer[:chunk_size])
del buffer[:chunk_size]
yield Response(
response.status,
headers=response_headers,
content=out,
request=setup,
)
if buffer:
yield Response( yield Response(
response.status, response.status,
headers=response_headers, headers=response_headers,
content=chunk, content=bytes(buffer),
request=setup, request=setup,
) )

View File

@@ -22,6 +22,7 @@ from nonebot.drivers import (
WebSocketClientMixin, WebSocketClientMixin,
WebSocketServerSetup, WebSocketServerSetup,
) )
from nonebot.drivers.aiohttp import Session as AiohttpSession
from nonebot.drivers.aiohttp import WebSocket as AiohttpWebSocket from nonebot.drivers.aiohttp import WebSocket as AiohttpWebSocket
from nonebot.exception import WebSocketClosed from nonebot.exception import WebSocketClosed
from nonebot.params import Depends from nonebot.params import Depends
@@ -596,6 +597,46 @@ async def test_http_client_session(driver: Driver, server_url: URL):
await anyio.sleep(1) await anyio.sleep(1)
@pytest.mark.anyio
async def test_aiohttp_stream_request_skip_empty_chunk() -> None:
class _FakeContent:
async def iter_chunked(self, _: int):
for chunk in (b"ab", b"", b"cd", b"e"):
yield chunk
class _FakeResponse:
def __init__(self) -> None:
self.status = 200
self.headers = {"x-test": "1"}
self.content = _FakeContent()
class _FakeRequestContext:
async def __aenter__(self) -> _FakeResponse:
return _FakeResponse()
async def __aexit__(self, *args: object) -> bool:
return False
class _FakeClient:
def request(self, *args: object, **kwargs: object) -> _FakeRequestContext:
return _FakeRequestContext()
session = AiohttpSession()
session._client = _FakeClient() # type: ignore[assignment]
chunks = []
async for resp in session.stream_request(
Request("GET", "https://example.com"), chunk_size=2
):
assert resp.status_code == 200
assert resp.content
chunks.append(resp.content)
assert chunks == [b"ab", b"cd", b"e"]
assert b"".join(chunks) == b"abcde"
assert all(len(chunk) == 2 for chunk in chunks[:-1])
@pytest.mark.anyio @pytest.mark.anyio
@pytest.mark.parametrize( @pytest.mark.parametrize(
"driver", "driver",