diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 3c65fc0bf06..3b921824195 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -2836,6 +2836,7 @@ async def ui_view_spend_logs( sql_conditions: Final[list[str]] = [] sql_params: Final[list[object]] = [] p = 1 # parameter index counter + window_literal_conditions: Final[list[str]] = [] # mutable-ok: filled once below from the time-window helper # Date range. Wrap the param side with `AT TIME ZONE 'UTC'` so comparison # against the plain `timestamp` column does not depend on the DB session @@ -2847,6 +2848,10 @@ async def ui_view_spend_logs( sql_conditions.append(f"\"startTime\" <= (${p}::timestamptz AT TIME ZONE 'UTC')") sql_params.append(end_date_obj) p += 1 + window_literal_conditions.extend( + f'"startTime" {op} {_utc_timestamp_literal(bound)}' + for op, bound in ((">=", start_date_obj), ("<=", end_date_obj)) + ) if search is not None and start_date_obj is not None and end_date_obj is not None: search_condition: Final = _build_spend_log_search_condition( @@ -2945,10 +2950,26 @@ async def ui_view_spend_logs( sql_params.append(json.dumps(used_client_oauth_token)) p += 1 - if status_filter is not None and group_by_session is True and not is_search_lookup: - session_filter_conditions: Final = " AND ".join(sql_conditions) or "TRUE" - sql_conditions.append( - f"""({_SESSION_GROUP_KEY_SQL}) IN ( + session_page_route: Final = ( + group_by_session is True + and not is_v2 + and not is_request_id_lookup + and not is_search_lookup + and sort_by == "startTime" + ) + # The session page reads a failure filter's sessions once, as a MATERIALIZED CTE: inlined in the read, + # each parallel worker re-sorts the whole window for it (3x the blocks and temp on 2M rows). A success + # filter keeps most sessions, and there the inlined form is faster. + materialize_status_sessions: Final = session_page_route and status_filter == "failure" + # The window goes in as literals (parse_date's UTC datetimes, checked by + # _utc_timestamp_literal). Prisma reuses its prepared statements, and from the 6th + # run Postgres plans a $N window as 0.5% of the table, then probes every + # session by api_key: ~2 min instead of ~2.5 s over 700k rows. + session_filter_conditions: Final = ( + " AND ".join([*window_literal_conditions, *sql_conditions[len(window_literal_conditions) :]]) or "TRUE" + ) + status_sessions_sql: Final = ( + f""" SELECT session_key, api_key FROM ( SELECT DISTINCT ON ({_SESSION_GROUP_KEY_SQL}) {_SESSION_KEY_EXPR} AS session_key, api_key, status @@ -2957,18 +2978,20 @@ async def ui_view_spend_logs( ORDER BY {_SESSION_GROUP_KEY_SQL}, {_SESSION_REPRESENTATIVE_ORDER_SQL} ) AS session_outcomes WHERE COALESCE(status, 'success') = ${p} - )""" + """ + if status_filter is not None and group_by_session is True and not is_search_lookup + else None + ) + if status_sessions_sql is not None: + sql_conditions.append( + f"({_SESSION_GROUP_KEY_SQL}) IN (SELECT session_key, api_key FROM status_sessions)" + if materialize_status_sessions + else f"({_SESSION_GROUP_KEY_SQL}) IN ({status_sessions_sql})" ) sql_params.append(status_filter) p += 1 - if ( - group_by_session is True - and not is_v2 - and not is_request_id_lookup - and not is_search_lookup - and sort_by == "startTime" - ): + if session_page_route: return await _ui_session_grouped_spend_logs( prisma_client=prisma_client, sql_conditions=sql_conditions, @@ -2978,6 +3001,8 @@ async def ui_view_spend_logs( page_size=page_size, sort_desc=order_direction != "asc", session_cursor=session_cursor, + status_filtered=status_filter is not None and not is_search_lookup, + status_sessions_sql=status_sessions_sql if materialize_status_sessions else None, ) # Build the ORDER BY expression. ttft_ms is computed from @@ -3083,6 +3108,47 @@ class _SessionPageRow(TypedDict): last_activity: ReadOnly[str] +class _SessionScanRow(TypedDict): + session_key: ReadOnly[str | None] + api_key: ReadOnly[str | None] + last_activity: ReadOnly[str | None] + settled: ReadOnly[bool | None] + rows_scanned: ReadOnly[int] + session_count: ReadOnly[int] + oldest: ReadOnly[str | None] + carried: ReadOnly[bool | None] + + +# Each attempt reads this many times the first read's rows, ``None`` reading them all; an attempt is exact or discarded. +_SESSION_SCAN_GROWTH: Final[tuple[int | None, ...]] = (1, 8, 64, None) +# Carried sessions' keys, api keys and last activity, plus the oldest row of the read that carried them. +_CarryParams: TypeAlias = tuple[Sequence[str], Sequence[str], Sequence[str], str] | tuple[()] + + +def _read_falls_short(scan_rows: int | None, session_count: int, rows_read: int, sessions_needed: int) -> bool: + """Whether ``scan_rows`` newest rows, as dense as the ``rows_read`` that held ``session_count`` sessions, hold too few.""" + return scan_rows is not None and scan_rows * session_count < sessions_needed * rows_read + + +def _utc_timestamp_literal(bound: datetime) -> str: + """SQL literal for a UTC datetime; anything else is refused, so no caller text reaches the SQL.""" + text: Final = bound.isoformat() if type(bound) is datetime and bound.tzinfo is timezone.utc else "" + if not text or not set(text) <= set("0123456789-:.T+"): + raise TypeError(f"expected a UTC datetime, got {type(bound).__name__}") + return f"('{text}'::timestamptz AT TIME ZONE 'UTC')" + + +def _newest_session_rows_sql(where_clause: str, scan_rows: int | None, not_after: str | None) -> str: + """The newest ``scan_rows`` matching rows (all of them for ``None``), none newer than ``not_after``.""" + return f""" + SELECT {_SESSION_KEY_EXPR} AS session_key, api_key, "startTime" + FROM "LiteLLM_SpendLogs" + WHERE {where_clause} + {f'AND "startTime" <= {not_after}' if not_after is not None else ""} + {f'ORDER BY "startTime" DESC LIMIT {scan_rows}' if scan_rows is not None else ""} + """ + + def _parse_session_cursor(session_cursor: str | None) -> tuple[str, str, str] | None: if session_cursor is None or session_cursor.count("|") < 2: return None @@ -3099,19 +3165,26 @@ async def _fetch_session_representatives( sql_params: Sequence[object], next_param_index: int, session_keys: Sequence[tuple[str, str]], + status_filter: object | None = None, ) -> list[dict[str, object]]: # mutable-ok: _build_ui_spend_logs_response writes session counts onto each row - """Fetch the final agent outcome, or newest non-MCP row, of each ``(session_key, api_key)`` session, in ``session_keys`` order.""" + """Fetch the final agent outcome, or newest non-MCP row, of each ``(session_key, api_key)`` session, in ``session_keys`` order. + + With ``status_filter``, a session whose outcome no longer has that status is dropped, as the session status + filter would drop it in this statement's snapshot. + """ rep_query: Final = f""" SELECT * FROM ( SELECT DISTINCT ON ({_SESSION_GROUP_KEY_SQL}) {_SPEND_LOG_LIST_COLUMNS} FROM "LiteLLM_SpendLogs" WHERE {where_clause} + AND (session_id = ANY(${next_param_index}::text[]) OR request_id = ANY(${next_param_index}::text[])) AND ({_SESSION_GROUP_KEY_SQL}) IN ( SELECT * FROM unnest(${next_param_index}::text[], ${next_param_index + 1}::text[]) ) ORDER BY {_SESSION_GROUP_KEY_SQL}, {_SESSION_REPRESENTATIVE_ORDER_SQL} ) AS session_representatives + {f"WHERE COALESCE(status, 'success') = ${next_param_index + 2}" if status_filter is not None else ""} """ rep_rows: Final[Sequence[dict[str, object]]] = await _query_raw( # mutable-ok: rows are enriched in place prisma_client, @@ -3119,6 +3192,7 @@ async def _fetch_session_representatives( *sql_params, [session_key for session_key, _ in session_keys], [api_key for _, api_key in session_keys], + *((status_filter,) if status_filter is not None else ()), ) rep_by_key: Final[Mapping[tuple[str, str], dict[str, object]]] = MappingProxyType( # mutable-ok: same rows {(str(row["session_id"] or row["request_id"]), str(row["api_key"])): row for row in rep_rows} @@ -3126,32 +3200,6 @@ async def _fetch_session_representatives( return [rep_by_key[key] for key in session_keys if key in rep_by_key] -async def _count_grouped_sessions( - prisma_client: "PrismaClient", - where_clause: str, - sql_params: Sequence[object], - next_param_index: int, -) -> tuple[int, bool]: - """Count the sessions matching the filter, returning ``(total, total_is_capped)`` bounded by the count cap.""" - count_query: Final = f""" - SELECT COUNT(*) AS total_count - FROM ( - SELECT 1 - FROM "LiteLLM_SpendLogs" - WHERE {where_clause} - GROUP BY {_SESSION_GROUP_KEY_SQL} - LIMIT ${next_param_index} - ) AS bounded_sessions - """ - count_rows: Final[Sequence[_SpendLogsCountRow]] = await _query_raw( - prisma_client, count_query, *sql_params, SPEND_LOGS_PAGINATION_COUNT_CAP + 1 - ) - raw_total: Final = int(count_rows[0]["total_count"]) if count_rows else 0 - return ( - (SPEND_LOGS_PAGINATION_COUNT_CAP, True) if raw_total > SPEND_LOGS_PAGINATION_COUNT_CAP else (raw_total, False) - ) - - async def _ui_session_grouped_spend_logs( prisma_client: "PrismaClient", sql_conditions: Sequence[str], @@ -3161,6 +3209,8 @@ async def _ui_session_grouped_spend_logs( page_size: int, sort_desc: bool, session_cursor: str | None, + status_filtered: bool = False, + status_sessions_sql: str | None = None, ) -> Mapping[str, object]: """ One row per session, keyset-paginated by session last activity. @@ -3186,12 +3236,23 @@ async def _ui_session_grouped_spend_logs( starts past the end says nothing about the total, so that one is counted. """ where_clause: Final = " AND ".join(sql_conditions) if sql_conditions else "TRUE" + # The status filter keeps or drops whole sessions, so it holds for every row of a session already on the + # page; the probe and the representatives skip it rather than regroup the window once per session. The + # representatives re-check it on the page's own sessions, since a row written after the page read can change a + # session's outcome. + row_conditions: Final = sql_conditions[:-1] if status_filtered else sql_conditions + row_params: Final = sql_params[:-1] if status_filtered else sql_params + row_where_clause: Final = " AND ".join(row_conditions) if row_conditions else "TRUE" cmp_op: Final = "<" if sort_desc else ">" direction: Final = "DESC" if sort_desc else "ASC" + # ``status_sessions_sql`` (a failure filter's sessions, which the last condition reads) is computed once per read. + status_cte: Final = ( + f"status_sessions AS MATERIALIZED ({status_sessions_sql})," if status_sessions_sql is not None else "" + ) cursor: Final = _parse_session_cursor(session_cursor) - having_clause: Final = ( - f'HAVING (MAX("startTime"), {_SESSION_GROUP_KEY_SQL}) {cmp_op} ' + cursor_clause: Final = ( + f"WHERE (last_activity, session_key, api_key) {cmp_op} " f"(${next_param_index}::timestamp, ${next_param_index + 1}, ${next_param_index + 2})" if cursor else "" @@ -3203,21 +3264,158 @@ async def _ui_session_grouped_spend_logs( offset_params: Final[tuple[int, ...]] = (offset,) if offset and page_limit > 0 else () offset_clause: Final = f"OFFSET ${limit_index + 1}" if offset_params else "" - page_query: Final = f""" - SELECT {_SESSION_KEY_EXPR} AS session_key, - api_key, - MAX("startTime")::text AS last_activity - FROM "LiteLLM_SpendLogs" - WHERE {where_clause} - GROUP BY {_SESSION_GROUP_KEY_SQL} - {having_clause} - ORDER BY MAX("startTime") {direction}, {_SESSION_KEY_EXPR} {direction}, api_key {direction} - LIMIT ${limit_index} {offset_clause} - """ + # Read only the newest matching rows, growing the read until it decides the page and the capped count: + # every row newer than the oldest one read is in the read, so a session whose last activity is newer + # than that row is complete in it. A descending page is exact once its (page_limit + 1)th session is, + # and the count once the read holds more sessions than the cap (or the read holds every matching row). + # A session after a descending cursor has no row newer than the cursor, so that read starts at the + # cursor; a session with a newer row ranks before the cursor, so it is dropped after the sort (only the + # page's candidates are probed). If the sessions after the cursor do not pass the cap, the window's + # count is read afterwards from the newest rows the same way, without the page. + cap: Final = int(SPEND_LOGS_PAGINATION_COUNT_CAP) + # offset + page_limit never passes the cap, so the read stays bounded however large the page number is. + first_read: Final = cap + offset + page_limit + 2 + from_cursor: Final = cursor is not None and sort_desc + # A cursor page needs the session after it; any other page needs one session past the cap for its count. + sessions_needed: Final = page_limit + 1 if from_cursor else cap + 1 + # An offset page's read that holds no session of the page hands its sessions to the next read, which then reads + # only the rows at or before that read's oldest row instead of reading the newest rows again. + carry_index: Final = limit_index + 1 + len(offset_params) + carried: Sequence[_SessionScanRow] = () # rebind-ok: each read replaces the sessions carried to the next read + carry_params: _CarryParams = () # rebind-ok: each read replaces the carried sessions' query params + found: Sequence[_SessionScanRow] = () # rebind-ok: each read replaces the sessions found so far + session_count = 0 # rebind-ok: each read replaces the total session count + rows_read = 0 # rebind-ok: each read replaces the rows-read tally + for growth in _SESSION_SCAN_GROWTH if sort_desc else (None,): + scan_rows: int | None = None if growth is None else first_read * growth + # Skip a read that the last read's rows per session say still holds too few sessions to decide anything. + if _read_falls_short(scan_rows, session_count, rows_read, sessions_needed): + continue + read_rows: int | None = None if scan_rows is None else scan_rows - rows_read if carried else scan_rows + not_after: str | None = ( + f"${next_param_index}::timestamp" if from_cursor else f"${carry_index + 3}::timestamp" if carried else None + ) + sessions_after: str = ( + f"""( + SELECT session_key, api_key, last_activity + FROM sessions + {cursor_clause} + ORDER BY last_activity DESC, session_key DESC, api_key DESC + ) AS after_cursor + WHERE ( + SELECT 1 FROM "LiteLLM_SpendLogs" + WHERE {row_where_clause} + AND "startTime" > ${next_param_index}::timestamp + AND (session_id = after_cursor.session_key OR request_id = after_cursor.session_key) + AND {_SESSION_KEY_EXPR} = after_cursor.session_key + AND api_key = after_cursor.api_key + LIMIT 1 + ) IS NULL""" + if from_cursor + else f"sessions {cursor_clause}" + ) + carried_keys: str = f"${carry_index}::text[], ${carry_index + 1}::text[]" + page_query: str = f""" + WITH {status_cte} scanned AS MATERIALIZED ( + {_newest_session_rows_sql(where_clause, read_rows, not_after)} + ), + sessions AS ( + SELECT session_key, api_key, MAX("startTime") AS last_activity + FROM scanned + {f"WHERE (session_key, api_key) NOT IN (SELECT * FROM unnest({carried_keys}))" if carried else ""} + GROUP BY session_key, api_key + {f"UNION ALL SELECT * FROM unnest({carried_keys}, ${carry_index + 2}::timestamp[])" if carried else ""} + ), + edge AS ( + SELECT COUNT(*) AS rows_scanned, + MIN("startTime") AS oldest, + (SELECT COUNT(*) FROM sessions) AS session_count + FROM scanned + ) + SELECT page.session_key, + page.api_key, + page.last_activity::text AS last_activity, + page.last_activity > edge.oldest AS settled, + edge.rows_scanned, + edge.session_count, + edge.oldest::text AS oldest, + page.carried + FROM edge + LEFT JOIN LATERAL ( + ( + SELECT session_key, api_key, last_activity, FALSE AS carried + FROM {sessions_after} + ORDER BY last_activity {direction}, session_key {direction}, api_key {direction} + LIMIT ${limit_index} {offset_clause} + ) + { + f'''UNION ALL + SELECT session_key, api_key, last_activity, TRUE FROM sessions + WHERE edge.session_count <= ${limit_index + 1} AND edge.rows_scanned = {read_rows}''' + if offset_params and read_rows is not None + else "" + } + ) AS page ON TRUE + """ + scan: Sequence[_SessionScanRow] = await _query_raw( + prisma_client, + page_query, + *sql_params, + *cursor_params, + max(page_limit + 1, 0), + *offset_params, + *carry_params, + ) + found = [row for row in scan if row["session_key"] is not None] + session_count = int(scan[0]["session_count"]) + read_everything: bool = read_rows is None or int(scan[0]["rows_scanned"]) < read_rows + rows_read = int(scan[0]["rows_scanned"]) + (rows_read if carried else 0) + carried = [row for row in scan if row.get("carried")] + carry_params = ( + ( + [str(row["session_key"]) for row in carried], + [str(row["api_key"]) for row in carried], + [str(row["last_activity"]) for row in carried], + str(scan[0]["oldest"]), + ) + if carried + else () + ) + page_is_exact: bool = ( + page_limit <= 0 or read_everything or (len(found) > page_limit and bool(found[-1]["settled"])) + ) + if page_is_exact and (from_cursor or read_everything or session_count > cap): + break + if from_cursor and session_count <= cap: + # Skip to the first read that the last read's rows per session say holds more than the cap's sessions: + # each read is exact or discarded, so skipping a read changes how much is read, not the count. + for growth in _SESSION_SCAN_GROWTH: + count_rows: int | None = None if growth is None else first_read * growth + if _read_falls_short(count_rows, session_count, rows_read, cap + 1): + continue + count_query: str = f""" + WITH {status_cte} scanned AS MATERIALIZED ({_newest_session_rows_sql(where_clause, count_rows, None)}) + SELECT COUNT(*) AS rows_scanned, + (SELECT COUNT(*) FROM (SELECT 1 FROM scanned GROUP BY session_key, api_key) AS s) + AS session_count + FROM scanned + """ + counted: Sequence[Mapping[str, object]] = await _query_raw(prisma_client, count_query, *sql_params) + session_count = int(str(counted[0]["session_count"])) + rows_read = int(str(counted[0]["rows_scanned"])) + if count_rows is None or rows_read < count_rows or session_count > cap: + break page_rows: Final[Sequence[_SessionPageRow]] = ( () if page_limit <= 0 - else await _query_raw(prisma_client, page_query, *sql_params, *cursor_params, page_limit + 1, *offset_params) + else [ + { + "session_key": str(row["session_key"]), + "api_key": str(row["api_key"]), + "last_activity": str(row["last_activity"]), + } + for row in found + ] ) has_more: Final = len(page_rows) > page_limit @@ -3233,17 +3431,22 @@ async def _ui_session_grouped_spend_logs( total_records, total_is_capped = ( (offset + len(page_rows), False) if page_ends_the_list - else await _count_grouped_sessions(prisma_client, where_clause, sql_params, next_param_index) + else ( + (SPEND_LOGS_PAGINATION_COUNT_CAP, True) + if session_count > SPEND_LOGS_PAGINATION_COUNT_CAP + else (session_count, False) + ) ) session_keys: Final = tuple((row["session_key"], row["api_key"]) for row in visible_rows) data: Final[list[dict[str, object]]] = ( # mutable-ok: _build_ui_spend_logs_response writes onto each row await _fetch_session_representatives( prisma_client=prisma_client, - where_clause=where_clause, - sql_params=sql_params, - next_param_index=next_param_index, + where_clause=row_where_clause, + sql_params=row_params, + next_param_index=len(row_params) + 1, session_keys=session_keys, + status_filter=sql_params[-1] if status_filtered else None, ) if session_keys else [] diff --git a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py index 8d53817f51e..ef7b6ac1b3f 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py @@ -3,9 +3,10 @@ import collections import datetime import hashlib import json +import operator import re import sqlite3 -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from datetime import timezone from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock, patch @@ -7176,12 +7177,20 @@ def test_parse_session_cursor(): SESSION_GROUP_KEY_SQL = "COALESCE(NULLIF(session_id, ''), request_id), api_key" +def _session_scan_rows(page_rows, session_count): + """Rows of the session page query: the page joined to the read's edge, as if the read held every matching row.""" + edge = {"rows_scanned": 0, "session_count": session_count} + if not page_rows: + return [{"session_key": None, "api_key": None, "last_activity": None, "settled": None, **edge}] + return [{**row, "settled": True, **edge} for row in page_rows] + + def _session_grouped_mock_prisma(session_page_rows, session_total, representative_rows): """Mock prisma dispatching the raw queries the keyset grouped path emits.""" async def mock_query_raw(sql_query, *params): - if "COUNT(*) AS total_count" in sql_query: - return [{"total_count": session_total}] + if "rows_scanned" in sql_query: + return _session_scan_rows(session_page_rows, session_total) if "DISTINCT ON" in sql_query: return representative_rows if "COALESCE(SUM(spend)" in sql_query: @@ -7199,7 +7208,7 @@ def _session_grouped_mock_prisma(session_page_rows, session_total, representativ "session_models": ["gpt-4o"], } ] - return session_page_rows + raise AssertionError(f"unexpected query:\n{sql_query}") mock_prisma = MagicMock() mock_prisma.db = MagicMock() @@ -7232,8 +7241,6 @@ def _session_grouped_paginating_prisma(sessions, counted_total=None): """Mock prisma serving the grouped page query out of ``sessions``, honoring the LIMIT and OFFSET it asks for.""" async def mock_query_raw(sql_query, *params): - if "COUNT(*) AS total_count" in sql_query: - return [{"total_count": min(len(sessions) if counted_total is None else counted_total, params[-1])}] if "DISTINCT ON" in sql_query: return [_session_representative_row(f"req-{session_key}", session_key) for session_key in params[-2]] if "COALESCE(SUM(spend)" in sql_query: @@ -7241,10 +7248,13 @@ def _session_grouped_paginating_prisma(sessions, counted_total=None): bounds = re.search(r"LIMIT \$(\d+)(?: OFFSET \$(\d+))?", sql_query) limit = params[int(bounds.group(1)) - 1] offset = params[int(bounds.group(2)) - 1] if bounds.group(2) else 0 - return [ - _session_page_row(session_key, last_activity) - for session_key, last_activity in sessions[offset : offset + limit] - ] + return _session_scan_rows( + [ + _session_page_row(session_key, last_activity) + for session_key, last_activity in sessions[offset : offset + limit] + ], + len(sessions) if counted_total is None else counted_total, + ) mock_prisma = MagicMock() mock_prisma.db = MagicMock() @@ -7296,30 +7306,31 @@ async def test_ui_view_spend_logs_group_by_session_first_page(client, monkeypatc assert data["data"][1]["session_total_count"] == 1 emitted = [call.args for call in mock_prisma.db.query_raw.await_args_list] + assert len(emitted) == 3, "one query decides the page and the count; no separate count query" page_query_sql = emitted[0][0] - assert f"GROUP BY {SESSION_GROUP_KEY_SQL}" in page_query_sql + assert "GROUP BY session_key, api_key" in page_query_sql + assert 'ORDER BY "startTime" DESC LIMIT' in page_query_sql, "the first read takes only the newest rows" assert "OFFSET" not in page_query_sql - assert "HAVING" not in page_query_sql + assert "(last_activity, session_key, api_key)" not in page_query_sql assert emitted[0][-1] == 3, "page query fetches page_size + 1 sessions to detect has_more" - count_sql = emitted[1][0] - assert f"GROUP BY {SESSION_GROUP_KEY_SQL}" in count_sql - assert "LIMIT" in count_sql and "FROM (" in count_sql, "the grouped count must stay bounded" - - rep_call = emitted[2] + rep_call = emitted[1] assert f"DISTINCT ON ({SESSION_GROUP_KEY_SQL})" in rep_call[0] assert ( f"ORDER BY {SESSION_GROUP_KEY_SQL}, " + spend_management_endpoints._SESSION_REPRESENTATIVE_ORDER_SQL ) in rep_call[0] assert rep_call[-2] == ["sess-1", "req-solo"] assert rep_call[-1] == ["hashed-key", "hashed-key"] + assert "session_id = ANY(" in rep_call[0] and "request_id = ANY(" in rep_call[0], ( + "representatives must narrow by the indexed id columns before grouping the window" + ) finally: app.dependency_overrides.pop(ps.user_api_key_auth, None) @pytest.mark.asyncio async def test_ui_view_spend_logs_group_by_session_cursor_page(client, monkeypatch): - """A session_cursor becomes a HAVING keyset predicate instead of an OFFSET.""" + """A session_cursor becomes a keyset predicate on the sessions instead of an OFFSET.""" page_rows = [_session_page_row("sess-2", "2026-08-29 07:00:00")] reps = [_session_representative_row("req-2", "sess-2")] mock_prisma = _session_grouped_mock_prisma(page_rows, 3, reps) @@ -7352,7 +7363,7 @@ async def test_ui_view_spend_logs_group_by_session_cursor_page(client, monkeypat page_query_call = mock_prisma.db.query_raw.await_args_list[0] page_query_sql = page_query_call.args[0] - assert f'HAVING (MAX("startTime"), {SESSION_GROUP_KEY_SQL}) <' in page_query_sql + assert "WHERE (last_activity, session_key, api_key) <" in page_query_sql assert "OFFSET" not in page_query_sql cursor_index = page_query_call.args.index("2026-08-29 09:00:00") assert page_query_call.args[cursor_index : cursor_index + 3] == ( @@ -7364,6 +7375,746 @@ async def test_ui_view_spend_logs_group_by_session_cursor_page(client, monkeypat app.dependency_overrides.pop(ps.user_api_key_auth, None) +# The first read of a page_size=1 first page: the cap's sessions plus the page and the session after it. +_FULL_FIRST_READ = spend_management_endpoints.SPEND_LOGS_PAGINATION_COUNT_CAP + 3 + + +@pytest.mark.parametrize( + "first_read, first_read_sessions, expected_reads, expected_total", + [ + pytest.param( + {"rows_scanned": _FULL_FIRST_READ, "session_count": 10_001, "settled": False}, + 2, + 2, + (10_000, True), + id="unsettled-page-reads-further", + ), + pytest.param( + {"rows_scanned": _FULL_FIRST_READ, "session_count": 3, "settled": True}, + 2, + 2, + (4, False), + id="settled-page-under-the-cap-reads-further-for-the-count", + ), + pytest.param( + {"rows_scanned": _FULL_FIRST_READ, "session_count": 10_001, "settled": True}, + 2, + 1, + (10_000, True), + id="settled-page-past-the-cap-stops", + ), + pytest.param( + {"rows_scanned": _FULL_FIRST_READ, "session_count": 10_001, "settled": True}, + 1, + 2, + (10_000, True), + id="read-without-the-next-session-reads-further", + ), + pytest.param( + {"rows_scanned": _FULL_FIRST_READ - 1, "session_count": 3, "settled": False}, + 2, + 1, + (3, False), + id="short-read-holds-every-row", + ), + pytest.param( + {"rows_scanned": _FULL_FIRST_READ - 1, "session_count": 10_000, "settled": True}, + 2, + 1, + (10_000, False), + id="exactly-the-cap-is-not-capped", + ), + ], +) +@pytest.mark.asyncio +async def test_ui_view_spend_logs_group_by_session_reads_newest_rows_until_exact( + client, monkeypatch, first_read, first_read_sessions, expected_reads, expected_total +): + """The page reads the newest rows first and only trusts that read once it decides both the page and the count.""" + monkeypatch.setattr(spend_management_endpoints, "_SESSION_SCAN_GROWTH", (1, None)) + page_rows = [ + _session_page_row("sess-1", "2026-08-29 10:00:00"), + _session_page_row("sess-2", "2026-08-29 09:00:00"), + ] + reads = [] + + async def mock_query_raw(sql_query, *params): + if "rows_scanned" in sql_query: + reads.append(sql_query) + if f'ORDER BY "startTime" DESC LIMIT {_FULL_FIRST_READ}' in sql_query: + return [{**row, **first_read} for row in page_rows[:first_read_sessions]] + whole_count = max(first_read["session_count"], 4) + edge = {"settled": True, "rows_scanned": 40, "session_count": whole_count} + return [{**row, **edge} for row in page_rows] + if "DISTINCT ON" in sql_query: + return [_session_representative_row("req-1", "sess-1")] + return [] + + mock_prisma = MagicMock() + mock_prisma.db = MagicMock() + mock_prisma.db.query_raw = AsyncMock(side_effect=mock_query_raw) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + monkeypatch.setattr( + "litellm.proxy.spend_tracking.spend_management_endpoints._is_admin_view_safe", + lambda user_api_key_dict: True, + ) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user" + ) + try: + start_date, end_date = _default_date_range() + response = client.get( + "/spend/logs/ui", + params={"start_date": start_date, "end_date": end_date, "group_by_session": "true", "page_size": 1}, + headers={"Authorization": "Bearer sk-test"}, + ) + assert response.status_code == 200, response.text + data = response.json() + assert len(reads) == expected_reads + assert (data["total"], data["total_is_capped"]) == expected_total + assert data["next_session_cursor"] == "2026-08-29 10:00:00|hashed-key|sess-1" + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.parametrize( + "params, first_read_rows, bounded_by_cursor", + [ + pytest.param( + {"page_size": 2}, spend_management_endpoints.SPEND_LOGS_PAGINATION_COUNT_CAP + 4, False, id="first-page" + ), + pytest.param( + {"page_size": 2, "page": 3}, + spend_management_endpoints.SPEND_LOGS_PAGINATION_COUNT_CAP + 8, + False, + id="offset-page", + ), + pytest.param( + {"page_size": 1000, "page": 20}, + 2 * spend_management_endpoints.SPEND_LOGS_PAGINATION_COUNT_CAP + 2, + False, + id="page-past-the-cap", + ), + pytest.param( + {"page_size": 1000, "page": 9_300_000_000_000_000}, + 2 * spend_management_endpoints.SPEND_LOGS_PAGINATION_COUNT_CAP + 2, + False, + id="page-number-near-the-bigint-limit", + ), + pytest.param( + {"page_size": 2, "session_cursor": "2026-08-29 09:00:00|hashed-key|req-solo"}, + spend_management_endpoints.SPEND_LOGS_PAGINATION_COUNT_CAP + 4, + True, + id="descending-cursor-page", + ), + pytest.param( + {"page_size": 2, "sort_order": "asc", "session_cursor": "2026-08-29 09:00:00|hashed-key|req-solo"}, + None, + False, + id="ascending-reads-every-row", + ), + ], +) +@pytest.mark.asyncio +async def test_ui_view_spend_logs_group_by_session_read_is_bounded( + client, monkeypatch, params, first_read_rows, bounded_by_cursor +): + """The first read takes only as many newest rows as its depth needs, never more than twice the cap however large + the page number, and a descending cursor page starts it at the cursor.""" + mock_prisma = _session_grouped_mock_prisma([_session_page_row("sess-2", "2026-08-29 07:00:00")], 3, []) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + monkeypatch.setattr( + "litellm.proxy.spend_tracking.spend_management_endpoints._is_admin_view_safe", + lambda user_api_key_dict: True, + ) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user" + ) + try: + start_date, end_date = _default_date_range() + response = client.get( + "/spend/logs/ui", + params={"start_date": start_date, "end_date": end_date, "group_by_session": "true", **params}, + headers={"Authorization": "Bearer sk-test"}, + ) + assert response.status_code == 200, response.text + page_query = mock_prisma.db.query_raw.await_args_list[0].args + sql = page_query[0] + scanned = sql[: sql.index("sessions AS (")] + if first_read_rows is None: + assert 'ORDER BY "startTime"' not in scanned, "an ascending page reads every matching row" + else: + assert f'ORDER BY "startTime" DESC LIMIT {first_read_rows}' in scanned + bounds = re.findall(r'"startTime" ([<>]=?) \$(\d+)::timestamp', sql) + if not bounded_by_cursor: + assert bounds == [] + return + assert {(op, page_query[int(index)]) for op, index in bounds} == { + ("<=", "2026-08-29 09:00:00"), + (">", "2026-08-29 09:00:00"), + }, "the read starts at the cursor and drops sessions with a newer row" + assert re.search(r'AND "startTime" <= \$\d+::timestamp\s+ORDER BY "startTime" DESC LIMIT', scanned), ( + "the cursor bound must sit inside the newest-rows read, before its LIMIT" + ) + count_query = mock_prisma.db.query_raw.await_args_list[1].args[0] + assert "rows_scanned" in count_query and "LATERAL" not in count_query + assert f'ORDER BY "startTime" DESC LIMIT {first_read_rows}' in count_query + assert not re.search(r"\$\d+::timestamp\s", count_query), "the window's count is not bounded by the cursor" + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.parametrize( + "page_read_rows, page_read_sessions, window_rows, count_limits", + [ + pytest.param(100, 100, 7, [10_004], id="short-read-of-one-row-sessions"), + pytest.param(10_004, 10_000, 7, [80_032], id="full-read-just-short-of-the-cap"), + pytest.param(10_004, 1_000, 7, [640_256], id="ten-row-sessions"), + pytest.param(10_004, 3, 7, [None], id="long-sessions"), + pytest.param(100, 100, 100_000, [10_004, None], id="a-full-count-read-grows-by-its-own-density"), + ], +) +@pytest.mark.asyncio +async def test_ui_view_spend_logs_group_by_session_cursor_count_starts_at_the_page_reads_density( + client, monkeypatch, page_read_rows, page_read_sessions, window_rows, count_limits +): + """A cursor page counts the window from the first read its own read's rows per session say can pass the cap.""" + page_rows = [ + _session_page_row("sess-1", "2026-08-29 08:00:00"), + _session_page_row("sess-2", "2026-08-29 07:00:00"), + _session_page_row("sess-3", "2026-08-29 06:00:00"), + ] + limits = [] + + async def mock_query_raw(sql_query, *params): + if "AS page ON TRUE" in sql_query: + edge = {"settled": True, "rows_scanned": page_read_rows, "session_count": page_read_sessions} + return [{**row, **edge} for row in page_rows] + if "rows_scanned" in sql_query: + limit = re.search(r'ORDER BY "startTime" DESC LIMIT (\d+)', sql_query) + limits.append(int(limit.group(1)) if limit else None) + return [ + {"rows_scanned": min(int(limit.group(1)) if limit else window_rows, window_rows), "session_count": 5} + ] + if "DISTINCT ON" in sql_query: + return [_session_representative_row("req-1", "sess-1")] + return [] + + mock_prisma = MagicMock() + mock_prisma.db = MagicMock() + mock_prisma.db.query_raw = AsyncMock(side_effect=mock_query_raw) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + monkeypatch.setattr( + "litellm.proxy.spend_tracking.spend_management_endpoints._is_admin_view_safe", + lambda user_api_key_dict: True, + ) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user" + ) + try: + start_date, end_date = _default_date_range() + response = client.get( + "/spend/logs/ui", + params={ + "start_date": start_date, + "end_date": end_date, + "group_by_session": "true", + "page_size": 2, + "session_cursor": "2026-08-29 09:00:00|hashed-key|req-solo", + }, + headers={"Authorization": "Bearer sk-test"}, + ) + assert response.status_code == 200, response.text + assert limits == count_limits + page_reads = [call for call in mock_prisma.db.query_raw.await_args_list if "AS page ON TRUE" in call.args[0]] + assert len(page_reads) == 1, "a settled cursor page is not read again for the count" + assert (response.json()["total"], response.json()["total_is_capped"]) == (5, False) + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.parametrize( + "rows_per_session, growths", + [ + pytest.param(7, [1, 8, 64, None], id="dense-enough-for-every-read-to-pass-the-cap"), + pytest.param(5_000, [1, None], id="long-sessions-skip-reads-that-cannot-pass-the-cap"), + ], +) +@pytest.mark.asyncio +async def test_ui_view_spend_logs_group_by_session_read_grows_geometrically( + client, monkeypatch, rows_per_session, growths +): + """A read that settles nothing grows 8x, then 64x, then reads every row, so a page takes at most four reads; a + read its predecessor's rows per session say cannot hold more sessions than the cap is skipped.""" + page_row = _session_page_row("sess-1", "2026-08-29 10:00:00") + limits = [] + + async def mock_query_raw(sql_query, *params): + if "rows_scanned" in sql_query: + limit = re.search(r'ORDER BY "startTime" DESC LIMIT (\d+)', sql_query) + limits.append(int(limit.group(1)) if limit else None) + rows = int(limit.group(1)) if limit else 7 + return [{**page_row, "settled": False, "rows_scanned": rows, "session_count": rows // rows_per_session}] + if "DISTINCT ON" in sql_query: + return [_session_representative_row("req-1", "sess-1")] + return [] + + mock_prisma = MagicMock() + mock_prisma.db = MagicMock() + mock_prisma.db.query_raw = AsyncMock(side_effect=mock_query_raw) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + monkeypatch.setattr( + "litellm.proxy.spend_tracking.spend_management_endpoints._is_admin_view_safe", + lambda user_api_key_dict: True, + ) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user" + ) + try: + start_date, end_date = _default_date_range() + response = client.get( + "/spend/logs/ui", + params={"start_date": start_date, "end_date": end_date, "group_by_session": "true", "page_size": 2}, + headers={"Authorization": "Bearer sk-test"}, + ) + assert response.status_code == 200, response.text + first = spend_management_endpoints.SPEND_LOGS_PAGINATION_COUNT_CAP + 4 + assert limits == [None if growth is None else growth * first for growth in growths] + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +class _NewestRowsSessionTable: + """Serves the descending session page query from in-memory rows the way Postgres runs it, recording how many + rows each read takes; ``session_rows`` gives each session's row count, oldest session first, rows + ``seconds_apart`` apart. Rows sharing a time come back lowest session key first, the opposite of the page order, + as a LIMIT may cut ties anywhere.""" + + def __init__(self, session_rows: Sequence[int], seconds_apart: int = 1) -> None: + start = datetime.datetime(2026, 10, 11) + firsts = [sum(session_rows[:session]) for session in range(len(session_rows))] + self.rows: Final = sorted( + sorted( + (start + datetime.timedelta(seconds=(first + row) * seconds_apart), f"sess-{session:05d}") + for session, (first, rows) in enumerate(zip(firsts, session_rows)) + for row in range(rows) + ), + key=lambda row: row[0], + reverse=True, + ) + self.rows_read: Final[list[int]] = [] + + async def query_raw(self, sql_query: str, *params: object) -> list[dict[str, object]]: + if "DISTINCT ON" in sql_query: + return [_session_representative_row(f"req-{key}", key) for key in params[-2]] + if "COALESCE(SUM(spend)" in sql_query: + return [] + + def param(pattern: str) -> object: + found = re.search(pattern, sql_query) + return params[int(found.group(1)) - 1] if found else None + + not_after = param(r'"startTime" <= \$(\d+)::timestamp') + newest = [row for row in self.rows if not_after is None or row[0] <= datetime.datetime.fromisoformat(not_after)] + limit = re.search(r'ORDER BY "startTime" DESC LIMIT (\d+)', sql_query) + scanned = newest[: int(limit.group(1))] if limit else newest + self.rows_read.append(len(scanned)) + carried = re.search( + r"UNION ALL SELECT \* FROM unnest\(\$(\d+)::text\[\], \$\d+::text\[\], \$\d+::timestamp\[\]\)\s*\),\s*edge AS", + sql_query, + ) + assert carried or "UNION ALL SELECT * FROM unnest" not in sql_query, "unrecognised carried sessions SQL" + settled = re.search(r"page\.last_activity (>=|>) edge\.oldest AS settled", sql_query) + assert settled, "unrecognised settled SQL" + is_settled = operator.ge if settled.group(1) == ">=" else operator.gt + excluded = re.search(r"NOT IN \(SELECT \* FROM unnest\(\$(\d+)::text\[\]", sql_query) + grouped: dict[str, datetime.datetime] = {} + for started, key in scanned: + if not excluded or key not in params[int(excluded.group(1)) - 1]: + grouped[key] = max(grouped.get(key, started), started) + sessions = [(at, key) for key, at in grouped.items()] + if carried: + keys, _, activities = params[int(carried.group(1)) - 1 : int(carried.group(1)) + 2] + sessions += [(datetime.datetime.fromisoformat(at), key) for key, at in zip(keys, activities)] + oldest = scanned[-1][0] + ranked = sorted(sessions, reverse=True) + offset = param(r"OFFSET \$(\d+)") or 0 + page = [(at, key, False) for at, key in ranked[offset : offset + param(r"LIMIT \$(\d+)")]] + carry = re.search(r"edge\.session_count <= \$(\d+) AND edge\.rows_scanned = (\d+)", sql_query) + if carry and len(sessions) <= params[int(carry.group(1)) - 1] and len(scanned) == int(carry.group(2)): + page += [(at, key, True) for at, key in ranked] + edge = {"rows_scanned": len(scanned), "session_count": len(sessions), "oldest": str(oldest)} + return [ + { + "session_key": key, + "api_key": "hashed-key", + "last_activity": str(at), + "settled": is_settled(at, oldest), + "carried": is_carried, + **edge, + } + for at, key, is_carried in page + ] or [{"session_key": None, "api_key": None, "last_activity": None, "settled": None, "carried": None, **edge}] + + +@pytest.mark.parametrize("page, base_reads", [(1, 2), (20, 2), (40, 1)], ids=["first-page", "middle-page", "last-page"]) +@pytest.mark.asyncio +async def test_ui_view_spend_logs_group_by_session_busy_key_reads_no_more_than_grouping_every_row( + client, monkeypatch, page, base_reads +): + """On a busy key (2000 sessions of 50 rows) a page reads no more rows than grouping every matching row did, once + for the page and once for the count unless the page ends the list: a read too sparse to pass the cap is skipped, + and a read that holds no session of the page hands its sessions on, so the next read continues below it.""" + table = _NewestRowsSessionTable([50] * 2_000) + mock_prisma = MagicMock() + mock_prisma.db = MagicMock() + mock_prisma.db.query_raw = AsyncMock(side_effect=table.query_raw) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + monkeypatch.setattr( + "litellm.proxy.spend_tracking.spend_management_endpoints._is_admin_view_safe", + lambda user_api_key_dict: True, + ) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user" + ) + try: + start_date, end_date = _default_date_range() + response = client.get( + "/spend/logs/ui", + params={ + "start_date": start_date, + "end_date": end_date, + "group_by_session": "true", + "page_size": 50, + "page": page, + }, + headers={"Authorization": "Bearer sk-test"}, + ) + assert response.status_code == 200, response.text + data = response.json() + newest = 1_999 - (page - 1) * 50 + assert [row["session_id"] for row in data["data"]] == [f"sess-{n:05d}" for n in range(newest, newest - 50, -1)] + assert (data["total"], data["total_is_capped"]) == (2_000, False) + # A read that continues below another re-reads only the one row at that read's oldest time. + assert sum(table.rows_read) <= base_reads * len(table.rows) + len(table.rows_read) - 1, table.rows_read + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.asyncio +async def test_ui_view_spend_logs_group_by_session_handed_on_reads_take_only_the_rows_left(client, monkeypatch): + """Under a cap of 100, page 9 needs 81 sessions: 64 three-row sessions fill the first read, the next read finds + only 14 of the 100-row sessions below it, so the third read takes just the rest of its 64x budget.""" + monkeypatch.setattr(spend_management_endpoints, "SPEND_LOGS_PAGINATION_COUNT_CAP", 100) + table = _NewestRowsSessionTable([100] * 200 + [3] * 64) + mock_prisma = MagicMock() + mock_prisma.db = MagicMock() + mock_prisma.db.query_raw = AsyncMock(side_effect=table.query_raw) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + monkeypatch.setattr( + "litellm.proxy.spend_tracking.spend_management_endpoints._is_admin_view_safe", + lambda user_api_key_dict: True, + ) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user" + ) + try: + start_date, end_date = _default_date_range() + response = client.get( + "/spend/logs/ui", + params={ + "start_date": start_date, + "end_date": end_date, + "group_by_session": "true", + "page_size": 10, + "page": 9, + }, + headers={"Authorization": "Bearer sk-test"}, + ) + assert response.status_code == 200, response.text + data = response.json() + assert [row["session_id"] for row in data["data"]] == [f"sess-{n:05d}" for n in range(183, 173, -1)] + assert (data["total"], data["total_is_capped"]) == (100, True) + first = 100 + 80 + 10 + 2 + assert table.rows_read == [first, 8 * first - first, 64 * first - 8 * first] + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.asyncio +async def test_ui_view_spend_logs_group_by_session_session_at_the_reads_oldest_time_reads_further(client, monkeypatch): + """Under a cap of 10, 30 one-row sessions share one time and the first read takes 17 of them: a session at the + read's oldest time may rank after a tied session the read missed, so the page reads every row before it is exact.""" + monkeypatch.setattr(spend_management_endpoints, "SPEND_LOGS_PAGINATION_COUNT_CAP", 10) + table = _NewestRowsSessionTable([1] * 30, seconds_apart=0) + mock_prisma = MagicMock() + mock_prisma.db = MagicMock() + mock_prisma.db.query_raw = AsyncMock(side_effect=table.query_raw) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + monkeypatch.setattr( + "litellm.proxy.spend_tracking.spend_management_endpoints._is_admin_view_safe", + lambda user_api_key_dict: True, + ) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user" + ) + try: + start_date, end_date = _default_date_range() + response = client.get( + "/spend/logs/ui", + params={ + "start_date": start_date, + "end_date": end_date, + "group_by_session": "true", + "page_size": 5, + "page": 1, + }, + headers={"Authorization": "Bearer sk-test"}, + ) + assert response.status_code == 200, response.text + data = response.json() + assert [row["session_id"] for row in data["data"]] == [f"sess-{n:05d}" for n in range(29, 24, -1)] + assert (data["total"], data["total_is_capped"]) == (10, True) + assert table.rows_read == [10 + 5 + 2, 30] + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.asyncio +async def test_ui_view_spend_logs_group_by_session_ascending_reads_every_row(client, monkeypatch): + """The oldest sessions are not among the newest rows, so an ascending page reads every matching row at once.""" + page_rows = [_session_page_row("sess-1", "2026-08-29 08:00:00")] + mock_prisma = _session_grouped_mock_prisma(page_rows, 1, [_session_representative_row("req-1", "sess-1")]) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + monkeypatch.setattr( + "litellm.proxy.spend_tracking.spend_management_endpoints._is_admin_view_safe", + lambda user_api_key_dict: True, + ) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user" + ) + try: + start_date, end_date = _default_date_range() + response = client.get( + "/spend/logs/ui", + params={ + "start_date": start_date, + "end_date": end_date, + "group_by_session": "true", + "sort_order": "asc", + "page_size": 2, + }, + headers={"Authorization": "Bearer sk-test"}, + ) + assert response.status_code == 200, response.text + assert [row["request_id"] for row in response.json()["data"]] == ["req-1"] + reads = [call.args[0] for call in mock_prisma.db.query_raw.await_args_list if "rows_scanned" in call.args[0]] + assert len(reads) == 1 + assert 'ORDER BY "startTime" DESC LIMIT' not in reads[0] + assert "ORDER BY last_activity ASC" in reads[0] + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +class _ShoutingDatetime(datetime.datetime): + def isoformat(self, sep: str = "T", timespec: str = "auto") -> str: + return "2026-01-01'); DROP TABLE x; --" + + +@pytest.mark.parametrize( + "bound", + [ + "2026-01-01'); DROP TABLE x; --", + datetime.datetime(2026, 1, 1), + datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone(datetime.timedelta(hours=5))), + _ShoutingDatetime(2026, 1, 1, tzinfo=timezone.utc), + datetime.date(2026, 1, 1), + ], + ids=["text", "naive", "not-utc", "datetime-subclass", "date"], +) +def test_utc_timestamp_literal_refuses_anything_but_a_utc_datetime(bound): + with pytest.raises(TypeError): + spend_management_endpoints._utc_timestamp_literal(bound) + + +def test_utc_timestamp_literal_formats_a_utc_datetime(): + literal = spend_management_endpoints._utc_timestamp_literal( + datetime.datetime(2026, 1, 2, 3, 4, 5, 6, tzinfo=timezone.utc) + ) + assert literal == "('2026-01-02T03:04:05.000006+00:00'::timestamptz AT TIME ZONE 'UTC')" + + +@pytest.mark.asyncio +async def test_ui_view_spend_logs_group_by_session_status_regroups_only_the_read(client, monkeypatch): + """The status filter keeps whole sessions, so the cursor probe and the representatives of the page's sessions skip its subquery.""" + page_rows = [_session_page_row("sess-1", "2026-08-29 08:00:00")] + mock_prisma = _session_grouped_mock_prisma(page_rows, 1, [_session_representative_row("req-1", "sess-1")]) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + monkeypatch.setattr( + "litellm.proxy.spend_tracking.spend_management_endpoints._is_admin_view_safe", + lambda user_api_key_dict: True, + ) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user" + ) + try: + start_date, end_date = _default_date_range() + response = client.get( + "/spend/logs/ui", + params={ + "start_date": start_date, + "end_date": end_date, + "group_by_session": "true", + "status_filter": "failure", + "model": "gpt-4o", + "session_cursor": "2026-08-29 09:00:00|hashed-key|req-solo", + }, + headers={"Authorization": "Bearer sk-test"}, + ) + assert response.status_code == 200, response.text + calls = [call.args for call in mock_prisma.db.query_raw.await_args_list] + page = next(args for args in calls if "AS page ON TRUE" in args[0]) + scanned, probe = page[0].split("AS after_cursor", 1) + assert scanned.count("AS session_outcomes") == 1, "the read itself keeps only failed sessions" + assert "session_outcomes" not in probe and "model = $3" in probe + assert "failure" in page + rep = next(args for args in calls if "AS session_representatives" in args[0]) + assert "session_outcomes" not in rep[0] and "model = $3" in rep[0] + assert re.search(r"ANY\(\$4::text\[\]\)", rep[0]) and rep[-3:] == (["sess-1"], ["hashed-key"], "failure") + assert "COALESCE(status, 'success') = $6" in rep[0], "the page's outcomes are re-checked, not the window's" + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.asyncio +async def test_ui_view_spend_logs_group_by_session_status_inlines_the_window(client, monkeypatch): + """The status subquery holds the window as literals: Prisma's reused statements get generic plans from the 6th run, and a $N window there is planned as tiny.""" + page_rows = [_session_page_row("sess-1", "2026-08-29 08:00:00")] + mock_prisma = _session_grouped_mock_prisma(page_rows, 1, [_session_representative_row("req-1", "sess-1")]) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + monkeypatch.setattr( + "litellm.proxy.spend_tracking.spend_management_endpoints._is_admin_view_safe", + lambda user_api_key_dict: True, + ) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user" + ) + try: + start_date, end_date = _default_date_range() + response = client.get( + "/spend/logs/ui", + params={ + "start_date": start_date, + "end_date": end_date, + "group_by_session": "true", + "status_filter": "failure", + }, + headers={"Authorization": "Bearer sk-test"}, + ) + assert response.status_code == 200, response.text + read = next(call for call in mock_prisma.db.query_raw.await_args_list if "rows_scanned" in call.args[0]) + before, outcomes = read.args[0].split("SELECT DISTINCT ON", 1) + outcomes, after = outcomes.split("AS session_outcomes", 1) + outer = before + after + for bound in (start_date, end_date): + literal = datetime.datetime.strptime(bound, "%Y-%m-%d %H:%M:%S").replace(tzinfo=timezone.utc).isoformat() + assert f"('{literal}'::timestamptz AT TIME ZONE 'UTC')" in outcomes + assert "$1::timestamptz" not in outcomes and "$2::timestamptz" not in outcomes + assert "$1::timestamptz" in outer and "$2::timestamptz" in outer + assert read.args[1:3] == tuple( + datetime.datetime.strptime(bound, "%Y-%m-%d %H:%M:%S").replace(tzinfo=timezone.utc) + for bound in (start_date, end_date) + ) + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status_filter", ["failure", "success"]) +async def test_ui_view_spend_logs_group_by_session_materializes_only_failed_sessions( + client, monkeypatch, status_filter +): + """A failure filter's sessions are read once as a MATERIALIZED CTE (inlined, each parallel worker re-sorts the window); success stays inlined.""" + page_rows = [_session_page_row("sess-1", "2026-08-29 08:00:00")] + mock_prisma = _session_grouped_mock_prisma(page_rows, 1, [_session_representative_row("req-1", "sess-1")]) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + monkeypatch.setattr( + "litellm.proxy.spend_tracking.spend_management_endpoints._is_admin_view_safe", + lambda user_api_key_dict: True, + ) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user" + ) + try: + start_date, end_date = _default_date_range() + response = client.get( + "/spend/logs/ui", + params={ + "start_date": start_date, + "end_date": end_date, + "group_by_session": "true", + "status_filter": status_filter, + "model": "gpt-4o", + "session_cursor": "2026-08-29 09:00:00|hashed-key|req-solo", + }, + headers={"Authorization": "Bearer sk-test"}, + ) + assert response.status_code == 200, response.text + assert [row["session_id"] for row in response.json()["data"]] == ["sess-1"] + reads = [call.args for call in mock_prisma.db.query_raw.await_args_list if "rows_scanned" in call.args[0]] + assert len(reads) == 2, "the cursor page's read and the window count both carry the status filter" + for read in reads: + assert read[0].count("AS session_outcomes") == 1 + assert "COALESCE(status, 'success') = $4" in read[0] and read[4] == status_filter + if status_filter == "failure": + assert read[0].lstrip().startswith("WITH status_sessions AS MATERIALIZED (") + assert "IN (SELECT session_key, api_key FROM status_sessions)" in read[0] + else: + assert "status_sessions" not in read[0] + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.asyncio +async def test_ui_view_spend_logs_group_by_session_failure_off_the_session_page_stays_inlined(client, monkeypatch): + """Only the session page declares the status_sessions CTE, so other grouped sorts keep the inlined failure subquery.""" + queries = [] + + async def mock_query_raw(sql_query, *params): + queries.append(sql_query) + return [{"total_count": 0}] if "total_count" in sql_query else [] + + mock_prisma = MagicMock() + mock_prisma.db = MagicMock() + mock_prisma.db.query_raw = AsyncMock(side_effect=mock_query_raw) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + monkeypatch.setattr( + "litellm.proxy.spend_tracking.spend_management_endpoints._is_admin_view_safe", + lambda user_api_key_dict: True, + ) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user" + ) + try: + start_date, end_date = _default_date_range() + response = client.get( + "/spend/logs/ui", + params={ + "start_date": start_date, + "end_date": end_date, + "group_by_session": "true", + "status_filter": "failure", + "sort_by": "spend", + }, + headers={"Authorization": "Bearer sk-test"}, + ) + assert response.status_code == 200, response.text + assert queries and all("status_sessions" not in query for query in queries) + assert any("AS session_outcomes" in query for query in queries) + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + @pytest.mark.asyncio async def test_ui_view_spend_logs_group_by_session_jumps_to_page_without_cursor(client, monkeypatch): """page > 1 with no session_cursor (the UI's last-page jump) serves the sessions that page starts at.""" @@ -7518,6 +8269,15 @@ async def test_ui_view_spend_logs_group_by_session_page_past_count_cap_is_empty( assert past_cap_data["has_more"] is False assert past_cap_data["total"] == cap assert mock_prisma.db.query_raw.await_count == 1, "only the bounded count query runs past the capped window" + + beyond_cap = client.get( + "/spend/logs/ui", + params={**params, "page": cap // 25 + 2}, + headers={"Authorization": "Bearer sk-test"}, + ) + assert beyond_cap.status_code == 200, beyond_cap.text + assert beyond_cap.json()["data"] == [] + assert beyond_cap.json()["total"] == cap finally: app.dependency_overrides.pop(ps.user_api_key_auth, None) @@ -8268,3 +9028,50 @@ async def test_global_spend_refresh_runs_the_refresh_on_the_writer_when_a_reader assert [statement for statement in journal.statements if statement[1].startswith("REFRESH")] == [ ("writer", 'REFRESH MATERIALIZED VIEW "MonthlyGlobalSpend";') ] + + +@pytest.mark.asyncio +async def test_ui_view_spend_logs_group_by_session_status_drops_a_session_whose_outcome_changed(client, monkeypatch): + """A row written between the page read and the representatives can flip a session's outcome; like the status + filter, the representatives statement then drops the session instead of returning a row of the other status.""" + flipped = {**_session_representative_row("req-new", "sess-1"), "status": "success"} + + async def mock_query_raw(sql_query, *params): + if "rows_scanned" in sql_query: + return _session_scan_rows([_session_page_row("sess-1", "2026-08-29 08:00:00")], 1) + if "AS session_representatives" in sql_query: + # The page read saw sess-1 failed; its newest row, written since, succeeded. + check = re.search(r"WHERE COALESCE\(status, 'success'\) = \$(\d+)\s*$", sql_query.strip()) + keeps = check is None or params[int(check.group(1)) - 1] == flipped["status"] + return [flipped] if keeps else [] + if "COALESCE(SUM(spend)" in sql_query: + return [] + raise AssertionError(f"unexpected query:\n{sql_query}") + + mock_prisma = MagicMock() + mock_prisma.db = MagicMock() + mock_prisma.db.query_raw = AsyncMock(side_effect=mock_query_raw) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + monkeypatch.setattr( + "litellm.proxy.spend_tracking.spend_management_endpoints._is_admin_view_safe", + lambda user_api_key_dict: True, + ) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user" + ) + try: + start_date, end_date = _default_date_range() + response = client.get( + "/spend/logs/ui", + params={ + "start_date": start_date, + "end_date": end_date, + "group_by_session": "true", + "status_filter": "failure", + }, + headers={"Authorization": "Bearer sk-test"}, + ) + assert response.status_code == 200, response.text + assert [row.get("status") for row in response.json()["data"]] == [] + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) diff --git a/tests/unit/proxy/spend_tracking/test_spend_query_optimization.py b/tests/unit/proxy/spend_tracking/test_spend_query_optimization.py index 93fae093340..49c25c943f9 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_query_optimization.py +++ b/tests/unit/proxy/spend_tracking/test_spend_query_optimization.py @@ -518,10 +518,7 @@ async def test_spend_logs_ui_group_by_session_paginates_sessions(monkeypatch): footer still claims N (issue #38060). """ from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth - from litellm.proxy.spend_tracking.spend_management_endpoints import ( - SPEND_LOGS_PAGINATION_COUNT_CAP, - ui_view_spend_logs, - ) + from litellm.proxy.spend_tracking.spend_management_endpoints import ui_view_spend_logs group_key = "COALESCE(NULLIF(session_id, ''), request_id), api_key" session_rows = [ @@ -534,11 +531,11 @@ async def test_spend_logs_ui_group_by_session_paginates_sessions(monkeypatch): ] async def mock_query_raw(sql_query, *params): - if "COUNT(*) AS total_count" in sql_query: - return [{"total_count": 60}] + if "rows_scanned" in sql_query: + return [{**row, "settled": True, "rows_scanned": 0, "session_count": 60} for row in session_rows] if "AS session_representatives" in sql_query: return representative_rows - return session_rows + return [] mock_prisma = MagicMock() mock_prisma.db = MagicMock() @@ -569,18 +566,13 @@ async def test_spend_logs_ui_group_by_session_paginates_sessions(monkeypatch): emitted = [call[0] for call in mock_prisma.db.query_raw.call_args_list] page_sql = emitted[0][0] - assert f"GROUP BY {group_key}" in page_sql, f"page must select sessions, not calls. SQL was:\n{page_sql}" + assert "GROUP BY session_key, api_key" in page_sql, f"page must select sessions, not calls. SQL was:\n{page_sql}" assert "OFFSET" not in page_sql, "the startTime page must be keyset-selected, not offset-selected" assert emitted[0][-1] == 51, "the page query fetches page_size + 1 sessions to detect has_more" + assert "COUNT(*) OVER ()" not in page_sql + assert "total_count" not in " ".join(call[0] for call in emitted), "the page query also counts the sessions" - count_call = emitted[1] - count_sql = count_call[0] - assert f"GROUP BY {group_key}" in count_sql, f"grouped total must count sessions. SQL was:\n{count_sql}" - assert "COUNT(*) OVER ()" not in count_sql - assert "LIMIT" in count_sql and "FROM (" in count_sql, "the grouped count must stay bounded" - assert count_call[-1] == SPEND_LOGS_PAGINATION_COUNT_CAP + 1 - - rep_sql = emitted[2][0] + rep_sql = emitted[1][0] assert f"DISTINCT ON ({group_key})" in rep_sql, f"page must return one row per session. SQL was:\n{rep_sql}" assert ( f"ORDER BY {group_key}, (call_type = 'asend_message') DESC, "