From 79f1531aa07410bb30fc69c1811a65d9cef565d0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 21 Oct 2026 00:07:58 -0700 Subject: [PATCH] fix(aiohttp): yield bytes already received before raising on a dropped stream aiohttp read methods raise the stored payload error before returning buffered bytes, so a stream whose connection dropped after the final SSE events lost them. The transport now drains the reader's buffer once and yields it before raising httpx.ReadError, and falls back to the old behavior if aiohttp's internals change --- .../llms/custom_httpx/aiohttp_transport.py | 41 ++++-- .../custom_httpx/test_aiohttp_transport.py | 127 ++++++++++++++++-- 2 files changed, 148 insertions(+), 20 deletions(-) diff --git a/litellm/llms/custom_httpx/aiohttp_transport.py b/litellm/llms/custom_httpx/aiohttp_transport.py index b8cb183c02f..57860d4332e 100644 --- a/litellm/llms/custom_httpx/aiohttp_transport.py +++ b/litellm/llms/custom_httpx/aiohttp_transport.py @@ -109,6 +109,32 @@ def map_aiohttp_exceptions() -> Generator[None, None, None]: raise mapped_exc(message) from exc +@contextlib.contextmanager +def _map_stream_read_errors() -> Generator[None, None, None]: + try: + yield + except RuntimeError as e: + if "Connection closed" not in str(e): + raise + raise httpx.ReadError(str(e)) from e + except aiohttp.http_exceptions.TransferEncodingError as e: + raise httpx.ReadError(str(e)) from e + except Exception: + with map_aiohttp_exceptions(): + raise + + +def _received_unread_bytes(content: aiohttp.StreamReader) -> bytes: + read_buffer: Final[Callable[[int], object] | None] = getattr(content, "_read_nowait", None) + if read_buffer is None: + return b"" + try: + unread: Final = read_buffer(-1) + except (AttributeError, TypeError, TimeoutError): + return b"" + return unread if isinstance(unread, bytes) else b"" + + class AiohttpResponseStream(httpx.AsyncByteStream): CHUNK_SIZE = 1024 * 16 @@ -116,18 +142,15 @@ class AiohttpResponseStream(httpx.AsyncByteStream): self._aiohttp_response = aiohttp_response async def __aiter__(self) -> typing.AsyncIterator[bytes]: + content: Final = self._aiohttp_response.content try: - async for chunk in self._aiohttp_response.content.iter_chunked(self.CHUNK_SIZE): + async for chunk in content.iter_chunked(self.CHUNK_SIZE): yield chunk - except RuntimeError as e: - if "Connection closed" not in str(e): - raise - raise httpx.ReadError(str(e)) from e - except aiohttp.http_exceptions.TransferEncodingError as e: - raise httpx.ReadError(str(e)) from e except Exception: - # For other exceptions, use the normal mapping - with map_aiohttp_exceptions(): + unread: Final = _received_unread_bytes(content) + if unread: + yield unread + with _map_stream_read_errors(): raise finally: # Release the aiohttp connection when iteration ends for any diff --git a/tests/unit/llms/custom_httpx/test_aiohttp_transport.py b/tests/unit/llms/custom_httpx/test_aiohttp_transport.py index 0765371dc4a..5696afcf89e 100644 --- a/tests/unit/llms/custom_httpx/test_aiohttp_transport.py +++ b/tests/unit/llms/custom_httpx/test_aiohttp_transport.py @@ -1276,6 +1276,20 @@ async def test_refused_connection_maps_to_httpx_connect_error( assert (cause.host, cause.port) == ("127.0.0.1", port) +async def _read_request_headers(reader: asyncio.StreamReader) -> bytes: + headers: Final = await reader.readuntil(b"\r\n\r\n") + content_length: Final = next( + ( + int(header.split(b":", 1)[1]) + for header in headers.split(b"\r\n") + if header.lower().startswith(b"content-length:") + ), + 0, + ) + await reader.readexactly(content_length) + return headers + + @pytest.mark.asyncio async def test_reset_after_request_is_not_replayed(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("NO_PROXY", "127.0.0.1") @@ -1283,17 +1297,7 @@ async def test_reset_after_request_is_not_replayed(monkeypatch: pytest.MonkeyPat received_requests: Final[list[bytes]] = [] # mutable-ok: local server captures requests async def reset_after_request(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: - headers: Final = await reader.readuntil(b"\r\n\r\n") - content_length: Final = next( - ( - int(header.split(b":", 1)[1]) - for header in headers.split(b"\r\n") - if header.lower().startswith(b"content-length:") - ), - 0, - ) - await reader.readexactly(content_length) - received_requests.append(headers) + received_requests.append(await _read_request_headers(reader)) transport_socket: Final = writer.get_extra_info("socket") assert transport_socket is not None transport_socket.setsockopt(socket.SOL_SOCKET, socket.SO_LINGER, struct.pack("ii", 1, 0)) @@ -1310,3 +1314,104 @@ async def test_reset_after_request_is_not_replayed(monkeypatch: pytest.MonkeyPat assert len(received_requests) == 1 finally: await handler.close() + + +_SSE_EVENTS_BEFORE_DROP: Final = ( + b'data: {"id":"chatcmpl-drop","object":"chat.completion.chunk","created":1,"model":"gpt-5.6",' + b'"choices":[{"index":0,"delta":{"role":"assistant","content":"1 2 "},"finish_reason":null}]}\n\n', + b'data: {"id":"chatcmpl-drop","object":"chat.completion.chunk","created":1,"model":"gpt-5.6",' + b'"choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}\n\n', +) + + +async def _stream_events_then_drop_before_done(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: + await _read_request_headers(reader) + chunked_events: Final = b"".join(b"%x\r\n%s\r\n" % (len(event), event) for event in _SSE_EVENTS_BEFORE_DROP) + writer.write( + b"HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nTransfer-Encoding: chunked\r\n\r\n" + chunked_events + ) + await writer.drain() + writer.close() + await writer.wait_closed() + + +async def _until_connection_drop_is_received(content: aiohttp.StreamReader) -> None: + while content.exception() is None: + await asyncio.sleep(0) + + +async def _read_until_error(stream: httpx.AsyncByteStream) -> tuple[bytes, Exception | None]: + chunks: Final[list[bytes]] = [] # mutable-ok: keeps the chunks read before the stream raises + try: + async for chunk in stream: + chunks.append(chunk) + except Exception as error: + return b"".join(chunks), error + return b"".join(chunks), None + + +@pytest.mark.asyncio +async def test_bytes_received_before_a_connection_drop_are_yielded_before_read_error() -> None: + server: Final = await asyncio.start_server(_stream_events_then_drop_before_done, "127.0.0.1", 0) + port: Final = server.sockets[0].getsockname()[1] + async with server, aiohttp.ClientSession() as session: + aiohttp_response: Final = await session.post(f"http://127.0.0.1:{port}/v1/chat/completions", data=b"{}") + await _until_connection_drop_is_received(aiohttp_response.content) + received, error = await _read_until_error(AiohttpResponseStream(aiohttp_response)) + + assert received == b"".join(_SSE_EVENTS_BEFORE_DROP) + assert isinstance(error, httpx.ReadError) + + +@pytest.mark.asyncio +async def test_finished_stream_dropped_before_done_ends_cleanly_through_acompletion( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("NO_PROXY", "127.0.0.1") + monkeypatch.setenv("no_proxy", "127.0.0.1") + upstream_responses: Final[list[aiohttp.ClientResponse]] = [] # mutable-ok: trace hook captures the response + + async def capture_response( + _session: ClientSession, _context: object, params: aiohttp.TraceRequestEndParams + ) -> None: + upstream_responses.append(params.response) + + trace_config: Final = aiohttp.TraceConfig() + trace_config.on_request_end.append(capture_response) + server: Final = await asyncio.start_server(_stream_events_then_drop_before_done, "127.0.0.1", 0) + port: Final = server.sockets[0].getsockname()[1] + async with server, aiohttp.ClientSession(trace_configs=[trace_config]) as session: + stream: Final = await litellm.acompletion( + model="openai/gpt-5.6", + api_base=f"http://127.0.0.1:{port}/v1", + api_key="sk-local-test", + messages=[{"role": "user", "content": "count to two"}], + stream=True, + shared_session=session, + ) + await _until_connection_drop_is_received(upstream_responses[0].content) + chunks: Final = [chunk async for chunk in stream] + + assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == "1 2 " + assert chunks[-1].choices[0].finish_reason == "stop" + + +class _ContentWithChangedBufferInternals(MockContent): + def _read_nowait(self, n: int) -> bytes: + raise TypeError(f"_read_nowait({n}) no longer drains the buffer") + + +@pytest.mark.asyncio +async def test_payload_error_still_maps_to_read_error_when_buffer_drain_breaks() -> None: + mock_response: Final = MockAiohttpResponse() + mock_response.content = _ContentWithChangedBufferInternals( + chunks=[b"chunk1", b"chunk2"], + exception_to_raise=aiohttp.ClientPayloadError("Response payload is not completed"), + exception_at_chunk=1, + ) + + received, error = await _read_until_error(AiohttpResponseStream(mock_response)) # pyright: ignore[reportArgumentType] # duck-typed aiohttp response double + + assert received == b"chunk1" + assert isinstance(error, httpx.ReadError) + assert mock_response.closed is True