diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 84278f4619d..5527d051458 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -3,6 +3,7 @@ import collections import json import os from collections.abc import Mapping, Sequence +from dataclasses import dataclass from datetime import date, datetime, timedelta, timezone from functools import partial from itertools import groupby @@ -50,6 +51,7 @@ from litellm.proxy.auth.authorization_dependencies import ( LogTeamLookupDependency, ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.db.db_span import db_span from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup from litellm.proxy.spend_tracking.spend_capture_rate import ( ProviderBillingCredentialMissing, @@ -88,6 +90,10 @@ _DECODED_JSON: Final = TypeAdapter(object) _SESSION_KEY_EXPR: Final = "COALESCE(NULLIF(session_id, ''), request_id)" _SESSION_GROUP_KEY_SQL: Final = f"{_SESSION_KEY_EXPR}, api_key" +_SESSION_WALK_FIRST_BATCH_ROWS: Final = 1000 +_SESSION_WALK_MAX_ROWS: Final = 20000 +_SESSION_COUNT_FIRST_BATCH_ROWS: Final = 50000 +_SESSION_COUNT_MAX_ROWS: Final = 200000 _MCP_CALL_TYPES_SQL: Final = "('call_mcp_tool', 'list_mcp_tools')" _AGENT_CALL_TYPE_SQL: Final = "'asend_message'" _SESSION_REPRESENTATIVE_ORDER_SQL: Final = ( @@ -184,6 +190,11 @@ class _SpendLogsCountRow(TypedDict): total_count: int +class _SessionCountBatchRow(TypedDict): + rows_read: ReadOnly[int] + sessions: ReadOnly[int] + + class _PgClassRow(TypedDict): relname: str relkind: str @@ -2847,6 +2858,7 @@ 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_conditions: Final = tuple(sql_conditions) 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( @@ -2971,6 +2983,7 @@ async def ui_view_spend_logs( ): return await _ui_session_grouped_spend_logs( prisma_client=prisma_client, + window_conditions=window_conditions, sql_conditions=sql_conditions, sql_params=sql_params, next_param_index=p, @@ -3093,6 +3106,16 @@ def _parse_session_cursor(session_cursor: str | None) -> tuple[str, str, str] | return (last_activity, session_key, api_key) +def _session_rows_sql(session_key: str, api_key: str) -> str: + """Every row that can belong to one session, found by its key alone. The OFFSET 0 keeps the caller's window + and filters out of this scan, so a generic plan cannot combine it with a scan of the whole window.""" + return f"""SELECT * + FROM "LiteLLM_SpendLogs" + WHERE (request_id = {session_key} OR session_id = {session_key}) + AND api_key = {api_key} + OFFSET 0""" + + async def _fetch_session_representatives( prisma_client: "PrismaClient", where_clause: str, @@ -3102,16 +3125,16 @@ async def _fetch_session_representatives( ) -> 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.""" rep_query: Final = f""" - SELECT * FROM ( - SELECT DISTINCT ON ({_SESSION_GROUP_KEY_SQL}) - {_SPEND_LOG_LIST_COLUMNS} - FROM "LiteLLM_SpendLogs" + SELECT rep.* + FROM unnest(${next_param_index}::text[], ${next_param_index + 1}::text[]) AS requested(session_key, key) + CROSS JOIN LATERAL ( + SELECT {_SPEND_LOG_LIST_COLUMNS} + FROM ({_session_rows_sql("requested.session_key", "requested.key")}) AS "LiteLLM_SpendLogs" WHERE {where_clause} - 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 + AND {_SESSION_KEY_EXPR} = requested.session_key + ORDER BY {_SESSION_REPRESENTATIVE_ORDER_SQL} + LIMIT 1 + ) AS rep """ rep_rows: Final[Sequence[dict[str, object]]] = await _query_raw( # mutable-ok: rows are enriched in place prisma_client, @@ -3126,6 +3149,211 @@ async def _fetch_session_representatives( return [rep_by_key[key] for key in session_keys if key in rep_by_key] +class _SessionWalkRow(TypedDict): + session_key: ReadOnly[str] + api_key: ReadOnly[str] + last_activity: ReadOnly[str] + sort_time: ReadOnly[str] + request_id: ReadOnly[str] + is_edge: ReadOnly[bool] + batch_rows: ReadOnly[int] + + +class _SessionKeyRow(TypedDict): + session_key: ReadOnly[str] + api_key: ReadOnly[str] + + +@dataclass(frozen=True, slots=True) +class _SessionWalkBatch: + first_rows: tuple[_SessionWalkRow, ...] + edge: _SessionWalkRow | None + rows_read: int + + +@dataclass(frozen=True, slots=True) +class _SessionHeadWalk: + """Walks the window newest first; the first row it meets of each session is that session's newest row.""" + + prisma_client: "PrismaClient" + window_clause: str + where_clause: str + sql_params: tuple[object, ...] + next_param_index: int + cursor: tuple[str, str, str] | None + page_limit: int + + def _cursor_bounds(self) -> str: + return f"(${self.next_param_index}::timestamp, ${self.next_param_index + 1}, ${self.next_param_index + 2})" + + def _params(self) -> tuple[object, ...]: + return (*self.sql_params, *(self.cursor or ())) + + async def _batch(self, after: _SessionWalkRow | None, size: int) -> _SessionWalkBatch: + params: Final = self._params() + keyset_clause: Final = ( + f'AND ("startTime", {_SESSION_KEY_EXPR}, api_key) < {self._cursor_bounds()}' if self.cursor else "" + ) + below_cursor_clause: Final = f'AND "startTime" <= ${self.next_param_index}::timestamp' if self.cursor else "" + after_clause: Final = ( + f'AND ("startTime", request_id) < (${len(params) + 1}::timestamp, ${len(params) + 2})' if after else "" + ) + after_params: Final = (after["last_activity"], after["request_id"]) if after else () + walk_columns: Final = ( + "session_key, api_key, start_time::text AS last_activity, " + "to_char(start_time, 'YYYY-MM-DD HH24:MI:SS.US') AS sort_time, request_id" + ) + query: Final = f""" + WITH batch AS MATERIALIZED ( + SELECT {_SESSION_KEY_EXPR} AS session_key, + api_key, + "startTime" AS start_time, + request_id, + COALESCE(({self.where_clause}) {keyset_clause}, FALSE) AS in_page + FROM ( + SELECT * + FROM "LiteLLM_SpendLogs" + WHERE {self.window_clause} + {below_cursor_clause} + {after_clause} + ORDER BY "startTime" DESC, request_id DESC + LIMIT ${len(params) + len(after_params) + 1} + ) AS "LiteLLM_SpendLogs" + ) + SELECT {walk_columns}, FALSE AS is_edge, 0 AS batch_rows + FROM ( + SELECT DISTINCT ON (session_key, api_key) * + FROM batch + WHERE in_page + ORDER BY session_key, api_key, start_time DESC, request_id DESC + ) AS first_rows + UNION ALL + SELECT {walk_columns}, TRUE AS is_edge, (SELECT COUNT(*) FROM batch)::int AS batch_rows + FROM (SELECT * FROM batch ORDER BY start_time, request_id LIMIT 1) AS edge + """ + rows: Final[Sequence[_SessionWalkRow]] = await _query_raw( + self.prisma_client, query, *params, *after_params, size + ) + edges: Final = tuple(row for row in rows if row["is_edge"]) + return _SessionWalkBatch( + first_rows=tuple( + sorted( + (row for row in rows if not row["is_edge"]), + key=lambda row: (row["sort_time"], row["request_id"]), + reverse=True, + ) + ), + edge=edges[0] if edges else None, + rows_read=edges[0]["batch_rows"] if edges else 0, + ) + + async def _shown_before_the_cursor(self, candidates: Sequence[_SessionWalkRow]) -> frozenset[tuple[str, str]]: + if not self.cursor or not candidates: + return frozenset() + params: Final = self._params() + query: Final = f""" + SELECT requested.session_key, requested.key AS api_key + FROM unnest(${len(params) + 1}::text[], ${len(params) + 2}::text[]) AS requested(session_key, key) + CROSS JOIN LATERAL ( + SELECT 1 + FROM ({_session_rows_sql("requested.session_key", "requested.key")}) AS "LiteLLM_SpendLogs" + WHERE {self.where_clause} + AND {_SESSION_KEY_EXPR} = requested.session_key + AND ("startTime", requested.session_key, requested.key) >= {self._cursor_bounds()} + LIMIT 1 + ) AS newer_row + """ + shown: Final[Sequence[_SessionKeyRow]] = await _query_raw( + self.prisma_client, + query, + *params, + [row["session_key"] for row in candidates], + [row["api_key"] for row in candidates], + ) + return frozenset((row["session_key"], row["api_key"]) for row in shown) + + async def page(self) -> tuple[_SessionPageRow, ...] | None: + """The first ``page_limit + 1`` sessions in page order, or None when the walk gives up before it can tell.""" + return await self._walk( + after=None, heads=(), seen=frozenset(), rows_read=0, size=_SESSION_WALK_FIRST_BATCH_ROWS + ) + + async def _walk( + self, + after: _SessionWalkRow | None, + heads: tuple[_SessionWalkRow, ...], + seen: frozenset[tuple[str, str]], + rows_read: int, + size: int, + ) -> tuple[_SessionPageRow, ...] | None: + batch: Final = await self._batch(after, size) + first_seen: Final = tuple(row for row in batch.first_rows if (row["session_key"], row["api_key"]) not in seen) + shown_before: Final = await self._shown_before_the_cursor(first_seen) + all_heads: Final = heads + tuple( + row for row in first_seen if (row["session_key"], row["api_key"]) not in shown_before + ) + window_exhausted: Final = batch.edge is None or batch.rows_read < size + page_is_settled: Final = ( + len(all_heads) > self.page_limit + and batch.edge is not None + and batch.edge["sort_time"] < all_heads[self.page_limit]["sort_time"] + ) + total_read: Final = rows_read + batch.rows_read + sessions_missing: Final = self.page_limit + 1 - len(all_heads) + next_size: Final = ( + max(_SESSION_WALK_FIRST_BATCH_ROWS, sessions_missing * total_read * 5 // (4 * len(all_heads))) + if all_heads + else _SESSION_WALK_MAX_ROWS + ) + if not (window_exhausted or page_is_settled): + if total_read + next_size > _SESSION_WALK_MAX_ROWS: + return None + return await self._walk( + after=batch.edge, + heads=all_heads, + seen=seen | frozenset((row["session_key"], row["api_key"]) for row in first_seen), + rows_read=total_read, + size=next_size, + ) + if len(frozenset(row["sort_time"] for row in all_heads)) < len(all_heads): + async with db_span("find_many", "LiteLLM_SpendLogs"): + page_rows: Final[Sequence[_SessionPageRow]] = await _query_raw( + self.prisma_client, + """SELECT session_key, api_key, last_activity + FROM unnest($1::text[], $2::text[], $3::text[]) AS heads(session_key, api_key, last_activity) + ORDER BY last_activity::timestamp DESC, session_key DESC, api_key DESC + LIMIT $4""", + [row["session_key"] for row in all_heads], + [row["api_key"] for row in all_heads], + [row["last_activity"] for row in all_heads], + self.page_limit + 1, + ) + return tuple(page_rows) + in_page_order: Final = sorted( + all_heads, key=lambda row: (row["sort_time"], row["session_key"], row["api_key"]), reverse=True + ) + return tuple( + _SessionPageRow(session_key=row["session_key"], api_key=row["api_key"], last_activity=row["last_activity"]) + for row in in_page_order[: self.page_limit + 1] + ) + + +def _grouped_session_page_sql( + where_clause: str, having_clause: str, direction: str, limit_index: int, offset_clause: str +) -> str: + return 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} + """ + + async def _count_grouped_sessions( prisma_client: "PrismaClient", where_clause: str, @@ -3152,8 +3380,55 @@ async def _count_grouped_sessions( ) +async def _count_sessions_newest_first( + prisma_client: "PrismaClient", + window_clause: str, + where_clause: str, + sql_params: Sequence[object], + next_param_index: int, + size: int = _SESSION_COUNT_FIRST_BATCH_ROWS, +) -> tuple[int, bool] | None: + """Count sessions in the newest ``size`` rows of the window, stopping as soon as the count cap is passed or the + window runs out. None when the sessions are too long to settle within ``_SESSION_COUNT_MAX_ROWS`` rows.""" + count_query: Final = f""" + WITH newest_rows AS MATERIALIZED ( + SELECT {_SESSION_KEY_EXPR} AS session_key, + api_key, + COALESCE(({where_clause}), FALSE) AS in_page + FROM ( + SELECT * + FROM "LiteLLM_SpendLogs" + WHERE {window_clause} + ORDER BY "startTime" DESC, request_id DESC + LIMIT ${next_param_index} + ) AS "LiteLLM_SpendLogs" + ) + SELECT (SELECT COUNT(*) FROM newest_rows)::int AS rows_read, + (SELECT COUNT(*) FROM (SELECT DISTINCT session_key, api_key FROM newest_rows WHERE in_page) AS s)::int + AS sessions + """ + rows: Final[Sequence[_SessionCountBatchRow]] = await _query_raw(prisma_client, count_query, *sql_params, size) + rows_read: Final = rows[0]["rows_read"] if rows else 0 + sessions: Final = rows[0]["sessions"] if rows else 0 + if sessions > SPEND_LOGS_PAGINATION_COUNT_CAP: + return (SPEND_LOGS_PAGINATION_COUNT_CAP, True) + if rows_read < size: + return (sessions, False) + next_size: Final = ( + max(2 * size, (SPEND_LOGS_PAGINATION_COUNT_CAP + 1) * rows_read * 5 // (4 * sessions)) + if sessions + else _SESSION_COUNT_MAX_ROWS + 1 + ) + if next_size > _SESSION_COUNT_MAX_ROWS: + return None + return await _count_sessions_newest_first( + prisma_client, window_clause, where_clause, sql_params, next_param_index, next_size + ) + + async def _ui_session_grouped_spend_logs( prisma_client: "PrismaClient", + window_conditions: Sequence[str], sql_conditions: Sequence[str], sql_params: Sequence[object], next_param_index: int, @@ -3171,7 +3446,11 @@ async def _ui_session_grouped_spend_logs( next ``page_size`` sessions ordered by ``(MAX(startTime), session_key, api_key)``, resumed from the ``session_cursor`` keyset ``'||'`` instead of an OFFSET, so - page depth does not degrade the query plan. A request for ``page > 1`` + page depth does not degrade the query plan. A newest-first page comes from + ``_SessionHeadWalk``, which stops after ``page_size + 1`` sessions instead + of grouping the whole window; ascending and OFFSET pages, and walks that + read ``_SESSION_WALK_MAX_ROWS`` rows without settling the page, group the + window instead. A request for ``page > 1`` without a cursor (the UI jumping straight to the last page, or back to a page it never walked through) falls back to ``OFFSET (page - 1) * page_size``, trimmed to the end of the ``SPEND_LOGS_PAGINATION_COUNT_CAP`` @@ -3182,10 +3461,14 @@ async def _ui_session_grouped_spend_logs( ``next_session_cursor`` / ``has_more`` while ``total`` counts sessions (capped like the flat total). A page that runs out of sessions while still holding some is itself the end of the list, so its ``total`` is - ``offset + len(page)`` and the grouped count query is skipped; a page that + ``offset + len(page)`` and the count is skipped; a page that starts past the end says nothing about the total, so that one is counted. + The count walks the window newest first like the page and stops once it + passes the cap, grouping the whole window only when the sessions are too + long for that walk to settle within ``_SESSION_COUNT_MAX_ROWS`` rows. """ where_clause: Final = " AND ".join(sql_conditions) if sql_conditions else "TRUE" + window_clause: Final = " AND ".join(window_conditions) if window_conditions else "TRUE" cmp_op: Final = "<" if sort_desc else ">" direction: Final = "DESC" if sort_desc else "ASC" @@ -3203,21 +3486,32 @@ 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} - """ + walked_rows: Final = ( + await _SessionHeadWalk( + prisma_client=prisma_client, + window_clause=window_clause, + where_clause=where_clause, + sql_params=tuple(sql_params), + next_param_index=next_param_index, + cursor=cursor, + page_limit=page_limit, + ).page() + if sort_desc and offset == 0 and page_limit > 0 + else None + ) 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 walked_rows + if walked_rows is not None + else await _query_raw( + prisma_client, + _grouped_session_page_sql(where_clause, having_clause, direction, limit_index, offset_clause), + *sql_params, + *cursor_params, + page_limit + 1, + *offset_params, + ) ) has_more: Final = len(page_rows) > page_limit @@ -3230,9 +3524,18 @@ async def _ui_session_grouped_spend_logs( page_starts_inside_the_list: Final = offset == 0 or len(page_rows) > 0 page_ends_the_list: Final = cursor is None and page_limit > 0 and not has_more and page_starts_inside_the_list + counted_newest_first: Final = ( + None + if page_ends_the_list + else await _count_sessions_newest_first( + prisma_client, window_clause, where_clause, sql_params, next_param_index + ) + ) total_records, total_is_capped = ( (offset + len(page_rows), False) if page_ends_the_list + else counted_newest_first + if counted_newest_first is not None else await _count_grouped_sessions(prisma_client, where_clause, sql_params, next_param_index) ) @@ -4730,9 +5033,13 @@ async def _build_ui_spend_logs_response( COALESCE(SUM(prompt_tokens), 0)::bigint AS session_total_prompt_tokens, COALESCE(SUM(completion_tokens), 0)::bigint AS session_total_completion_tokens, COALESCE(SUM(total_tokens), 0)::bigint AS session_total_tokens - FROM "LiteLLM_SpendLogs" - WHERE session_id = ANY($1::text[]) - AND api_key = ANY($2::text[]) + FROM unnest($1::text[]) AS requested(session_key) + CROSS JOIN LATERAL ( + SELECT * FROM "LiteLLM_SpendLogs" + WHERE session_id = requested.session_key + AND api_key = ANY($2::text[]) + OFFSET 0 + ) AS session_rows GROUP BY session_id, api_key ) s LEFT JOIN LATERAL ( diff --git a/tests/integration/spend/test_session_view_page_cost.py b/tests/integration/spend/test_session_view_page_cost.py new file mode 100644 index 00000000000..056ccf5e4e2 --- /dev/null +++ b/tests/integration/spend/test_session_view_page_cost.py @@ -0,0 +1,524 @@ +"""The session-grouped /spend/logs/ui page lists each session once and reads a bounded slice of the window. + +The page stops after page_size + 1 sessions and the session count stops once it +passes the count cap, so a window with fewer sessions than the cap is read at +most once. Tuple reads come from pg_stat, and the ordering, representative and +total assertions pin what the page returns. +""" + +import os +import time +import uuid +from collections.abc import Callable, Iterator, Mapping +from dataclasses import dataclass +from datetime import datetime, timedelta +from typing import Final + +import psycopg +import pytest +from integration._support.client import Gateway, string_value +from integration._support.database import ROWS, read_rows +from pydantic import JsonValue + +from litellm.proxy.spend_tracking.spend_management_endpoints import SPEND_LOGS_PAGINATION_COUNT_CAP + +SEED_ROWS: Final = 40_000 +LONG_SESSION_ROWS: Final = 5_000 +SESSION_LEN: Final = 5 +BULK_ROWS: Final = SEED_ROWS - LONG_SESSION_ROWS +PAGE_SIZE: Final = 50 +MAX_TUPLES_PER_PAGE: Final = SEED_ROWS * 2 +_BASE: Final = datetime(2001, 1, 1) +M1_NEWEST_SECOND: Final = SEED_ROWS + 8 +M2_NEWEST_SECOND: Final = SEED_ROWS + 5 +BLOCKER_NEWEST_SECOND: Final = SEED_ROWS + 3 +WINDOW_START: Final = _BASE.strftime("%Y-%m-%d %H:%M:%S") +WINDOW_END: Final = (_BASE + timedelta(seconds=M1_NEWEST_SECOND)).strftime("%Y-%m-%d %H:%M:%S") +_BULK_SESSIONS: Final = len({n // SESSION_LEN for n in range(1, BULK_ROWS + 1)}) +EXPECTED_SESSIONS: Final = _BULK_SESSIONS + 5 + +_STATS_SETTLE_S: Final = 4.0 +_STATS_STABLE_POLLS: Final = 4 +_STATS_TIMEOUT_S: Final = 30.0 +_QUIET_SETTLE_S: Final = 12.0 +_QUIET_STABLE_POLLS: Final = 5 +_QUIET_TIMEOUT_S: Final = 60.0 + + +def _at(seconds: int) -> str: + return (_BASE + timedelta(seconds=seconds)).strftime("%Y-%m-%d %H:%M:%S") + + +KEY_ALIAS: Final = "intg-sesswin-alias" +_SMALL_BASE: Final = datetime(2002, 6, 1) +SMALL_WINDOW_START: Final = _SMALL_BASE.strftime("%Y-%m-%d %H:%M:%S") +SMALL_WINDOW_END: Final = (_SMALL_BASE + timedelta(seconds=100)).strftime("%Y-%m-%d %H:%M:%S") + + +@dataclass(frozen=True, slots=True) +class _SeedRow: + request_id: str + session_id: str | None + start_time: str + api_key: str = "K" + call_type: str = "acompletion" + end_time: str | None = None + + +def _hand_placed_rows(marker: str) -> tuple[_SeedRow, ...]: + prefix: Final = f"intg-sesswin-{marker}-" + blocker_request_id: Final = f"{prefix}x" + return ( + _SeedRow(blocker_request_id, None, _at(SEED_ROWS + 1)), + _SeedRow(f"{prefix}x1", blocker_request_id, _at(SEED_ROWS + 2)), + _SeedRow(f"{prefix}x2", blocker_request_id, _at(BLOCKER_NEWEST_SECOND)), + _SeedRow(f"{prefix}m2a", f"{marker}-m2", _at(SEED_ROWS + 4)), + _SeedRow(f"{prefix}m2b", f"{marker}-m2", _at(M2_NEWEST_SECOND)), + _SeedRow(f"{prefix}m1a", f"{marker}-m1", _at(SEED_ROWS + 6)), + _SeedRow(f"{prefix}m1b", f"{marker}-m1", _at(SEED_ROWS + 7)), + _SeedRow(f"{prefix}m1c", f"{marker}-m1", _at(M1_NEWEST_SECOND)), + _SeedRow(f"intg-straddle-{marker}-in", f"{marker}-straddle", _at(1000)), + _SeedRow(f"intg-straddle-{marker}-out", f"{marker}-straddle", _at(SEED_ROWS + 2_000_000)), + ) + + +def _seed(marker: str) -> None: + with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as connection: + connection.execute( + 'INSERT INTO "LiteLLM_SpendLogs" (' + 'request_id, call_type, api_key, session_id, "startTime", "endTime", model, metadata, messages, response' + ") " + "SELECT %s || n, 'acompletion', 'K', " + "CASE WHEN n > %s THEN %s || '-long' ELSE %s || '-s-' || (n / %s) END, " + "ts, ts, 'e2e-sesswin', jsonb_build_object('user_api_key_alias', %s::text), '{}'::jsonb, '{}'::jsonb " + "FROM generate_series(1, %s) AS g(n) " + "CROSS JOIN LATERAL (SELECT %s::timestamp + (n || ' seconds')::interval AS ts) AS t", + (f"intg-sesswin-{marker}-", BULK_ROWS, marker, marker, SESSION_LEN, KEY_ALIAS, SEED_ROWS, WINDOW_START), + ) + _insert_rows(_hand_placed_rows(marker)) + + +def _insert_rows(rows: tuple[_SeedRow, ...]) -> None: + with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as connection: + connection.cursor().executemany( + 'INSERT INTO "LiteLLM_SpendLogs" (' + 'request_id, call_type, api_key, session_id, "startTime", "endTime", model, metadata, messages, response' + ") VALUES (%s, %s, %s, %s, %s, %s, 'm', jsonb_build_object('user_api_key_alias', %s::text), '{}', '{}')", + [ + ( + row.request_id, + row.call_type, + row.api_key, + row.session_id, + row.start_time, + row.end_time or row.start_time, + KEY_ALIAS, + ) + for row in rows + ], + ) + connection.execute('VACUUM ANALYZE "LiteLLM_SpendLogs"') + + +def _settle_autovacuum() -> None: + with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as connection: + connection.execute('VACUUM ANALYZE "LiteLLM_SpendLogs"') + connection.execute('ALTER TABLE "LiteLLM_SpendLogs" SET (autovacuum_enabled = false)') + + +def _cleanup(marker: str) -> None: + with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as connection: + try: + connection.execute( + 'DELETE FROM "LiteLLM_SpendLogs" WHERE request_id LIKE %s OR request_id LIKE %s', + (f"intg-sesswin-{marker}-%", f"intg-straddle-{marker}-%"), + ) + connection.execute('VACUUM "LiteLLM_SpendLogs"') + finally: + connection.execute('ALTER TABLE "LiteLLM_SpendLogs" RESET (autovacuum_enabled)') + _wait_for_stats_quiet() + + +def _wait_for_stats_quiet() -> None: + _wait_until_stable( + _tuples_read, + settle_s=_STATS_SETTLE_S, + stable_polls=_STATS_STABLE_POLLS, + timeout_s=_STATS_TIMEOUT_S, + ) + + +def _tuples_read() -> int: + """Heap tuples Postgres reports read on LiteLLM_SpendLogs: seq_tup_read + plus idx_tup_fetch (heap rows fetched through index scans). Index entries + traversed do not count, so bulk INSERTs never inflate the reading.""" + with psycopg.connect(os.environ["DATABASE_URL"]) as connection: + row = connection.execute( + "SELECT t.seq_tup_read, t.idx_tup_fetch FROM pg_stat_user_tables t WHERE t.relname = 'LiteLLM_SpendLogs'", + ).fetchone() + assert row is not None + return int(row[0] or 0) + int(row[1] or 0) + + +def _wait_until_stable(read: Callable[[], int], *, settle_s: float, stable_polls: int, timeout_s: float) -> int: + """pg_stat readings once the collector has gone quiet. Each pooled backend + flushes its stats independently, so the counter climbs in delayed jumps: + wait at least ``settle_s`` seconds, then take the reading once it is + unchanged for ``stable_polls`` consecutive polls.""" + deadline: Final = time.monotonic() + timeout_s + min_wait: Final = time.monotonic() + settle_s + stable: int = 0 + last: int = read() + while time.monotonic() < deadline: + current: int = read() + stable = stable + 1 if current == last else 0 + last = current + if stable >= stable_polls and time.monotonic() > min_wait: + return current + time.sleep(0.5) + raise AssertionError("pg_stat counters never settled") + + +def _tuples_read_quiet() -> int: + return _wait_until_stable( + _tuples_read, + settle_s=_QUIET_SETTLE_S, + stable_polls=_QUIET_STABLE_POLLS, + timeout_s=_QUIET_TIMEOUT_S, + ) + + +def _data_rows(page: Mapping[str, JsonValue]) -> list[dict[str, JsonValue]]: + return ROWS.validate_python(page["data"]) + + +def _session_page(gateway: Gateway, **extra: str) -> dict[str, JsonValue]: + return gateway.get( + "/spend/logs/ui", + params={ + "group_by_session": "true", + "start_date": WINDOW_START, + "end_date": WINDOW_END, + "sort_by": "startTime", + "sort_order": "desc", + "page_size": str(PAGE_SIZE), + **extra, + }, + ) + + +@pytest.fixture +def seeded(gateway: Gateway) -> Iterator[str]: + marker: Final = uuid.uuid4().hex + with gateway.scenario() as scenario: + scenario.cleanups.callback(_cleanup, marker) + _settle_autovacuum() + _seed(marker) + _wait_for_stats_quiet() + yield marker + + +def test_session_view_first_page_reads_a_bounded_slice(gateway: Gateway, seeded: str) -> None: + prefix: Final = f"intg-sesswin-{seeded}-" + before: Final = _tuples_read_quiet() + page: Final = _session_page(gateway, page="1") + tuples_read: Final = _tuples_read_quiet() - before + + assert page["has_more"] is True + assert page["total"] == EXPECTED_SESSIONS + assert EXPECTED_SESSIONS < SPEND_LOGS_PAGINATION_COUNT_CAP + assert page["total_is_capped"] is False + data: Final = _data_rows(page) + assert len(data) == PAGE_SIZE + assert data[0]["request_id"] == f"{prefix}m1c", "multi-row session must be represented by its newest row" + assert data[1]["request_id"] == f"{prefix}m2b" + assert data[2]["request_id"] == f"{prefix}x2", ( + "a NULL session_id row and the rows adopting its request_id as session_id form one session" + ) + blocker_keys: Final = {f"{prefix}x", f"{prefix}x1", f"{prefix}x2"} + assert sum(1 for row in data if row["request_id"] in blocker_keys) == 1 + assert data[3]["request_id"] == f"{prefix}{SEED_ROWS}", "the long session must be represented by its newest row" + expected_fivers: Final = [BULK_ROWS, BULK_ROWS - 1, BULK_ROWS - SESSION_LEN - 1] + assert [row["request_id"] for row in data[4:7]] == [f"{prefix}{n}" for n in expected_fivers], ( + "each 5-row session must be represented by its newest row, newest session first" + ) + start_times: Final = [string_value(row["startTime"]) for row in data] + assert start_times == sorted(start_times, reverse=True) + assert tuples_read < MAX_TUPLES_PER_PAGE, ( + f"first page read {tuples_read} tuples for a window of {SEED_ROWS} rows and page_size {PAGE_SIZE}; " + "the page must stop after page_size + 1 sessions, leaving the capped count as the only full pass" + ) + + +def test_session_view_cursor_page_reads_a_bounded_slice(gateway: Gateway, seeded: str) -> None: + first: Final = _session_page(gateway, page="1") + cursor: Final = first["next_session_cursor"] + assert isinstance(cursor, str) + + before: Final = _tuples_read_quiet() + second: Final = _session_page(gateway, page="2", session_cursor=cursor) + tuples_read: Final = _tuples_read_quiet() - before + + second_data: Final = _data_rows(second) + assert len(second_data) == PAGE_SIZE + oldest_first: Final = string_value(_data_rows(first)[-1]["startTime"]) + assert all(string_value(row["startTime"]) < oldest_first for row in second_data) + assert tuples_read < MAX_TUPLES_PER_PAGE, ( + f"cursor page read {tuples_read} tuples for a window of {SEED_ROWS} rows and page_size {PAGE_SIZE}; " + "the page must stop after page_size + 1 sessions, leaving the capped count as the only full pass" + ) + + +def test_session_view_repeated_page_stays_bounded_under_the_generic_plan(gateway: Gateway, seeded: str) -> None: + """The proxy runs these as prepared statements and Postgres switches to a generic plan after five executions, so a plan that only looks good with literal values shows up here.""" + for _ in range(19): + _session_page(gateway, page="1") + before: Final = _tuples_read_quiet() + _session_page(gateway, page="1") + tuples_read: Final = _tuples_read_quiet() - before + assert tuples_read < MAX_TUPLES_PER_PAGE, ( + f"the 20th page read {tuples_read} tuples for a window of {SEED_ROWS} rows and page_size {PAGE_SIZE}" + ) + + +def test_session_view_uses_newest_row_inside_the_window(gateway: Gateway, seeded: str) -> None: + page: Final = _session_page(gateway, page="1", session_id=f"{seeded}-straddle") + assert page["total"] == 1 + data: Final = _data_rows(page) + assert [row["request_id"] for row in data] == [f"intg-straddle-{seeded}-in"] + + +def test_session_view_filtered_first_page_stays_bounded(gateway: Gateway, seeded: str) -> None: + """A filter on a column no index covers is where the planner once drove every per-row probe off the wrong index.""" + before: Final = _tuples_read_quiet() + page: Final = _session_page(gateway, page="1", key_alias=KEY_ALIAS) + tuples_read: Final = _tuples_read_quiet() - before + + assert page["total"] == EXPECTED_SESSIONS + assert _data_rows(page)[0]["request_id"] == f"intg-sesswin-{seeded}-m1c" + assert tuples_read < MAX_TUPLES_PER_PAGE, ( + f"filtered first page read {tuples_read} tuples for a window of {SEED_ROWS} rows and page_size {PAGE_SIZE}" + ) + + +MANY_SESSIONS_ROWS: Final = 150_000 +MANY_SESSIONS_LEN: Final = 2 +_MANY_BASE: Final = datetime(2003, 3, 1) +MANY_WINDOW_START: Final = _MANY_BASE.strftime("%Y-%m-%d %H:%M:%S") +MANY_WINDOW_END: Final = (_MANY_BASE + timedelta(seconds=MANY_SESSIONS_ROWS + 1)).strftime("%Y-%m-%d %H:%M:%S") + + +@pytest.fixture +def many_sessions(gateway: Gateway) -> Iterator[str]: + marker: Final = uuid.uuid4().hex + with gateway.scenario() as scenario: + scenario.cleanups.callback(_cleanup, marker) + _settle_autovacuum() + with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as connection: + connection.execute( + 'INSERT INTO "LiteLLM_SpendLogs" (' + 'request_id, call_type, api_key, session_id, "startTime", "endTime", model, metadata, messages, ' + "response) " + "SELECT %s || n, 'acompletion', 'K', %s || '-s-' || (n / %s), ts, ts, 'e2e-sesswin', '{}'::jsonb, " + "'{}'::jsonb, '{}'::jsonb " + "FROM generate_series(1, %s) AS g(n) " + "CROSS JOIN LATERAL (SELECT %s::timestamp + (n || ' seconds')::interval AS ts) AS t", + (f"intg-sesswin-{marker}-", marker, MANY_SESSIONS_LEN, MANY_SESSIONS_ROWS, MANY_WINDOW_START), + ) + connection.execute('VACUUM ANALYZE "LiteLLM_SpendLogs"') + _wait_for_stats_quiet() + yield marker + + +def test_session_view_count_stops_at_the_cap_on_a_window_of_many_sessions(gateway: Gateway, many_sessions: str) -> None: + """A window holding far more sessions than the count cap is counted from its newest rows, not grouped whole.""" + assert MANY_SESSIONS_ROWS // MANY_SESSIONS_LEN > SPEND_LOGS_PAGINATION_COUNT_CAP + before: Final = _tuples_read_quiet() + page: Final = gateway.get( + "/spend/logs/ui", + params={ + "group_by_session": "true", + "start_date": MANY_WINDOW_START, + "end_date": MANY_WINDOW_END, + "sort_by": "startTime", + "sort_order": "desc", + "page_size": str(PAGE_SIZE), + "page": "1", + }, + ) + tuples_read: Final = _tuples_read_quiet() - before + + assert page["total"] == SPEND_LOGS_PAGINATION_COUNT_CAP + assert page["total_is_capped"] is True + assert _data_rows(page)[0]["request_id"] == f"intg-sesswin-{many_sessions}-{MANY_SESSIONS_ROWS}" + assert tuples_read < MANY_SESSIONS_ROWS // 2, ( + f"the page read {tuples_read} tuples of a {MANY_SESSIONS_ROWS} row window; the page and the capped count " + "must both stop long before the end of a window with more sessions than the cap" + ) + + +def _small_rows(marker: str) -> tuple[_SeedRow, ...]: + def at(seconds: int) -> str: + return (_SMALL_BASE + timedelta(seconds=seconds)).strftime("%Y-%m-%d %H:%M:%S") + + return ( + _SeedRow(f"{marker}-adopter", f"{marker}-root", at(10)), + _SeedRow(f"{marker}-root", None, at(20)), + _SeedRow(f"{marker}-adopter2", f"{marker}-root2", at(30)), + _SeedRow(f"{marker}-root2", "", at(40)), + _SeedRow(f"{marker}-shared-a", f"{marker}-shared", at(50), api_key=f"{marker}-key-a"), + _SeedRow(f"{marker}-shared-b", f"{marker}-shared", at(60), api_key=f"{marker}-key-b"), + _SeedRow(f"{marker}-tie-1", f"{marker}-tie", at(70)), + _SeedRow(f"{marker}-tie-2", f"{marker}-tie", at(70)), + _SeedRow(f"{marker}-mcp-old", f"{marker}-mcp", at(45)), + _SeedRow(f"{marker}-mcp-llm", f"{marker}-mcp", at(80)), + _SeedRow(f"{marker}-mcp-tool", f"{marker}-mcp", at(90), call_type="call_mcp_tool"), + _SeedRow(f"{marker}-mcp-agent-final", f"{marker}-mcp", at(50), call_type="asend_message", end_time=at(85)), + _SeedRow(f"{marker}-mcp-agent-newer", f"{marker}-mcp", at(70), call_type="asend_message", end_time=at(75)), + _SeedRow(f"{marker}-shared", f"{marker}-other", at(95), api_key=f"{marker}-key-b"), + ) + + +@pytest.fixture +def small_seed(gateway: Gateway) -> Iterator[str]: + marker: Final = f"intg-sesssmall-{uuid.uuid4().hex}" + with gateway.scenario() as scenario: + scenario.cleanups.callback(_delete_small_rows, marker) + _insert_rows(_small_rows(marker)) + yield marker + + +def _delete_small_rows(marker: str) -> None: + with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as connection: + connection.execute('DELETE FROM "LiteLLM_SpendLogs" WHERE request_id LIKE %s', (f"{marker}-%",)) + + +def _small_page(gateway: Gateway, **extra: str) -> dict[str, JsonValue]: + return gateway.get( + "/spend/logs/ui", + params={ + "group_by_session": "true", + "start_date": SMALL_WINDOW_START, + "end_date": SMALL_WINDOW_END, + "sort_by": "startTime", + **extra, + }, + ) + + +def _sessions(page: Mapping[str, JsonValue]) -> list[tuple[str, str]]: + return [(_tie_as_one(string_value(row["request_id"])), string_value(row["api_key"])) for row in _data_rows(page)] + + +def _tie_as_one(request_id: str) -> str: + return request_id.removesuffix("-1").removesuffix("-2") if "-tie-" in request_id else request_id + + +def _newest_first(marker: str) -> list[tuple[str, str]]: + return [ + (f"{marker}-shared", f"{marker}-key-b"), + (f"{marker}-mcp-agent-final", "K"), + (f"{marker}-tie", "K"), + (f"{marker}-shared-b", f"{marker}-key-b"), + (f"{marker}-shared-a", f"{marker}-key-a"), + (f"{marker}-root2", "K"), + (f"{marker}-root", "K"), + ] + + +def test_session_view_lists_each_session_once(gateway: Gateway, small_seed: str) -> None: + page: Final = _small_page(gateway, sort_order="desc", page_size="50", page="1") + + assert _sessions(page) == _newest_first(small_seed), ( + "each session once, preferring the final agent outcome, then the newest non-MCP row: " + "a row without a session_id newer than the rows " + "adopting its request_id joins them, two rows at one startTime are one session, one session_id under two " + "api_keys is two sessions, and a row whose request_id equals another session's key stays in its own session" + ) + assert page["total"] == 7 + assert page["has_more"] is False + + +def test_session_view_cursor_walk_counts_sessions_per_api_key(gateway: Gateway, small_seed: str) -> None: + first: Final = _small_page(gateway, sort_order="desc", page_size="2", page="1") + cursor_pages: Final = [first] + while cursor_pages[-1]["has_more"] is True: + cursor_pages.append( + _small_page( + gateway, + sort_order="desc", + page_size="2", + page=str(len(cursor_pages) + 1), + session_cursor=string_value(cursor_pages[-1]["next_session_cursor"]), + ) + ) + walked: Final = [session for page in cursor_pages for session in _sessions(page)] + + assert first["total"] == 7, "the capped count must count one session_id under two api_keys as two sessions" + assert walked == _newest_first(small_seed), ( + "a session whose older rows fall below the cursor was already shown and must not come back on a later page" + ) + assert [page["has_more"] for page in cursor_pages] == [True, True, True, False] + + +def test_session_view_ascending_and_offset_pages_match_the_newest_first_order( + gateway: Gateway, small_seed: str +) -> None: + ascending: Final = _small_page(gateway, sort_order="asc", page_size="50", page="1") + offset_page: Final = _small_page(gateway, sort_order="desc", page_size="2", page="2") + + assert _sessions(ascending) == _newest_first(small_seed)[::-1] + assert ascending["total"] == 7 + assert _sessions(offset_page) == _newest_first(small_seed)[2:4] + assert offset_page["has_more"] is True + + +def test_session_view_filtered_page_only_lists_matching_sessions(gateway: Gateway, small_seed: str) -> None: + page: Final = _small_page(gateway, sort_order="desc", page_size="50", page="1", api_key=f"{small_seed}-key-b") + + assert _sessions(page) == [ + (f"{small_seed}-shared", f"{small_seed}-key-b"), + (f"{small_seed}-shared-b", f"{small_seed}-key-b"), + ] + assert page["total"] == 2 + + +@pytest.mark.parametrize("identities", ((("a", "K"), ("B", "K")), (("shared", "a"), ("shared", "B")))) +def test_session_view_cursor_pages_preserve_database_order_for_tied_sessions( + gateway: Gateway, identities: tuple[tuple[str, str], tuple[str, str]] +) -> None: + marker: Final = f"intg-sesscollation-{uuid.uuid4().hex}" + rows: Final = tuple( + _SeedRow(f"{marker}-{index}", f"{marker}-{session}", SMALL_WINDOW_START, api_key=key) + for index, (session, key) in enumerate(identities) + ) + + with gateway.scenario() as scenario: + scenario.cleanups.callback(_delete_small_rows, marker) + _insert_rows(rows) + expected: Final = tuple( + (string_value(row["request_id"]), string_value(row["api_key"])) + for row in read_rows( + 'SELECT request_id, api_key FROM "LiteLLM_SpendLogs" WHERE request_id LIKE %s ' + 'ORDER BY "startTime" DESC, session_id DESC, api_key DESC', + (f"{marker}-%",), + ) + ) + + first: Final = _small_page(gateway, session_id=marker, page_size="1", page="1") + second: Final = _small_page( + gateway, + session_id=marker, + page_size="1", + page="2", + session_cursor=string_value(first["next_session_cursor"]), + ) + walked: Final = tuple( + (string_value(row["request_id"]), string_value(row["api_key"])) + for row in _data_rows(first) + _data_rows(second) + ) + + assert walked == expected, "cursor pages must match database ordering without omissions or duplicates" + assert first["total"] == second["total"] == len(rows) + assert first["has_more"] is True + assert second["has_more"] is False 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..b83640a8cb1 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py @@ -7173,16 +7173,46 @@ def test_parse_session_cursor(): ) -SESSION_GROUP_KEY_SQL = "COALESCE(NULLIF(session_id, ''), request_id), api_key" +SESSION_KEY_EXPR = "COALESCE(NULLIF(session_id, ''), request_id)" +SESSION_GROUP_KEY_SQL = f"{SESSION_KEY_EXPR}, api_key" + + +def _session_walk_batch(page_rows, rows_read): + """Walk-batch rows for ``page_rows`` (newest first): one first row per session plus the batch's oldest row.""" + first_rows = [ + { + "session_key": row["session_key"], + "api_key": row["api_key"], + "last_activity": row["last_activity"], + "sort_time": row["last_activity"], + "request_id": f"req-{row['session_key']}", + "is_edge": False, + "batch_rows": 0, + } + for row in page_rows + ] + edge = [{**first_rows[-1], "is_edge": True, "batch_rows": rows_read}] if first_rows else [] + return first_rows + edge + + +def _session_count_walk(sessions): + """Count-walk result for a window that ran out of rows after holding ``sessions`` sessions.""" + return [{"rows_read": 0, "sessions": sessions}] 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 "AS newer_row" in sql_query: + return [] + if "WITH batch AS MATERIALIZED" in sql_query: + return _session_walk_batch(session_page_rows, 0) + if "newest_rows" in sql_query: + return _session_count_walk(session_total) if "COUNT(*) AS total_count" in sql_query: return [{"total_count": session_total}] - if "DISTINCT ON" in sql_query: + if "AS rep" in sql_query or "DISTINCT ON" in sql_query: return representative_rows if "COALESCE(SUM(spend)" in sql_query: return [ @@ -7232,9 +7262,15 @@ 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 "AS newer_row" in sql_query: + return [] + if "WITH batch AS MATERIALIZED" in sql_query: + return _session_walk_batch([_session_page_row(key, activity) for key, activity in sessions], 0) + if "newest_rows" in sql_query: + return _session_count_walk(len(sessions) if counted_total is None else counted_total) 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: + if "AS rep" in sql_query or "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: return [] @@ -7294,32 +7330,13 @@ async def test_ui_view_spend_logs_group_by_session_first_page(client, monkeypatc assert request_ids == ["req-1", "req-solo"], "representatives must follow the page order" assert data["data"][0]["session_total_count"] == 3 assert data["data"][1]["session_total_count"] == 1 - - emitted = [call.args for call in mock_prisma.db.query_raw.await_args_list] - page_query_sql = emitted[0][0] - assert f"GROUP BY {SESSION_GROUP_KEY_SQL}" in page_query_sql - assert "OFFSET" not in page_query_sql - assert "HAVING" 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] - 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"] 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 cursor continues after the previously displayed session.""" 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) @@ -7349,21 +7366,214 @@ async def test_ui_view_spend_logs_group_by_session_cursor_page(client, monkeypat assert data["has_more"] is False assert data["next_session_cursor"] is None assert [row["request_id"] for row in data["data"]] == ["req-2"] - - 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 "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] == ( - "2026-08-29 09:00:00", - "req-solo", - "hashed-key", - ), "cursor params must line up with (MAX(startTime), session key, api_key)" finally: app.dependency_overrides.pop(ps.user_api_key_auth, None) +def _session_walk_mock_prisma(walked_rows, window_rows_per_batch, grouped_rows): + """Mock prisma whose walk batches keep returning ``walked_rows`` and report ``window_rows_per_batch`` rows read.""" + + async def mock_query_raw(sql_query, *params): + if "WITH batch AS MATERIALIZED" in sql_query: + return _session_walk_batch(walked_rows, min(window_rows_per_batch, params[-1])) + if "newest_rows" in sql_query: + return _session_count_walk(3) + if "COUNT(*) AS total_count" in sql_query: + return [{"total_count": 3}] + if "AS rep" 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: + return [] + return grouped_rows + + mock_prisma = MagicMock() + mock_prisma.db = MagicMock() + mock_prisma.db.query_raw = AsyncMock(side_effect=mock_query_raw) + return mock_prisma + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("window_rows_per_batch", "expected_request_ids", "expected_has_more"), + [ + pytest.param(10**9, ["req-grouped-1", "req-grouped-2"], True, id="walk-that-cannot-settle-groups-the-window"), + pytest.param(10, ["req-walked-1"], False, id="walk-that-read-the-whole-window-is-trusted"), + ], +) +async def test_ui_view_spend_logs_group_by_session_falls_back_when_the_walk_cannot_settle( + client, monkeypatch, window_rows_per_batch, expected_request_ids, expected_has_more +): + """A walk that keeps finding too few sessions serves the grouped page; one that ran out of rows is exact.""" + mock_prisma = _session_walk_mock_prisma( + walked_rows=[_session_page_row("walked-1", "2026-08-29 10:00:00")], + window_rows_per_batch=window_rows_per_batch, + grouped_rows=[ + _session_page_row("grouped-1", "2026-08-29 10:00:00"), + _session_page_row("grouped-2", "2026-08-29 09:00:00"), + _session_page_row("grouped-3", "2026-08-29 08:00:00"), + ], + ) + 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 + data = response.json() + assert [row["request_id"] for row in data["data"]] == expected_request_ids + assert data["has_more"] is expected_has_more + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +def _walk_row(session_key, sort_time, is_edge=False, batch_rows=0): + return { + "session_key": session_key, + "api_key": "hashed-key", + "last_activity": sort_time, + "sort_time": sort_time, + "request_id": f"req-{session_key}", + "is_edge": is_edge, + "batch_rows": batch_rows, + } + + +@pytest.mark.asyncio +async def test_ui_view_spend_logs_group_by_session_walk_finishes_a_tied_timestamp_before_settling(client, monkeypatch): + """Sessions tied on their newest startTime order by session key, so the walk reads past the tie before settling.""" + tied = "2026-08-29 10:00:00" + + async def mock_query_raw(sql_query, *params): + if "AS heads(session_key, api_key, last_activity)" in sql_query: + return [_session_page_row(key, tied) for key in ("Z", "B", "a") if key in params[0]][: params[-1]] + if "WITH batch AS MATERIALIZED" in sql_query: + if '("startTime", request_id) <' not in sql_query: + return [ + _walk_row("B", tied), + _walk_row("a", tied), + _walk_row("a", tied, is_edge=True, batch_rows=params[-1]), + ] + return [ + {**_walk_row("Z", tied), "request_id": "req-0"}, + _walk_row("Z", "2026-08-29 09:00:00", is_edge=True, batch_rows=2), + ] + if "newest_rows" in sql_query: + return _session_count_walk(3) + if "COUNT(*) AS total_count" in sql_query: + return [{"total_count": 3}] + if "AS rep" in sql_query: + return [_session_representative_row(f"req-{session_key}", session_key) for session_key in params[-2]] + 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 [row["request_id"] for row in data["data"]] == ["req-Z"] + assert data["next_session_cursor"] == f"{tied}|hashed-key|Z" + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +def _session_window(session_count, rows_per_session, newer_filtered_out_rows=0): + """Session keys of a window's rows, newest first, each session's rows back to back; None marks a newer row the + request's filters exclude.""" + return [None] * newer_filtered_out_rows + [ + f"sess-{index}" for index in range(session_count) for _ in range(rows_per_session) + ] + + +def _counting_prisma(window): + """Mock prisma answering the count walk and the grouped count from ``window``, recording the rows each walk read.""" + rows_read_per_walk = [] + + async def mock_query_raw(sql_query, *params): + if "newest_rows" in sql_query: + newest = window[: params[-1]] + rows_read_per_walk.append(len(newest)) + return [{"rows_read": len(newest), "sessions": len(set(newest) - {None})}] + if "COUNT(*) AS total_count" in sql_query: + return [{"total_count": min(len(set(window) - {None}), params[-1])}] + raise AssertionError(f"unexpected query: {sql_query}") + + mock_prisma = MagicMock() + mock_prisma.db = MagicMock() + mock_prisma.db.query_raw = AsyncMock(side_effect=mock_query_raw) + return mock_prisma, rows_read_per_walk + + +_COUNT_CAP = spend_management_endpoints.SPEND_LOGS_PAGINATION_COUNT_CAP + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("window", "expected_total", "walk_settles", "walk_reads_every_row"), + [ + pytest.param(_session_window(30000, 5), (_COUNT_CAP, True), True, False, id="many-sessions-stop-at-the-cap"), + pytest.param( + _session_window(_COUNT_CAP + 2000, 5), (_COUNT_CAP, True), True, True, id="cap-passed-after-the-first-batch" + ), + pytest.param(_session_window(_COUNT_CAP, 5), (_COUNT_CAP, False), True, True, id="exactly-the-cap-is-exact"), + pytest.param(_session_window(1000, 5), (1000, False), True, True, id="small-window-is-counted-exactly"), + pytest.param(_session_window(20, 20000), (20, False), False, False, id="long-sessions-fall-back-to-grouping"), + pytest.param( + _session_window(10, 5, newer_filtered_out_rows=400000), + (10, False), + False, + False, + id="filter-matching-no-recent-rows-falls-back-at-once", + ), + ], +) +async def test_session_count_matches_the_grouped_count_and_stops_early( + window, expected_total, walk_settles, walk_reads_every_row +): + """The newest-first count gives the grouped count's answer, reading only part of a window of many sessions.""" + mock_prisma, rows_read_per_walk = _counting_prisma(window) + + counted = await spend_management_endpoints._count_sessions_newest_first(mock_prisma, "TRUE", "TRUE", (), 1) + total = ( + counted + if counted is not None + else await spend_management_endpoints._count_grouped_sessions(mock_prisma, "TRUE", (), 1) + ) + + assert total == expected_total + assert (counted is not None) is walk_settles + assert (max(rows_read_per_walk) == len(window)) is walk_reads_every_row + assert max(rows_read_per_walk) <= spend_management_endpoints._SESSION_COUNT_MAX_ROWS + if not walk_settles: + assert rows_read_per_walk == [spend_management_endpoints._SESSION_COUNT_FIRST_BATCH_ROWS], ( + "a walk that cannot reach the cap within its row budget must stop after its first batch" + ) + + @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.""" 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..aabce025a10 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_query_optimization.py +++ b/tests/unit/proxy/spend_tracking/test_spend_query_optimization.py @@ -510,20 +510,15 @@ async def test_global_spend_report_team_group_forwards_team_id(monkeypatch): async def test_spend_logs_ui_group_by_session_paginates_sessions(monkeypatch): """ With group_by_session=true, /spend/logs/ui must page and count SESSIONS, - not raw calls: the page is one representative row per session (DISTINCT - ON the session group key, preferring non-MCP calls, newest first), the - bounded count counts groups, and on the default startTime sort the page is + not raw calls: the page is one representative row per session, newest + session first, the count stays bounded, and on the default startTime sort the page is selected by keyset (no OFFSET) so deep pages do not degrade. Otherwise the UI collapses a server page of N calls into fewer visible rows while the 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 = [ {"session_key": f"req-{index}", "api_key": "k", "last_activity": f"2026-02-16 10:{59 - index:02d}:00"} for index in range(51) @@ -534,9 +529,17 @@ async def test_spend_logs_ui_group_by_session_paginates_sessions(monkeypatch): ] async def mock_query_raw(sql_query, *params): + if "WITH batch AS MATERIALIZED" in sql_query: + first_rows = [ + {**row, "sort_time": row["last_activity"], "request_id": row["session_key"], "is_edge": False} + for row in session_rows + ] + return [*first_rows, {**first_rows[-1], "is_edge": True, "batch_rows": 0}] + if "newest_rows" in sql_query: + return [{"rows_read": 0, "sessions": 60}] if "COUNT(*) AS total_count" in sql_query: return [{"total_count": 60}] - if "AS session_representatives" in sql_query: + if "AS rep" in sql_query or "DISTINCT ON" in sql_query: return representative_rows return session_rows @@ -567,28 +570,6 @@ async def test_spend_logs_ui_group_by_session_paginates_sessions(monkeypatch): session_cursor=None, ) - 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 "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" - - 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] - 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, " - "CASE WHEN call_type = 'asend_message' THEN \"endTime\" END DESC NULLS LAST, " - "call_type IN ('call_mcp_tool', 'list_mcp_tools'), \"startTime\" DESC" - ) in rep_sql, "the session representative must prefer the final agent outcome, then the newest non-MCP call" - assert "COUNT(*) OVER ()" not in rep_sql - assert [row["request_id"] for row in response["data"]] == ["req-1", "req-2"] assert response["total"] == 60 assert response["total_is_capped"] is False