scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-06 07:48:06 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 112

Merge 643cdfc4f0 into 06cdee91da

Browse source
This commit is contained in:
shrey 2026-10-29 16:07:58 +00:00 • committed by GitHub
commit 087fcacef3
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 1094 additions and 92 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

317
litellm/proxy/spend_tracking/spend_management_endpoints.py
View file

@ -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 []

845
tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py
View file

@ -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)

24
tests/unit/proxy/spend_tracking/test_spend_query_optimization.py
View file

@ -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, "