From 61efbd071e81a1cc59a147aed29569a9ac74e420 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 21:18:02 -0700 Subject: [PATCH] fix(aiohttp): bill the usage a stream sent before its connection dropped, and deflake the CyberArk override test (#47838) * fix(aiohttp): hand out bytes buffered before a dropped connection so a cut stream bills its partial usage test(integration): stop the CyberArk override test failing when a short config value lands inside ciphertext Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(aiohttp): cover bytes already read before a dropped connection Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llms/custom_httpx/aiohttp_transport.py | 15 +++++++ .../security/test_config_overrides.py | 5 ++- .../custom_httpx/test_aiohttp_transport.py | 43 +++++++++++++++++++ 3 files changed, 62 insertions(+), 1 deletion(-) diff --git a/litellm/llms/custom_httpx/aiohttp_transport.py b/litellm/llms/custom_httpx/aiohttp_transport.py index 9297e76624d..c32e75bafa8 100644 --- a/litellm/llms/custom_httpx/aiohttp_transport.py +++ b/litellm/llms/custom_httpx/aiohttp_transport.py @@ -110,6 +110,15 @@ def map_aiohttp_exceptions() -> Generator[None, None, None]: raise mapped_exc(message) from exc +def _unread_payload(content: object) -> bytes: + """Bytes the upstream sent before the connection dropped that aiohttp buffered but never handed out, since its + reads raise the stored payload error ahead of the buffer.""" + if not isinstance(content, aiohttp.StreamReader): + return b"" + buffered: Final = b"".join(content._buffer) # pyright: ignore[reportPrivateUsage] # no public read drains it + return buffered[content._buffer_offset :] # pyright: ignore[reportPrivateUsage] # see above + + def _release_transport_anchor(_transport: object) -> None: """Release a transport anchored to a response stream. See ``_anchor_transport_to``.""" @@ -144,6 +153,12 @@ class AiohttpResponseStream(httpx.AsyncByteStream): raise httpx.ReadError(str(e)) from e except aiohttp.http_exceptions.TransferEncodingError as e: raise httpx.ReadError(str(e)) from e + except aiohttp.ClientPayloadError: + unread: Final = _unread_payload(self._aiohttp_response.content) + if unread: + yield unread + with map_aiohttp_exceptions(): + raise except Exception: # For other exceptions, use the normal mapping with map_aiohttp_exceptions(): diff --git a/tests/integration/security/test_config_overrides.py b/tests/integration/security/test_config_overrides.py index 70f9226ea79..6b8e8791b93 100644 --- a/tests/integration/security/test_config_overrides.py +++ b/tests/integration/security/test_config_overrides.py @@ -715,8 +715,11 @@ def test_cyberark_incremental_update_encrypts_reloads_on_restart_and_delete_clea }, saved.text stored_full: Final = _stored_override(database_url, "cyberark") assert set(stored_full) == set(full_config), stored_full + assert all(stored_full[key] != value for key, value in full_config.items()), stored_full serialized_full: Final = json.dumps(stored_full) - assert all(value not in serialized_full for value in full_config.values()), serialized_full + assert all(value not in serialized_full for value in (cyberark_wire.url, account, username, api_key)), ( + serialized_full + ) model_one: Final = "integration-cyberark-" + uuid.uuid4().hex _new_model(first, model_one, provider_wire.url, secret_names[0]) diff --git a/tests/unit/llms/custom_httpx/test_aiohttp_transport.py b/tests/unit/llms/custom_httpx/test_aiohttp_transport.py index aed808223bd..4e4d5df9fb0 100644 --- a/tests/unit/llms/custom_httpx/test_aiohttp_transport.py +++ b/tests/unit/llms/custom_httpx/test_aiohttp_transport.py @@ -144,6 +144,49 @@ async def test_client_payload_error_mid_stream_raises_read_error(): assert mock_response.closed is True +@pytest.mark.asyncio +async def test_bytes_buffered_before_a_dropped_connection_are_yielded_before_the_read_error(): + loop: Final = asyncio.get_running_loop() + content: Final = aiohttp.StreamReader(aiohttp.base_protocol.BaseProtocol(loop), 2**16, loop=loop) + content.feed_data(b"event: message_start\n\n") + content.set_exception(aiohttp.ClientPayloadError("Response payload is not completed")) + mock_response: Final = MockAiohttpResponse() + mock_response.content = content + received: Final[list[bytes]] = [] + + async def _drain() -> None: + async for chunk in AiohttpResponseStream(mock_response): # pyright: ignore[reportArgumentType] # mock response + received.append(chunk) + + with pytest.raises(httpx.ReadError): + await _drain() + + assert received == [b"event: message_start\n\n"] + assert mock_response.closed is True + + +@pytest.mark.asyncio +async def test_bytes_already_read_before_a_dropped_connection_are_not_yielded_again(): + loop: Final = asyncio.get_running_loop() + content: Final = aiohttp.StreamReader(aiohttp.base_protocol.BaseProtocol(loop), 2**16, loop=loop) + content.feed_data(b"event: message_start\n\n") + content.feed_data(b"event: ping\n\n") + assert await content.read(7) == b"event: " + content.set_exception(aiohttp.ClientPayloadError("Response payload is not completed")) + mock_response: Final = MockAiohttpResponse() + mock_response.content = content + received: Final[list[bytes]] = [] + + async def _drain() -> None: + async for chunk in AiohttpResponseStream(mock_response): # pyright: ignore[reportArgumentType] # mock response + received.append(chunk) + + with pytest.raises(httpx.ReadError): + await _drain() + + assert received == [b"message_start\n\nevent: ping\n\n"] + + @pytest.mark.asyncio async def test_client_payload_error_before_first_chunk_raises_read_error(): """A connection reset before any body byte must surface, not yield an empty 200 body"""