From 9a5127b4f5f2d0d7e1d37361d4c7f7ad26bb74a5 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Thu, 25 Jun 2026 11:35:29 +0300 Subject: [PATCH] fix(vertex): stop O(n^2) re-parse of accumulated Gemini stream JSON handle_accumulated_json_chunk re-ran json.loads on the entire accumulated buffer after every fragment. For a streaming response fragmented across many chunks that is O(n^2) total work in a single GIL-holding C call, so a large enough Gemini response freezes the asyncio event loop for seconds, liveness probes time out, and the proxy pod gets killed and restarted. A complete Gemini stream value is a JSON object or array, so the buffer can only become parseable once its last non-whitespace byte can close one. Gate the json.loads attempt on that, which makes the common fragmented-response case parse roughly once instead of once per fragment. An 8MB payload drops from a 6.9s event-loop freeze to ~0.3s with identical parsed output. Resolves LIT-3503 Fixes #26181 --- .../vertex_and_google_ai_studio_gemini.py | 12 +++- ...test_vertex_and_google_ai_studio_gemini.py | 71 +++++++++++++++++++ 2 files changed, 80 insertions(+), 3 deletions(-) diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index 423a5dd5d179..32bdb17840d4 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -3601,16 +3601,22 @@ def handle_accumulated_json_chunk( chunk = litellm.CustomStreamWrapper._strip_sse_data_from_chunk(chunk) or "" message = chunk.replace("\n\n", "") - # Accumulate JSON data self.accumulated_json += message - # Try to parse the accumulated JSON + # json.loads on the whole buffer after every fragment is O(n^2) and + # holds the GIL, freezing the event loop for seconds on large responses + # (https://github.com/BerriAI/litellm/issues/26181). A complete Gemini + # chunk is a JSON object/array, so only attempt the parse once the + # buffer's last non-whitespace byte can close one. + stripped = self.accumulated_json.rstrip() + if not stripped or stripped[-1] not in "}]": + return None + try: _data = json.loads(self.accumulated_json) self.accumulated_json = "" # reset after successful parsing return self.chunk_parser(chunk=_data) except json.JSONDecodeError: - # If it's not valid JSON yet, continue to the next event return None def _common_chunk_parsing_logic( diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 1a2d0d86810e..20e7aec377b6 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -2799,6 +2799,77 @@ def test_partial_json_chunk_on_first_chunk(): ), "Should switch to accumulated_json mode" +def test_accumulated_json_does_not_reparse_every_fragment(): + """ + Regression test for https://github.com/BerriAI/litellm/issues/26181 + + handle_accumulated_json_chunk used to call json.loads on the entire + accumulated buffer after EVERY fragment. For a large response fragmented + across many chunks that is O(n^2) work in a single GIL-holding C call, + which freezes the asyncio event loop for seconds and kills liveness probes. + + The buffer only becomes a complete JSON object on the final fragment, so a + correct implementation parses it ~once, not once per fragment. We assert the + full chunk still parses correctly AND that json.loads is not called on every + fragment (which is what made it quadratic). + """ + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + ModelResponseIterator, + ) + + iterator = ModelResponseIterator( + streaming_response=MagicMock(), + sync_stream=True, + logging_obj=MagicMock(), + ) + iterator.chunk_type = "accumulated_json" + + text = "x" * 200_000 # no braces/brackets so only the final fragment closes + blob = json.dumps( + {"candidates": [{"content": {"role": "model", "parts": [{"text": text}]}}]} + ) + fragments = [blob[i : i + 4096] for i in range(0, len(blob), 4096)] + assert len(fragments) > 10, "need a multi-fragment payload to exercise the bug" + + parsed = None + with patch("json.loads", wraps=json.loads) as spy: + for fragment in fragments: + out = iterator.handle_accumulated_json_chunk(chunk=fragment) + if out is not None: + parsed = out + parse_calls = spy.call_count + + assert parsed is not None, "the reassembled chunk must still parse" + assert parsed.choices[0].delta.content == text, "content must be preserved intact" + + assert parse_calls <= 2, ( + f"json.loads was called {parse_calls} times for {len(fragments)} " + "fragments; the O(n^2) per-fragment re-parse has regressed" + ) + + +def test_accumulated_json_partial_fragment_returns_none_without_parsing(): + """A fragment that cannot complete the JSON must not trigger a json.loads + parse of the whole growing buffer (issue #26181).""" + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + ModelResponseIterator, + ) + + iterator = ModelResponseIterator( + streaming_response=MagicMock(), + sync_stream=True, + logging_obj=MagicMock(), + ) + iterator.chunk_type = "accumulated_json" + + with patch("json.loads", wraps=json.loads) as spy: + result = iterator.handle_accumulated_json_chunk( + chunk='{"candidates": [{"content": {"parts": [{"text": "partial' + ) + assert result is None + assert spy.call_count == 0, "incomplete buffer should not be parsed" + + def test_google_ai_studio_presence_penalty_supported(): """ Test that presence_penalty is supported for Google AI Studio Gemini.