diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 5ba96f87c53..4edaad38983 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -706,7 +706,7 @@ class _BrotliDecoder: produced += len(piece) # rebind-ok: see above pending = b"" # rebind-ok: see above yield piece - if self._inflate.can_accept_more_data(): + if not piece or self._inflate.can_accept_more_data(): return return b"".join(drained()) diff --git a/tests/unit/llms/custom_httpx/test_http_handler.py b/tests/unit/llms/custom_httpx/test_http_handler.py index d7c105bec74..1c1f6dc4c90 100644 --- a/tests/unit/llms/custom_httpx/test_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_http_handler.py @@ -26,6 +26,7 @@ from litellm.llms.custom_httpx.http_handler import ( HTTPHandler, HTTPResponseLimitError, MaskedHTTPStatusError, + _BrotliDecoder, _get_httpx_client, get_ssl_configuration, ) @@ -1934,6 +1935,27 @@ async def test_bounded_get_decodes_a_compressed_body_under_the_cap_and_rejects_a await handler.close() +def test_brotli_decoder_stops_draining_when_process_returns_empty(): + pieces: Final = iter((b"notes", b" more", b"")) + calls: Final = [] + + class Inflate: + def process(self, data: bytes, output_buffer_limit: int) -> bytes: + calls.append(data) + if len(calls) > 8: + raise AssertionError("brotli drain loop did not stop after an empty process()") + return next(pieces, b"") + + def can_accept_more_data(self) -> bool: + return False + + def is_finished(self) -> bool: + return True + + assert _BrotliDecoder(Inflate()).decode(b"wire", max_output=1024) == b"notes more" + assert calls == [b"wire", b"", b""] + + @pytest.mark.asyncio async def test_bounded_get_refuses_an_encoding_it_cannot_decode_under_the_cap(respx_mock, monkeypatch): monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True")