From f690c5aa757897571f526adda400a30216a502b2 Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Wed, 19 Aug 2026 11:22:01 -0400 Subject: [PATCH] fix(streaming): make JSONFragmentAccumulator.pop_next_value O(1) per call pop_next_value previously rebuilt the full remaining string and sliced a new remainder on every call, so draining N concatenated JSON values already sitting in one buffer cost O(n^2) total. Replace the rebuild-and- slice with a materialize-once cursor: pending fragments are joined into the buffer only when new ones have arrived since the last pop, and consumed values are dropped by advancing an offset instead of copying the remaining string. --- .../json_fragment_accumulator.py | 59 ++++++++++++++----- .../test_json_fragment_accumulator.py | 24 ++++++++ 2 files changed, 68 insertions(+), 15 deletions(-) diff --git a/litellm/litellm_core_utils/json_fragment_accumulator.py b/litellm/litellm_core_utils/json_fragment_accumulator.py index d94b4510830..e829309b43b 100644 --- a/litellm/litellm_core_utils/json_fragment_accumulator.py +++ b/litellm/litellm_core_utils/json_fragment_accumulator.py @@ -1,25 +1,36 @@ import json from typing import Final, cast # noqa: TID251 # raw_decode returns tuple[Any, int]; no cast-free unpack +_JSON_WHITESPACE: Final = " \t\n\r" # matches json.decoder.WHITESPACE's character class + class JSONFragmentAccumulator: """ Buffers a JSON value that arrives piecemeal over a stream (SSE data split across TCP packets, one shard per network read, etc) without the O(n^2) - cost of repeated `buffer += fragment` string concatenation. + cost of repeated `buffer += fragment` string concatenation, and without + the O(n^2) cost of re-copying the unconsumed remainder on every peeled + value when one payload holds many concatenated JSON values. Fragments are appended to a list in O(1). The buffer is only rebuilt into a single string, and only decoded, when a caller asks for a value via `pop_next_value`, and `could_close_json` lets callers skip that rebuild - entirely for fragments that plainly cannot close a JSON value yet. + entirely for fragments that plainly cannot close a JSON value yet. Once + rebuilt, consumed values are dropped by advancing a cursor rather than + slicing a new string, so draining N concatenated values already sitting + in the buffer costs O(n) total, not O(n^2). """ def __init__(self) -> None: self._chunks: list[str] = [] # mutable-ok: O(1) append; string concat would copy the buffer each time + self._buffer: str = ( + "" # mutable-ok: lazily materialized join of _chunks, rebuilt only when _chunks is non-empty + ) + self._offset: int = 0 # mutable-ok: cursor past already-consumed values; avoids re-slicing on every pop self._could_close: bool = False # mutable-ok: cached heuristic; rescanning past fragments was itself O(n^2) def __bool__(self) -> bool: - return bool(self._chunks) + return bool(self._chunks) or self._offset < len(self._buffer) def append(self, fragment: str) -> None: self._chunks.append(fragment) # mutable-ok: see __init__ @@ -37,34 +48,52 @@ class JSONFragmentAccumulator: """ return self._could_close + def _materialize(self) -> None: + if not self._chunks: + return + unconsumed: Final = self._buffer[self._offset :] + self._buffer = unconsumed + "".join(self._chunks) # mutable-ok: merge pending fragments, once per append batch + self._offset = 0 # mutable-ok: see __init__ + self._chunks = [] # mutable-ok: see __init__ + def pop_next_value(self) -> tuple[bool, object]: """ Attempt to decode one complete JSON value from the front of the - buffer. On success, advances the buffer past that value (keeping any - unconsumed tail, e.g. a second concatenated value) and returns - (True, value). If the buffer is empty or holds no complete value - yet, it is left untouched and this returns (False, None). + buffer. On success, advances a cursor past that value (keeping any + unconsumed tail, e.g. a second concatenated value, in place rather + than copying it) and returns (True, value). If the buffer is empty + or holds no complete value yet, it is left untouched and this + returns (False, None). """ - full: Final = "".join(self._chunks).strip() - if not full: + self._materialize() + length: Final = len(self._buffer) + start = self._offset + while start < length and self._buffer[start] in _JSON_WHITESPACE: + start += 1 + if start >= length: + self._offset = start # mutable-ok: see __init__ return False, None decoder: Final = json.JSONDecoder() try: - raw_value: Final = decoder.raw_decode(full) + raw_value: Final = decoder.raw_decode(self._buffer, start) except json.JSONDecodeError: return False, None decoded, end_index = cast("tuple[object, int]", raw_value) # cast-ok: raw_decode returns tuple[Any, int] - remainder: Final = full[end_index:].strip() - self._chunks = [remainder] if remainder else [] # mutable-ok: replace buffer with the unconsumed tail - if not remainder: + self._offset = end_index # mutable-ok: see __init__ + if self._offset >= len(self._buffer): + self._buffer = "" # mutable-ok: see __init__ + self._offset = 0 # mutable-ok: see __init__ self._could_close = False # mutable-ok: buffer is empty, nothing can close return True, decoded def snapshot(self) -> str: - return "".join(self._chunks) + self._materialize() + return self._buffer[self._offset :] def set(self, value: str) -> None: """Replace the buffer's contents with a single fragment.""" - self._chunks = [value] if value else [] # mutable-ok: see __init__ + self._chunks = [] # mutable-ok: see __init__ + self._buffer = value # mutable-ok: see __init__ + self._offset = 0 # mutable-ok: see __init__ stripped: Final = value.rstrip() self._could_close = bool(stripped) and stripped[-1] in ("}", "]") # mutable-ok: see __init__ diff --git a/tests/test_litellm/litellm_core_utils/test_json_fragment_accumulator.py b/tests/test_litellm/litellm_core_utils/test_json_fragment_accumulator.py index d94f2a2fa48..af141115103 100644 --- a/tests/test_litellm/litellm_core_utils/test_json_fragment_accumulator.py +++ b/tests/test_litellm/litellm_core_utils/test_json_fragment_accumulator.py @@ -156,6 +156,30 @@ def test_accumulation_of_many_fragments_is_not_quadratic(): assert elapsed_ms < 50, f"1000-fragment append took {elapsed_ms:.1f} ms (expected < 50 ms); O(n^2) regression?" +def test_draining_many_concatenated_values_is_not_quadratic(): + """ + Regression guard: peeling N JSON values already sitting in one buffer, + one pop_next_value() call per value with no new fragments in between, + must be O(n) total. Re-copying the shrinking remainder on every pop + (slicing a new string instead of advancing a cursor) makes total drain + time scale with the square of the buffer size. + """ + accumulator = JSONFragmentAccumulator() + accumulator.append('{"a": 1}' * 80_000) + + start = time.perf_counter() + drained = 0 + while True: + found, _ = accumulator.pop_next_value() + if not found: + break + drained += 1 + elapsed_ms = (time.perf_counter() - start) * 1000 + + assert drained == 80_000 + assert elapsed_ms < 150, f"draining 80000 concatenated values took {elapsed_ms:.1f} ms (expected < 150 ms); O(n^2) regression?" + + def test_could_close_json_after_many_blank_fragments_is_not_quadratic(): """ Regression test: a hostile upstream can send malformed JSON that never