perf(spend): append staged turns instead of rebuilding the buffer per turn

Staging rebuilt a session's turn tuple on every arrival, so a caller reusing one
session_id copied every turn already staged on each request, while holding the
queue's lock. The work grows with the square of the session's length within an
interval, and the staging cap bounds that at 50,000 turns rather than preventing
it.

ToolDiscoveryQueue already had the answer: an appendable list that the flush
drains wholesale. _Pending now holds one, appended to under the lock, which
makes staging cost the same for the thousandth turn of a session as for the
first. Nothing observes the buffer mid-append, since the only reader is the
flush that swapped it out.

A test pins the property by asserting the buffer object is the same one across
appends, so rebuilding it is caught rather than merely slower

litellm/proxy/spend_tracking/auto_router_session_queue.py

@@ -48,11 +48,19 @@ MAX_CACHED_SESSIONS = 10_000
 SESSIONS_PER_STATEMENT = 1_000
 
 
-@dataclass(frozen=True, slots=True)
+@dataclass(slots=True)
 class _Pending:
+    """One session's staged turns, appended to in place the way ``ToolDiscoveryQueue`` stages items.
+
+    Rebuilding this as a frozen value per turn would copy every turn already
+    staged, so a caller reusing one ``session_id`` would pay for the whole
+    interval on each request. It is mutated only under the queue's lock and
+    drained wholesale on flush, so nothing observes it mid-append.
+    """
+
     router_kind: str
     baseline_model: str | None
-    turns: tuple[TurnFacts, ...]
+    turns: list[TurnFacts]  # mutable-ok: appended per turn under the lock, drained wholesale on flush
 
 
 _Chunk = Mapping[SessionKey, _Pending]
@@ -145,23 +153,22 @@ class AutoRouterSessionQueue:
             if self._staged_turns >= self._max_staged_turns:
                 _warn_staging_full(self._max_staged_turns)
                 return
-            self._stage(key, router_kind, baseline_model, (turn,))
+            self._stage(key, router_kind, baseline_model, [turn])  # mutable-ok: becomes this session's staging buffer
 
-    def _stage(
-        self, key: SessionKey, router_kind: str, baseline_model: str | None, turns: tuple[TurnFacts, ...]
-    ) -> None:
+    def _stage(self, key: SessionKey, router_kind: str, baseline_model: str | None, turns: list[TurnFacts]) -> None:
         """Add turns to whatever this session already has staged; callers hold the lock.
 
         Arriving turns and a replayed batch stage identically, because the fold
         sorts by start time and so does not care which of the two came first.
         """
-        current = self._pending.get(key)
-        self._pending[key] = _Pending(
-            router_kind=router_kind,
-            baseline_model=baseline_model or (current.baseline_model if current is not None else None),
-            turns=(current.turns if current is not None else ()) + turns,
-        )
         self._staged_turns += len(turns)
+        current = self._pending.get(key)
+        if current is None:
+            self._pending[key] = _Pending(router_kind=router_kind, baseline_model=baseline_model, turns=turns)
+            return
+        current.router_kind = router_kind
+        current.baseline_model = baseline_model or current.baseline_model
+        current.turns.extend(turns)
 
     async def flush(self, prisma_client: "PrismaClient") -> int:
         """Fold and write every staged session. Returns rows written.

litellm/proxy/spend_tracking/auto_router_sessions.py

@@ -17,7 +17,7 @@ the model that served the turn, and a rollup row has already summed across
 models. This is the same reason ``savings.py`` prices in the spend writer.
 """
 
-from collections.abc import Mapping
+from collections.abc import Mapping, Sequence
 from dataclasses import dataclass, fields
 from datetime import datetime, timezone
 from functools import reduce
@@ -461,7 +461,7 @@ def merge_deltas(left: TurnDelta, right: TurnDelta) -> TurnDelta:
     )
 
 
-def fold_session(state: SessionState, turns: tuple[TurnFacts, ...], rates: RateLookup = cache_rates) -> TurnDelta:
+def fold_session(state: SessionState, turns: Sequence[TurnFacts], rates: RateLookup = cache_rates) -> TurnDelta:
     """Fold an interval of one session's turns onto its prior state, oldest first.
 
     Ordering happens here rather than at arrival because completion order is not

tests/test_litellm/proxy/spend_tracking/test_auto_router_sessions.py

@@ -579,6 +579,41 @@ class TestFlushDurability:
         assert await queue.flush(prisma) == 0
 
 
+@pytest.mark.asyncio
+class TestStagingCostsTheSamePerTurn:
+    """Staging must not re-copy the turns already staged for that session.
+
+    Rebuilding the buffer per turn is quadratic in the length of a session, so a
+    caller reusing one `session_id` pays for the whole interval on every request
+    while holding the queue's lock.
+    """
+
+    async def test_a_turn_is_appended_to_the_existing_buffer_not_a_fresh_copy(self):
+        from litellm.proxy.spend_tracking.auto_router_session_queue import AutoRouterSessionQueue
+
+        queue = AutoRouterSessionQueue()
+        key = ("s1", "g")
+        await queue.record_turn(key, "complexity", None, _turn_at(0))
+        buffer = queue._pending[key].turns
+
+        for at in (60, 120, 180):
+            await queue.record_turn(key, "complexity", None, _turn_at(at))
+
+        assert queue._pending[key].turns is buffer
+        assert [turn.started_at for turn in buffer] == [0, 60, 120, 180]
+
+    async def test_a_long_session_still_folds_every_turn_it_staged(self):
+        from litellm.proxy.spend_tracking.auto_router_session_queue import AutoRouterSessionQueue
+
+        table = _RecordingTable()
+        queue = AutoRouterSessionQueue()
+        for at in range(500):
+            await queue.record_turn(("s1", "g"), "complexity", None, _turn_at(at * 60))
+
+        assert await queue.flush(_RecordingPrisma(table)) == 1
+        assert table.upserts[0][1]["create"]["turns"] == 500
+
+
 @pytest.mark.asyncio
 class TestStagingIsBounded:
     """`session_id` is caller-controlled, so what is held between flushes needs a ceiling."""