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.

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__

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