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