diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index facfa5eace5..e58c6d108ec 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -346,7 +346,6 @@ class DBSpendUpdateWriter: cache_creation_tokens=_extract_cache_creation_tokens(usage_obj), usage_object=usage_obj, ), - prisma_client=prisma_client, ) async def _enqueue_tool_usage_transaction( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 822ff8cb154..23bd8bb9096 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -16,7 +16,7 @@ import time import traceback import warnings from collections.abc import AsyncGenerator, Callable, Mapping -from datetime import datetime, timedelta, timezone +from datetime import date, datetime, timedelta, timezone from types import UnionType from typing import ( TYPE_CHECKING, @@ -16483,8 +16483,8 @@ async def get_adaptive_router_state( dependencies=[Depends(user_api_key_auth)], # mutable-ok: FastAPI takes a list of dependencies ) async def get_auto_router_benchmarks( - start_date: str, - end_date: str, + start_date: date, + end_date: date, user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], ): """Session-level benchmarks for every configured auto-router. @@ -16494,10 +16494,11 @@ async def get_auto_router_benchmarks( counterfactual baseline, and how the provider prompt cache behaved. Reads the per-session rollup, never the per-request spend logs. - ``start_date`` / ``end_date`` are ``YYYY-MM-DD``; the window is clamped to - the most recent ``BENCHMARKS_MAX_WINDOW_DAYS`` days and the response echoes - the window actually served. Sessions are counted whole when they were active - in the window. Returns 404 when no auto-router is configured. + ``start_date`` / ``end_date`` are ``YYYY-MM-DD``, rejected by the framework + when malformed; the window is clamped to the most recent + ``BENCHMARKS_MAX_WINDOW_DAYS`` days and the response echoes the window + actually served. Sessions are counted whole when they were active in the + window. Returns 404 when no auto-router is configured. """ from litellm.proxy.spend_tracking.auto_router_benchmarks import compute_benchmarks from litellm.proxy.spend_tracking.auto_router_sessions import auto_router_group_kinds @@ -16509,6 +16510,13 @@ async def get_auto_router_benchmarks( "error": CommonProxyErrors.not_allowed_access.value }, ) + if end_date < start_date: + raise HTTPException( + status_code=400, + detail={ # mutable-ok: HTTPException takes a dict detail + "error": "end_date must not be earlier than start_date." + }, + ) if llm_router is None: raise HTTPException( status_code=404, diff --git a/litellm/proxy/spend_tracking/auto_router_benchmarks.py b/litellm/proxy/spend_tracking/auto_router_benchmarks.py index 91557869297..a9f298d2ee4 100644 --- a/litellm/proxy/spend_tracking/auto_router_benchmarks.py +++ b/litellm/proxy/spend_tracking/auto_router_benchmarks.py @@ -15,7 +15,7 @@ One aggregate query covers every auto-router, rather than four per router. """ from collections.abc import Mapping -from datetime import datetime, timedelta, timezone +from datetime import date, timedelta from types import MappingProxyType from typing import TYPE_CHECKING, NamedTuple @@ -158,18 +158,16 @@ GROUP BY model_group """ -def clamp_window(start_date: str, end_date: str) -> _Window: - """Parse the range and enforce ``start >= end - BENCHMARKS_MAX_WINDOW_DAYS``. +def clamp_window(start_date: date, end_date: date) -> _Window: + """Enforce ``start >= end - BENCHMARKS_MAX_WINDOW_DAYS``. The returned start reflects the window actually served, which the response echoes so the dashboard can label what it is showing rather than what it - asked for. + asked for. The dates arrive already parsed, because a malformed one is the + route's contract to reject rather than this module's to discover. """ - start = datetime.fromisoformat(start_date).replace(tzinfo=timezone.utc) - end = datetime.fromisoformat(end_date).replace(tzinfo=timezone.utc) - floor = (end - timedelta(days=BENCHMARKS_MAX_WINDOW_DAYS)).replace(hour=0, minute=0, second=0, microsecond=0) - clamped = max(start, floor) - return _Window(start=clamped.date().isoformat(), end=end.date().isoformat()) + floor = end_date - timedelta(days=BENCHMARKS_MAX_WINDOW_DAYS) + return _Window(start=max(start_date, floor).isoformat(), end=end_date.isoformat()) def _rate_pct(part: int, whole: int) -> float: @@ -259,8 +257,8 @@ def summarize_group(row: _GroupRow, router_kind: str) -> AutoRouterGroupBenchmar async def compute_benchmarks( prisma_client: "PrismaClient", group_kinds: Mapping[str, str], - start_date: str, - end_date: str, + start_date: date, + end_date: date, ) -> AutoRouterBenchmarksResponse: """Aggregate the session rollup for every configured auto-router.""" window = clamp_window(start_date, end_date) diff --git a/litellm/proxy/spend_tracking/auto_router_session_queue.py b/litellm/proxy/spend_tracking/auto_router_session_queue.py index 52db6c78800..f51bbc472cb 100644 --- a/litellm/proxy/spend_tracking/auto_router_session_queue.py +++ b/litellm/proxy/spend_tracking/auto_router_session_queue.py @@ -1,18 +1,19 @@ -"""In-memory aggregation and durable flush for auto-router session rollups. +"""In-memory staging and durable flush for auto-router session rollups. -Follows ``AdaptiveRouterUpdateQueue``: the logging path only folds into memory, -and a background task drains the aggregate into Postgres with atomic increment -upserts, so two pods writing the same session compose rather than overwrite. +Follows ``AdaptiveRouterUpdateQueue``: the logging path only stages into memory +and never touches the database, and a background task drains what it staged into +Postgres with atomic increment upserts, so two pods writing the same session +compose rather than overwrite. -The one departure is that this queue also caches the session state the fold reads -from. A pod that has never seen a session loads its row once and classifies from -memory thereafter, which is what keeps a session correct across a restart or a -move between pods without paying a read per turn. +Classifying a turn depends on the session's prior state, so the fold happens in +the flusher rather than at arrival. That is the one place the stored state can be +read, folded onto and written back as a single unit of work; a read that fails +there has no half-finished write to corrupt, and the turns simply stage again. """ import asyncio from collections import OrderedDict -from dataclasses import dataclass, replace +from dataclasses import dataclass from datetime import datetime, timezone from functools import lru_cache from typing import TYPE_CHECKING @@ -21,11 +22,12 @@ from litellm._logging import verbose_proxy_logger from litellm.proxy.spend_tracking.auto_router_sessions import ( EMPTY_SESSION_STATE, SessionState, + StateLookup, + StateUnavailable, TurnDelta, TurnFacts, counters_of, - fold_turn, - merge_deltas, + fold_session, state_column, state_from_row, ) @@ -36,50 +38,39 @@ if TYPE_CHECKING: SessionKey = tuple[str, str] -DEFAULT_MAX_TRACKED_SESSIONS = 10_000 +DEFAULT_MAX_STAGED_TURNS = 50_000 +MAX_CACHED_SESSIONS = 10_000 @dataclass(frozen=True, slots=True) class _Pending: router_kind: str baseline_model: str | None - first_turn_at: float - last_turn_at: float - delta: TurnDelta + turns: tuple[TurnFacts, ...] @lru_cache(maxsize=1) -def _warn_pending_full(cap: int) -> None: +def _warn_staging_full(cap: int) -> None: verbose_proxy_logger.warning( - "auto_router_sessions: %d sessions staged for the next flush; new sessions are not being recorded " + "auto_router_sessions: %d turns staged for the next flush; further turns are not being recorded " "until it drains. Benchmarks will undercount until then", cap, ) -def _merge_pending(earlier: _Pending, later: _Pending) -> _Pending: - """Fold two staged batches for one session, oldest first.""" - return _Pending( - router_kind=later.router_kind, - baseline_model=later.baseline_model or earlier.baseline_model, - first_turn_at=min(earlier.first_turn_at, later.first_turn_at), - last_turn_at=max(earlier.last_turn_at, later.last_turn_at), - delta=merge_deltas(earlier.delta, later.delta), - ) - - def _epoch_to_datetime(value: float) -> datetime: return datetime.fromtimestamp(value, tz=timezone.utc) class AutoRouterSessionQueue: - """Folds auto-routed turns in memory and flushes them to the session rollup.""" + """Stages auto-routed turns in memory and folds them into the rollup on flush.""" - def __init__(self, max_tracked_sessions: int = DEFAULT_MAX_TRACKED_SESSIONS) -> None: + def __init__(self, max_staged_turns: int = DEFAULT_MAX_STAGED_TURNS) -> None: self._pending: dict[SessionKey, _Pending] = {} # mutable-ok: drained and replaced wholesale on flush self._state: OrderedDict[SessionKey, SessionState] = OrderedDict() # mutable-ok: bounded LRU cache + self._staged_turns = 0 self._lock = asyncio.Lock() - self._max_tracked_sessions = max_tracked_sessions + self._max_staged_turns = max_staged_turns async def record_turn( self, @@ -87,58 +78,95 @@ class AutoRouterSessionQueue: router_kind: str, baseline_model: str | None, turn: TurnFacts, - prisma_client: "PrismaClient", ) -> None: - """Classify one turn against its session and stage the increments. + """Stage one turn against its session. Does no I/O; the fold happens on flush. - The session id is caller-controlled, so the staged aggregate is capped: - past the cap a session that is already staged keeps accumulating, but a - new one is dropped rather than admitted. Without that bound a caller - sending a fresh id per request grows the aggregate without limit between - flushes, and benchmark rows are not worth an out-of-memory kill. + ``session_id`` is caller-controlled, so the staging is capped on turns + held rather than on sessions seen, which is the quantity that actually + bounds the memory. Past the cap a turn is dropped and logged, because + benchmark rows are not worth an out-of-memory kill. """ - loaded = await self._session_state(key, prisma_client) async with self._lock: - current = self._pending.get(key) - if current is None and len(self._pending) >= self._max_tracked_sessions: - _warn_pending_full(self._max_tracked_sessions) + if self._staged_turns >= self._max_staged_turns: + _warn_staging_full(self._max_staged_turns) return - cached = self._state.get(key) - delta = fold_turn(cached if cached is not None else loaded, turn) - self._remember(key, delta.state) - self._pending[key] = ( - _Pending( - router_kind=router_kind, - baseline_model=baseline_model, - first_turn_at=turn.started_at, - last_turn_at=turn.started_at, - delta=delta, - ) - if current is None - else replace( - current, - baseline_model=baseline_model or current.baseline_model, - first_turn_at=min(current.first_turn_at, turn.started_at), - last_turn_at=max(current.last_turn_at, turn.started_at), - delta=merge_deltas(current.delta, delta), - ) - ) + self._stage(key, router_kind, baseline_model, (turn,)) - async def _session_state(self, key: SessionKey, prisma_client: "PrismaClient") -> SessionState: - """The session's state, from memory when this pod has seen it before. + def _stage( + self, key: SessionKey, router_kind: str, baseline_model: str | None, turns: tuple[TurnFacts, ...] + ) -> None: + """Add turns to whatever this session already has staged; callers hold the lock. - Loading outside the lock keeps a slow read from stalling every other - session's fold; a concurrent loader for the same key at worst repeats the - read, since both resolve to the same stored row. + 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) + + async def flush(self, prisma_client: "PrismaClient") -> int: + """Fold and write every staged session. Returns rows written. + + A session that could not be read or written is staged again rather than + dropped. Reading, folding and writing are one unit here, so a failure + means nothing landed and replaying it cannot double-count; draining first + and swallowing the error would lose that interval's turns, tokens and + spend permanently on any transient database fault. """ async with self._lock: - cached = self._state.get(key) - if cached is not None: - self._state.move_to_end(key) - return cached - return await self._load_state(key, prisma_client) + batch = self._pending + self._pending = {} # mutable-ok: fresh staging for the next interval + self._staged_turns = 0 - async def _load_state(self, key: SessionKey, prisma_client: "PrismaClient") -> SessionState: + failed = { # mutable-ok: built once from the sessions that did not land + key: batch[key] for key in sorted(batch.keys()) if not await self._commit(key, batch[key], prisma_client) + } + if failed: + verbose_proxy_logger.warning( + "auto_router_sessions: %d of %d sessions could not be folded; re-staging them for the next flush", + len(failed), + len(batch), + ) + async with self._lock: + for key, pending in failed.items(): + self._stage(key, pending.router_kind, pending.baseline_model, pending.turns) + return len(batch) - len(failed) + + async def _commit(self, key: SessionKey, pending: _Pending, prisma_client: "PrismaClient") -> bool: + """Read the session's state, fold its staged turns onto it, write both back. + + The cache advances only once the write has landed, so a replay folds from + the state the failed attempt did rather than one that was never persisted. + An evicted session is not lost either; its next flush reloads the row it + was already written to, which costs one read and folds identically. + """ + state = await self._session_state(key, prisma_client) + if isinstance(state, StateUnavailable): + return False + delta = fold_session(state, pending.turns) + if not await self._write(key, pending, delta, prisma_client): + return False + self._state[key] = delta.state + self._state.move_to_end(key) + while len(self._state) > MAX_CACHED_SESSIONS: + self._state.popitem(last=False) + return True + + async def _session_state(self, key: SessionKey, prisma_client: "PrismaClient") -> StateLookup: + """The session's state: from memory if this pod has folded it before, else its row. + + Only the flusher touches this cache, so it needs no lock. A session this + pod has not seen costs one read, which is what keeps it correct across a + restart or a move between pods. + """ + cached = self._state.get(key) + if cached is not None: + self._state.move_to_end(key) + return cached session_id, model_group = key try: row = await AutoRouterSessionRepository(prisma_client).table.find_unique( @@ -149,63 +177,20 @@ class AutoRouterSessionQueue: } } ) - except Exception as e: # noqa: BLE001 # a read fault must not fail the spend write - verbose_proxy_logger.warning( - "auto_router_sessions: could not load session state for %s (%s); treating as a new session", key, e - ) - return EMPTY_SESSION_STATE + except Exception as e: # noqa: BLE001 # a read fault re-stages the turns rather than failing the flush + verbose_proxy_logger.warning("auto_router_sessions: could not load session state for %s (%s)", key, e) + return StateUnavailable() if row is None: return EMPTY_SESSION_STATE return state_from_row(row.last_model, row.last_turn_at, row.model_state) - def _remember(self, key: SessionKey, state: SessionState) -> None: - """Cache the session's next state, evicting the least recently used. - - An evicted session is not lost; its next turn reloads the row it was - already flushed to, which costs one read and classifies identically. - """ - self._state[key] = state - self._state.move_to_end(key) - while len(self._state) > self._max_tracked_sessions: - self._state.popitem(last=False) - - async def flush(self, prisma_client: "PrismaClient") -> int: - """Drain the aggregate into the session rollup. Returns rows written. - - A session whose write fails is staged again rather than dropped. Draining - first and swallowing the error would lose that interval's turns, tokens - and spend permanently on any transient database fault, and because the - upsert is atomic a failure means nothing landed, so replaying it cannot - double-count. - """ - async with self._lock: - batch = self._pending - self._pending = {} # mutable-ok: fresh aggregate for the next interval - - failed = { # mutable-ok: built once from the writes that did not land - key: batch[key] for key in sorted(batch.keys()) if not await self._write(key, batch[key], prisma_client) - } - if failed: - verbose_proxy_logger.warning( - "auto_router_sessions: %d of %d session writes failed; re-staging them for the next flush", - len(failed), - len(batch), - ) - async with self._lock: - for key, pending in failed.items(): - current = self._pending.get(key) - # The retried batch is older than anything staged since, so it - # merges underneath it and the state of the newer one wins. - self._pending[key] = pending if current is None else _merge_pending(pending, current) - return len(batch) - len(failed) - - async def _write(self, key: SessionKey, pending: _Pending, prisma_client: "PrismaClient") -> bool: + async def _write(self, key: SessionKey, pending: _Pending, delta: TurnDelta, prisma_client: "PrismaClient") -> bool: session_id, model_group = key - counters = counters_of(pending.delta) + counters = counters_of(delta) shared = { # mutable-ok: prisma's write API takes dict payloads - "last_turn_at": _epoch_to_datetime(pending.last_turn_at), - "last_model": pending.delta.state.last_model, - "model_state": state_column(pending.delta.state), + "last_turn_at": _epoch_to_datetime(max(turn.started_at for turn in pending.turns)), + "last_model": delta.state.last_model, + "model_state": state_column(delta.state), "baseline_model": pending.baseline_model, } try: @@ -223,7 +208,7 @@ class AutoRouterSessionQueue: "session_id": session_id, "model_group": model_group, "router_kind": pending.router_kind, - "first_turn_at": _epoch_to_datetime(pending.first_turn_at), + "first_turn_at": _epoch_to_datetime(min(turn.started_at for turn in pending.turns)), **shared, **counters, }, diff --git a/litellm/proxy/spend_tracking/auto_router_sessions.py b/litellm/proxy/spend_tracking/auto_router_sessions.py index 8f7f7e7f946..3ef357630f7 100644 --- a/litellm/proxy/spend_tracking/auto_router_sessions.py +++ b/litellm/proxy/spend_tracking/auto_router_sessions.py @@ -20,6 +20,7 @@ models. This is the same reason ``savings.py`` prices in the spend writer. from collections.abc import Mapping from dataclasses import dataclass, fields from datetime import datetime, timezone +from functools import reduce from types import MappingProxyType from typing import TYPE_CHECKING, Literal, Protocol @@ -80,6 +81,19 @@ EMPTY_SESSION_STATE = SessionState( ) +@dataclass(frozen=True, slots=True) +class StateUnavailable: + """The session's stored state could not be read, so there is nothing to fold onto. + + Distinct from a session with no history: an absent row means start from empty + and write the result, an unreadable one means try again rather than replace a + history that really happened with one derived from a single turn. + """ + + +StateLookup = SessionState | StateUnavailable + + @dataclass(frozen=True, slots=True) class TurnFacts: """One auto-routed request, as the spend writer sees it. @@ -448,3 +462,17 @@ def merge_deltas(left: TurnDelta, right: TurnDelta) -> TurnDelta: name: getattr(left, name) + getattr(right, name) for name in COUNTER_FIELDS }, # mutable-ok: a JSON object is a dict by definition ) + + +def fold_session(state: SessionState, turns: tuple[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 + start order: turns in flight together finish in whichever order the providers + answer, and the earlier one would otherwise read as a late arrival. + """ + return reduce( + lambda folded, turn: merge_deltas(folded, fold_turn(folded.state, turn, rates)), + sorted(turns, key=lambda turn: turn.started_at), + TurnDelta(state=state), + ) diff --git a/tests/test_litellm/proxy/spend_tracking/test_auto_router_benchmarks.py b/tests/test_litellm/proxy/spend_tracking/test_auto_router_benchmarks.py index fe2dae72c61..d402655bd34 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_auto_router_benchmarks.py +++ b/tests/test_litellm/proxy/spend_tracking/test_auto_router_benchmarks.py @@ -1,4 +1,4 @@ -from datetime import datetime, timedelta, timezone +from datetime import date, timedelta import pytest @@ -67,23 +67,25 @@ def group_row(**overrides): return row +START = date(2026, 7, 2) +END = date(2026, 8, 1) + + async def benchmarks_for(**overrides): prisma = _FakePrisma([group_row(**overrides)]) - return await compute_benchmarks(prisma, GROUP_KINDS, "2026-07-02", "2026-08-01") + return await compute_benchmarks(prisma, GROUP_KINDS, START, END) class TestWindowClamping: def test_a_wider_request_is_clamped_to_the_maximum_window(self): - window = clamp_window("2020-01-01", "2026-08-01") - expected = (datetime(2026, 8, 1, tzinfo=timezone.utc) - timedelta(days=BENCHMARKS_MAX_WINDOW_DAYS)).date() - assert window.start == expected.isoformat() + window = clamp_window(date(2020, 1, 1), END) + assert window.start == (END - timedelta(days=BENCHMARKS_MAX_WINDOW_DAYS)).isoformat() def test_a_narrower_request_is_served_as_asked(self): - assert clamp_window("2026-07-25", "2026-08-01").start == "2026-07-25" + assert clamp_window(date(2026, 7, 25), END).start == "2026-07-25" def test_the_response_echoes_the_window_actually_served(self): - window = clamp_window("2020-01-01", "2026-08-01") - assert window.end == "2026-08-01" + assert clamp_window(date(2020, 1, 1), END).end == "2026-08-01" @pytest.mark.asyncio @@ -189,7 +191,7 @@ class TestWarmingEstimate: class TestReadPathSource: async def test_the_dashboard_query_never_touches_the_spend_logs(self): prisma = _FakePrisma([group_row()]) - await compute_benchmarks(prisma, GROUP_KINDS, "2026-07-02", "2026-08-01") + await compute_benchmarks(prisma, GROUP_KINDS, START, END) sql = prisma.db.queries[0][0] assert "LiteLLM_SpendLogs" not in sql assert "LiteLLM_AutoRouterSession" in sql @@ -197,14 +199,14 @@ class TestReadPathSource: async def test_one_query_covers_every_configured_auto_router(self): prisma = _FakePrisma([group_row(), group_row(model_group="claude-router-2")]) result = await compute_benchmarks( - prisma, {"claude-auto": "semantic", "claude-router-2": "complexity"}, "2026-07-02", "2026-08-01" + prisma, {"claude-auto": "semantic", "claude-router-2": "complexity"}, START, END ) assert len(prisma.db.queries) == 1 assert {g.model_group for g in result.groups} == {"claude-auto", "claude-router-2"} async def test_each_group_is_labelled_with_its_router_kind(self): prisma = _FakePrisma([group_row(model_group="claude-router-2")]) - result = await compute_benchmarks(prisma, {"claude-router-2": "complexity"}, "2026-07-02", "2026-08-01") + result = await compute_benchmarks(prisma, {"claude-router-2": "complexity"}, START, END) assert result.groups[0].router_kind == "complexity" 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 c48f741846d..9ccc1806702 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 @@ -8,6 +8,7 @@ from litellm.proxy.spend_tracking.auto_router_sessions import ( PROMPT_CACHE_TTL_SECONDS, SessionState, TurnFacts, + fold_session, fold_turn, merge_deltas, state_from_row, @@ -201,6 +202,36 @@ class TestTtlSelection: assert inside_the_hour.stale_return_misses == 0 +class TestFoldSession: + """An interval's turns are folded in start order, not in arrival order.""" + + def test_turns_that_arrived_out_of_order_still_all_classify(self): + """Two turns in flight together finish in whichever order the providers answer.""" + arrival_order = (turn(MODEL_B, at=60, created=5000), turn(MODEL_A, at=0, created=5000)) + folded = fold_session(EMPTY_SESSION_STATE, arrival_order, rates=rates) + assert folded.turns == 2 + assert buckets(folded) == 2 + assert folded.state.last_model == MODEL_B + + def test_folding_an_interval_matches_folding_its_turns_one_at_a_time(self): + turns = ( + turn(MODEL_A, at=0, created=5000), + turn(MODEL_A, at=60, read=5000), + turn(MODEL_B, at=400, created=5000), + turn(MODEL_A, at=800, created=5000), + ) + deltas, state = fold_all(turns) + folded = fold_session(EMPTY_SESSION_STATE, turns, rates=rates) + assert folded.turns == sum(d.turns for d in deltas) + assert folded.return_turns == sum(d.return_turns for d in deltas) + assert folded.replay_spend == pytest.approx(sum(d.replay_spend for d in deltas)) + assert folded.state == state + + def test_an_empty_interval_folds_to_the_state_it_was_given(self): + _, state = fold_all((turn(MODEL_A, at=0, created=5000),)) + assert fold_session(state, (), rates=rates).state is state + + class TestOutOfOrderTurns: def test_a_late_turn_keeps_its_spend_but_not_its_classification(self): first = fold_turn(EMPTY_SESSION_STATE, turn(MODEL_A, at=100, created=5000, spend=0.02), rates=rates) @@ -336,18 +367,33 @@ def test_state_from_row_without_a_timestamp_starts_at_the_epoch(): ) -class _RecordingTable: - """A session-rollup table that can be told to fail.""" +class _StoredRow: + """A session rollup as prisma hands it back.""" - def __init__(self, fail: bool = False): - self.fail = fail + def __init__(self, last_model: str, last_turn_at: float, model_state: dict): + self.last_model = last_model + self.last_turn_at = datetime.fromtimestamp(last_turn_at, tz=timezone.utc) + self.model_state = model_state + + +class _RecordingTable: + """A session-rollup table that can be told to fail on either side.""" + + def __init__(self, fail_read: bool = False, fail_write: bool = False, row=None): + self.fail_read = fail_read + self.fail_write = fail_write + self.row = row + self.reads = 0 self.upserts: list = [] async def find_unique(self, where): - return None + self.reads += 1 + if self.fail_read: + raise RuntimeError("transient database fault") + return self.row async def upsert(self, where, data): - if self.fail: + if self.fail_write: raise RuntimeError("transient database fault") self.upserts.append((where, data)) @@ -361,74 +407,156 @@ def _turn_at(at: float, model: str = MODEL_A) -> TurnFacts: return turn(model, at=at, created=5000) +def _stored_session() -> _StoredRow: + """A session last served on MODEL_B that has already used MODEL_A.""" + return _StoredRow( + last_model=MODEL_B, + last_turn_at=60.0, + model_state={ + MODEL_A: {"last_used_at": 0.0, "provisioned_replay_spend": 0.0}, + MODEL_B: {"last_used_at": 60.0, "provisioned_replay_spend": 0.0}, + }, + ) + + +@pytest.mark.asyncio +class TestTheLoggingPathNeverTouchesTheDatabase: + """Staging a turn must not put a database round trip in front of spend tracking.""" + + async def test_turns_can_be_recorded_with_no_database_in_sight(self): + from litellm.proxy.spend_tracking.auto_router_session_queue import AutoRouterSessionQueue + + queue = AutoRouterSessionQueue() + for at in (0, 60, 120): + await queue.record_turn(("s1", "g"), "complexity", None, _turn_at(at)) + + table = _RecordingTable() + assert await queue.flush(_RecordingPrisma(table)) == 1 + assert table.upserts[0][1]["create"]["turns"] == 3 + + async def test_a_session_costs_one_read_per_flush_at_most_not_one_per_turn(self): + from litellm.proxy.spend_tracking.auto_router_session_queue import AutoRouterSessionQueue + + table = _RecordingTable() + prisma = _RecordingPrisma(table) + queue = AutoRouterSessionQueue() + for at in (0, 60, 120): + await queue.record_turn(("s1", "g"), "complexity", None, _turn_at(at)) + await queue.flush(prisma) + + assert table.reads == 1 + + await queue.record_turn(("s1", "g"), "complexity", None, _turn_at(180)) + await queue.flush(prisma) + assert table.reads == 1 + + @pytest.mark.asyncio class TestFlushDurability: - """A transient write fault must not silently delete an interval of traffic.""" + """A transient database fault must not delete an interval of traffic or its history.""" async def test_a_failed_write_is_restaged_rather_than_dropped(self): from litellm.proxy.spend_tracking.auto_router_session_queue import AutoRouterSessionQueue - table = _RecordingTable(fail=True) + table = _RecordingTable(fail_write=True) prisma = _RecordingPrisma(table) queue = AutoRouterSessionQueue() - await queue.record_turn(("s1", "g"), "complexity", None, _turn_at(0), prisma) + await queue.record_turn(("s1", "g"), "complexity", None, _turn_at(0)) assert await queue.flush(prisma) == 0 - table.fail = False + table.fail_write = False assert await queue.flush(prisma) == 1 assert table.upserts[0][1]["create"]["turns"] == 1 async def test_a_restaged_batch_merges_under_turns_staged_since(self): from litellm.proxy.spend_tracking.auto_router_session_queue import AutoRouterSessionQueue - table = _RecordingTable(fail=True) + table = _RecordingTable(fail_write=True) prisma = _RecordingPrisma(table) queue = AutoRouterSessionQueue() - await queue.record_turn(("s1", "g"), "complexity", None, _turn_at(0), prisma) + await queue.record_turn(("s1", "g"), "complexity", None, _turn_at(0)) await queue.flush(prisma) - table.fail = False - await queue.record_turn(("s1", "g"), "complexity", None, _turn_at(60), prisma) + table.fail_write = False + await queue.record_turn(("s1", "g"), "complexity", None, _turn_at(60)) assert await queue.flush(prisma) == 1 - # Both turns land, once each assert table.upserts[0][1]["create"]["turns"] == 2 + async def test_a_failed_state_read_writes_nothing_at_all(self): + """The corrupting move was folding onto an empty state and persisting it.""" + from litellm.proxy.spend_tracking.auto_router_session_queue import AutoRouterSessionQueue + + table = _RecordingTable(fail_read=True, row=_stored_session()) + prisma = _RecordingPrisma(table) + queue = AutoRouterSessionQueue() + await queue.record_turn(("s1", "g"), "complexity", None, _turn_at(120)) + + assert await queue.flush(prisma) == 0 + assert table.upserts == [] + + async def test_history_survives_a_failed_read_and_still_classifies_the_turn(self): + """Returning to a tier this session already used is a return, not a first visit.""" + from litellm.proxy.spend_tracking.auto_router_session_queue import AutoRouterSessionQueue + + table = _RecordingTable(fail_read=True, row=_stored_session()) + prisma = _RecordingPrisma(table) + queue = AutoRouterSessionQueue() + await queue.record_turn(("s1", "g"), "complexity", None, _turn_at(120)) + await queue.flush(prisma) + + table.fail_read = False + assert await queue.flush(prisma) == 1 + update = table.upserts[0][1]["update"] + assert update["return_turns"] == {"increment": 1} + assert update["first_visit_turns"] == {"increment": 0} + assert update["turns"] == {"increment": 1} + async def test_a_successful_flush_stages_nothing_back(self): from litellm.proxy.spend_tracking.auto_router_session_queue import AutoRouterSessionQueue prisma = _RecordingPrisma(_RecordingTable()) queue = AutoRouterSessionQueue() - await queue.record_turn(("s1", "g"), "complexity", None, _turn_at(0), prisma) + await queue.record_turn(("s1", "g"), "complexity", None, _turn_at(0)) assert await queue.flush(prisma) == 1 assert await queue.flush(prisma) == 0 @pytest.mark.asyncio -class TestPendingIsBounded: - """`session_id` is caller-controlled, so the staged aggregate needs a ceiling.""" +class TestStagingIsBounded: + """`session_id` is caller-controlled, so what is held between flushes needs a ceiling.""" - async def test_new_sessions_are_refused_once_the_aggregate_is_full(self): + async def test_turns_are_refused_once_the_staging_is_full(self): from litellm.proxy.spend_tracking.auto_router_session_queue import AutoRouterSessionQueue - prisma = _RecordingPrisma(_RecordingTable()) - queue = AutoRouterSessionQueue(max_tracked_sessions=2) + table = _RecordingTable() + queue = AutoRouterSessionQueue(max_staged_turns=2) for i in range(5): - await queue.record_turn((f"s{i}", "g"), "complexity", None, _turn_at(0), prisma) + await queue.record_turn((f"s{i}", "g"), "complexity", None, _turn_at(0)) - assert await queue.flush(prisma) == 2 + assert await queue.flush(_RecordingPrisma(table)) == 2 - async def test_a_session_already_staged_keeps_accumulating_at_the_cap(self): - """Refusing new keys must not stall the conversations already in flight.""" + async def test_the_ceiling_counts_turns_not_sessions(self): + """One caller replaying a long session must not evade the memory bound.""" + from litellm.proxy.spend_tracking.auto_router_session_queue import AutoRouterSessionQueue + + table = _RecordingTable() + queue = AutoRouterSessionQueue(max_staged_turns=2) + for at in (0, 60, 120, 180): + await queue.record_turn(("s1", "g"), "complexity", None, _turn_at(at)) + + await queue.flush(_RecordingPrisma(table)) + assert table.upserts[0][1]["create"]["turns"] == 2 + + async def test_staging_reopens_once_it_has_drained(self): from litellm.proxy.spend_tracking.auto_router_session_queue import AutoRouterSessionQueue table = _RecordingTable() prisma = _RecordingPrisma(table) - queue = AutoRouterSessionQueue(max_tracked_sessions=1) - await queue.record_turn(("s1", "g"), "complexity", None, _turn_at(0), prisma) - await queue.record_turn(("s2", "g"), "complexity", None, _turn_at(0), prisma) - await queue.record_turn(("s1", "g"), "complexity", None, _turn_at(60), prisma) - + queue = AutoRouterSessionQueue(max_staged_turns=1) + await queue.record_turn(("s1", "g"), "complexity", None, _turn_at(0)) await queue.flush(prisma) - assert len(table.upserts) == 1 - assert table.upserts[0][1]["create"]["turns"] == 2 + + await queue.record_turn(("s1", "g"), "complexity", None, _turn_at(60)) + assert await queue.flush(prisma) == 1 + assert table.upserts[1][1]["update"]["turns"] == {"increment": 1} diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 67e76585485..1353bc83cc4 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -10825,3 +10825,50 @@ def test_startup_is_silent_when_mock_testing_params_disabled(caplog): ProxyStartupEvent._warn_if_mock_testing_params_enabled(general_settings={}) assert MOCK_TESTING_CONFIG_KEY not in caplog.text + + +def _auto_router_benchmarks_client(monkeypatch): + import litellm.proxy.proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + monkeypatch.setattr(ps, "llm_router", None) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_id="u", user_role=LitellmUserRoles.PROXY_ADMIN + ) + return TestClient(app) + + +@pytest.mark.parametrize( + "params, expected_status", + [ + ({"start_date": "not-a-date", "end_date": "2026-08-01"}, 422), + ({"start_date": "2026-08-01", "end_date": "2026-13-45"}, 422), + ({"start_date": "2026-08-01", "end_date": "2026-07-02"}, 400), + ], +) +def test_auto_router_benchmarks_rejects_windows_it_cannot_serve( + monkeypatch, params, expected_status +): + """A malformed date used to surface as a 500 and an inverted range as an empty + dashboard, which reads to an operator as "no auto-router traffic" rather than as a + bad request. Both are the caller's error and must say so. 404 (no auto-router + configured) is checked last in the handler, so reaching it means the window passed.""" + client = _auto_router_benchmarks_client(monkeypatch) + try: + resp = client.get("/auto_router/benchmarks", params=params) + assert resp.status_code == expected_status, resp.text + finally: + app.dependency_overrides.clear() + + +def test_auto_router_benchmarks_accepts_a_well_formed_window(monkeypatch): + """The guard must not reject the windows the dashboard actually sends.""" + client = _auto_router_benchmarks_client(monkeypatch) + try: + resp = client.get( + "/auto_router/benchmarks", + params={"start_date": "2026-07-02", "end_date": "2026-08-01"}, + ) + assert resp.status_code == 404, resp.text + finally: + app.dependency_overrides.clear() diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 9df6e11098a..bdeb7bbf0c6 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -776,10 +776,11 @@ export interface paths { * counterfactual baseline, and how the provider prompt cache behaved. * * Reads the per-session rollup, never the per-request spend logs. - * ``start_date`` / ``end_date`` are ``YYYY-MM-DD``; the window is clamped to - * the most recent ``BENCHMARKS_MAX_WINDOW_DAYS`` days and the response echoes - * the window actually served. Sessions are counted whole when they were active - * in the window. Returns 404 when no auto-router is configured. + * ``start_date`` / ``end_date`` are ``YYYY-MM-DD``, rejected by the framework + * when malformed; the window is clamped to the most recent + * ``BENCHMARKS_MAX_WINDOW_DAYS`` days and the response echoes the window + * actually served. Sessions are counted whole when they were active in the + * window. Returns 404 when no auto-router is configured. */ get: operations["get_auto_router_benchmarks_auto_router_benchmarks_get"]; put?: never;