From 44792295f055c18711183e7fb6567ea653bdbd67 Mon Sep 17 00:00:00 2001 From: shrey-berri <333930225+shrey-berri@users.noreply.github.com> Date: Sun, 27 Sep 2026 01:07:33 +0000 Subject: [PATCH] perf(spend): stop the session count once it passes the cap The session view's total came from grouping every row in the date window, even though it is capped at 10,000. It now counts sessions from the newest rows of the window and stops once it has seen more than the cap, so a window with many sessions is no longer read to the end. Windows whose sessions are too long to settle within 200,000 rows still use the grouped count. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_015S4LA6NcP6gp22R1iwH6xw --- .../spend_management_endpoints.py | 72 ++++++++++++++- .../spend/test_session_view_page_cost.py | 65 +++++++++++++- .../test_spend_management_endpoints.py | 89 ++++++++++++++++++- .../test_spend_query_optimization.py | 13 +-- 4 files changed, 219 insertions(+), 20 deletions(-) diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index a71b2f3e04e..bd96d11f39e 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -77,6 +77,8 @@ _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'" _BATCH_CALL_TYPES_SQL: Final = "('acreate_batch', 'create_batch', 'aretrieve_batch', 'retrieve_batch')" @@ -168,6 +170,11 @@ class _SpendLogsCountRow(TypedDict): total_count: int +class _SessionCountBatchRow(TypedDict): + rows_read: ReadOnly[int] + sessions: ReadOnly[int] + + class _PgClassRow(TypedDict): relname: str relkind: str @@ -3320,6 +3327,52 @@ 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], @@ -3343,7 +3396,7 @@ async def _ui_session_grouped_spend_logs( 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 every ``_SESSION_WALK_BATCH_ROWS`` batch without settling the page, group the + 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) * @@ -3355,10 +3408,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" @@ -3379,7 +3436,7 @@ async def _ui_session_grouped_spend_logs( walked_rows: Final = ( await _SessionHeadWalk( prisma_client=prisma_client, - window_clause=" AND ".join(window_conditions) if window_conditions else "TRUE", + window_clause=window_clause, where_clause=where_clause, sql_params=tuple(sql_params), next_param_index=next_param_index, @@ -3414,9 +3471,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) ) diff --git a/tests/integration/spend/test_session_view_page_cost.py b/tests/integration/spend/test_session_view_page_cost.py index b4988466791..9986698ae04 100644 --- a/tests/integration/spend/test_session_view_page_cost.py +++ b/tests/integration/spend/test_session_view_page_cost.py @@ -1,8 +1,9 @@ -"""The session-grouped /spend/logs/ui page lists each session once and reads the window at most once. +"""The session-grouped /spend/logs/ui page lists each session once and reads a bounded slice of the window. -The page itself stops after page_size + 1 sessions; the capped session count is -the one full pass over the window. Tuple reads come from pg_stat, and the -ordering, representative and total assertions pin what the page returns. +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 @@ -291,6 +292,62 @@ def test_session_view_filtered_first_page_stays_bounded(gateway: Gateway, seeded ) +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") diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index 1fbc762bf68..dbad3834459 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -7024,6 +7024,11 @@ def _session_walk_batch(page_rows, rows_read): 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.""" @@ -7032,6 +7037,8 @@ def _session_grouped_mock_prisma(session_page_rows, session_total, representativ 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 "CROSS JOIN LATERAL" in sql_query or "DISTINCT ON" in sql_query: @@ -7088,6 +7095,8 @@ def _session_grouped_paginating_prisma(sessions, counted_total=None): 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 "CROSS JOIN LATERAL" in sql_query or "DISTINCT ON" in sql_query: @@ -7156,9 +7165,6 @@ async def test_ui_view_spend_logs_group_by_session_first_page(client, monkeypatc assert "OFFSET" not in page_query_sql assert "HAVING" not in page_query_sql - count_sql = emitted[1][0] - assert "LIMIT" in count_sql and "FROM (" in count_sql, "the grouped count must stay bounded" - rep_call = emitted[2] assert rep_call[-2] == ["sess-1", "req-solo"] assert rep_call[-1] == ["hashed-key", "hashed-key"] @@ -7219,6 +7225,8 @@ def _session_walk_mock_prisma(walked_rows, window_rows_per_batch, grouped_rows): 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 "CROSS JOIN LATERAL" in sql_query: @@ -7303,6 +7311,8 @@ async def test_ui_view_spend_logs_group_by_session_walk_finishes_a_tied_timestam _walk_row("a", tied, is_edge=True, batch_rows=params[-1]), ] return [_walk_row("c", tied), _walk_row("c", "2026-08-29 09:00:00", is_edge=True, batch_rows=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 "CROSS JOIN LATERAL" in sql_query: @@ -7335,6 +7345,79 @@ async def test_ui_view_spend_logs_group_by_session_walk_finishes_a_tied_timestam 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/test_litellm/proxy/spend_tracking/test_spend_query_optimization.py b/tests/test_litellm/proxy/spend_tracking/test_spend_query_optimization.py index 7b358015ec9..a56ede47629 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_query_optimization.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_query_optimization.py @@ -520,10 +520,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 session_rows = [ {"session_key": f"req-{index}", "api_key": "k", "last_activity": f"2026-02-16 10:{59 - index:02d}:00"} @@ -541,6 +538,8 @@ async def test_spend_logs_ui_group_by_session_paginates_sessions(monkeypatch): 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 "CROSS JOIN LATERAL" in sql_query or "DISTINCT ON" in sql_query: @@ -577,12 +576,6 @@ async def test_spend_logs_ui_group_by_session_paginates_sessions(monkeypatch): page_sql = emitted[0][0] assert "OFFSET" not in page_sql, "the startTime page must be keyset-selected, not offset-selected" - count_call = emitted[1] - count_sql = count_call[0] - 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 emitted[2][-2] == [row["session_key"] for row in session_rows[:50]] assert emitted[2][-1] == ["k"] * 50