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>

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():

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])

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"""