diff --git a/litellm/proxy/spend_tracking/auto_router_session_queue.py b/litellm/proxy/spend_tracking/auto_router_session_queue.py index c4b12e04033..a0c94d2000c 100644 --- a/litellm/proxy/spend_tracking/auto_router_session_queue.py +++ b/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. diff --git a/litellm/proxy/spend_tracking/auto_router_sessions.py b/litellm/proxy/spend_tracking/auto_router_sessions.py index 120732dd5a9..c54cc068818 100644 --- a/litellm/proxy/spend_tracking/auto_router_sessions.py +++ b/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 diff --git a/tests/test_litellm/proxy/spend_tracking/test_auto_router_sessions.py b/tests/test_litellm/proxy/spend_tracking/test_auto_router_sessions.py index 0fc93702d0c..995e4a137a0 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_auto_router_sessions.py +++ b/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."""