From 9c4096d0e3e3cd2cc394e8ecc4ae3c74d1d008ac Mon Sep 17 00:00:00 2001 From: shrey-berri Date: Wed, 28 Oct 2026 09:03:39 -0700 Subject: [PATCH 1/3] fix(spend): read only the newest rows for the session view page (LIT-8513) The session view grouped every row in the date window to show 50 sessions and count them, so each page read the whole window. Read the newest rows that match the filters instead, group those, and keep only sessions that are complete in what was read: a session whose last activity is newer than the oldest row read has every row in the read. The first read is cap + offset + page size + 2 rows, and at most 2 x cap + 2 for a page past the cap, so a huge page number cannot overflow the read's LIMIT. It grows 8x, then 64x, then reads everything, only while the page or the count is not settled; a read that does not settle is discarded, so results are the same as before and the count still stops at 10,000. Ascending order still reads every row in one query. On a busy key, where few long sessions fill the newest rows, a read that the last read's rows per session say cannot hold enough sessions is skipped. An offset page's read that holds no session of the page hands its sessions to the next read, which reads only the rows at or below that read's oldest row, up to the rest of its budget. So a deep page reads each row about once instead of re-reading the newest rows at every step: the last page of a key with 2,000 sessions of 50 rows read about 307-313k rows before and about 195-205k now, against 205k for grouping every row. A descending cursor page reads from the cursor's startTime down and drops sessions that also have a newer row. When the sessions after the cursor do not pass the cap, the window's count is read separately from the newest rows, starting at the first read size that the last read's rows per session say can pass the cap, so an unindexed filter such as key_alias or status does not pay a full scan per step. The status filter keeps or drops whole sessions, so the newer-row probe skips it for sessions already on the page. The representative fetch checks it again on just those sessions, since a row written after the page read can change a session's outcome; as before, such a session is dropped from the page. Its subquery gets the date window as literals: Prisma reuses prepared statements, and from the 6th run Postgres plans a $N window as tiny, which turned the subquery into a ~2 min nested loop. The literal is built only from a UTC datetime; anything else raises. --- .../spend_management_endpoints.py | 275 +++++-- .../test_spend_management_endpoints.py | 758 +++++++++++++++++- .../test_spend_query_optimization.py | 24 +- 3 files changed, 976 insertions(+), 81 deletions(-) diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 84278f4619d..d4b7eb3a7dc 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: list[str] = [] # 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( @@ -2946,7 +2951,13 @@ async def ui_view_spend_logs( 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" + # 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" + ) sql_conditions.append( f"""({_SESSION_GROUP_KEY_SQL}) IN ( SELECT session_key, api_key FROM ( @@ -2978,6 +2989,7 @@ 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, ) # Build the ORDER BY expression. ttft_ms is computed from @@ -3083,6 +3095,45 @@ 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) + + +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 +3150,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 +3177,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 +3185,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 +3194,7 @@ async def _ui_session_grouped_spend_logs( page_size: int, sort_desc: bool, session_cursor: str | None, + status_filtered: bool = False, ) -> Mapping[str, object]: """ One row per session, keyset-paginated by session last activity. @@ -3186,12 +3220,19 @@ 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" 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 +3244,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] = () + carry_params: tuple[list[str], list[str], list[str], str] | tuple[()] = () + found: Sequence[_SessionScanRow] = () + session_count = 0 + rows_read = 0 + 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 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 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 +3411,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..e309c7eb00a 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,659 @@ 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]) + outer, outcomes = read.args[0].split("SELECT DISTINCT ON", 1) + outcomes = outcomes.split("AS session_outcomes", 1)[0] + 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 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 +8182,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 +8941,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, " From 2756af97872a35cddffe3f57b96e24a8404e5bb1 Mon Sep 17 00:00:00 2001 From: shrey-berri Date: Wed, 28 Oct 2026 13:45:08 -0700 Subject: [PATCH 2/3] chore(spend): give reasons for new mutable and rebound locals in session paging --- .../spend_tracking/spend_management_endpoints.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index d4b7eb3a7dc..34f5bf31f05 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -2836,7 +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: list[str] = [] + 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 @@ -3108,6 +3108,8 @@ class _SessionScanRow(TypedDict): # 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: @@ -3261,11 +3263,11 @@ async def _ui_session_grouped_spend_logs( # 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] = () - carry_params: tuple[list[str], list[str], list[str], str] | tuple[()] = () - found: Sequence[_SessionScanRow] = () - session_count = 0 - rows_read = 0 + 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. From 643cdfc4f00a1649c5e07721456608e6f764c972 Mon Sep 17 00:00:00 2001 From: shrey-berri Date: Thu, 29 Oct 2026 09:06:43 -0700 Subject: [PATCH 3/3] fix(spend): read a failure filter's sessions once on the session page With status_filter=failure, the session page inlined the session-outcomes subquery in its read, so each parallel worker re-sorted the whole window for it: on 2M rows, 2.2x slower than base with 3x the blocks and temp. The failure filter now reads those sessions once, as a `status_sessions AS MATERIALIZED` CTE, in the page read and the count. A success filter keeps most sessions and stays inlined, where it is faster. Other routes keep the inlined subquery. Results are unchanged. --- .../spend_management_endpoints.py | 58 ++++++++---- .../test_spend_management_endpoints.py | 91 ++++++++++++++++++- 2 files changed, 127 insertions(+), 22 deletions(-) diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 34f5bf31f05..fe71dc06d57 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -2950,16 +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: - # 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" - ) - 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 @@ -2968,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, @@ -2990,6 +3002,7 @@ async def ui_view_spend_logs( 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 @@ -3197,6 +3210,7 @@ async def _ui_session_grouped_spend_logs( 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. @@ -3231,6 +3245,10 @@ async def _ui_session_grouped_spend_logs( 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) cursor_clause: Final = ( @@ -3298,7 +3316,7 @@ async def _ui_session_grouped_spend_logs( ) carried_keys: str = f"${carry_index}::text[], ${carry_index + 1}::text[]" page_query: str = f""" - WITH scanned AS MATERIALIZED ( + WITH {status_cte} scanned AS MATERIALIZED ( {_newest_session_rows_sql(where_clause, read_rows, not_after)} ), sessions AS ( @@ -3376,7 +3394,7 @@ async def _ui_session_grouped_spend_logs( if _read_falls_short(count_rows, session_count, rows_read, cap + 1): continue count_query: str = f""" - WITH scanned AS MATERIALIZED ({_newest_session_rows_sql(where_clause, count_rows, None)}) + 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 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 e309c7eb00a..ef7b6ac1b3f 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py @@ -8013,8 +8013,9 @@ async def test_ui_view_spend_logs_group_by_session_status_inlines_the_window(cli ) 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]) - outer, outcomes = read.args[0].split("SELECT DISTINCT ON", 1) - outcomes = outcomes.split("AS session_outcomes", 1)[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 @@ -8028,6 +8029,92 @@ async def test_ui_view_spend_logs_group_by_session_status_inlines_the_window(cli 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."""