From 04d4a19c9d22e2066fdd99cbca2086718e333326 Mon Sep 17 00:00:00 2001 From: shrey-berri <333930225+shrey-berri@users.noreply.github.com> Date: Sat, 26 Sep 2026 17:59:43 +0000 Subject: [PATCH 01/10] perf(schema): index LiteLLM_SpendLogs by session for the session view Add (session_id, api_key, startTime DESC, request_id) so checking whether a spend log row is the newest row of its session is one index probe. The migration builds it CONCURRENTLY so spend log inserts keep flowing during a rolling deploy. --- .../migration.sql | 2 ++ litellm-proxy-extras/litellm_proxy_extras/schema.prisma | 1 + litellm/proxy/schema.prisma | 1 + schema.prisma | 1 + 4 files changed, 5 insertions(+) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260924204413_add_spend_logs_session_head_index/migration.sql diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260924204413_add_spend_logs_session_head_index/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260924204413_add_spend_logs_session_head_index/migration.sql new file mode 100644 index 00000000000..da26297b59c --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260924204413_add_spend_logs_session_head_index/migration.sql @@ -0,0 +1,2 @@ +-- CreateIndex (CONCURRENTLY): keep this a single statement so Prisma Migrate applies it outside a transaction +CREATE INDEX CONCURRENTLY IF NOT EXISTS "LiteLLM_SpendLogs_session_id_api_key_startTime_request_id_idx" ON "LiteLLM_SpendLogs"("session_id", "api_key", "startTime" DESC, "request_id"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 85996430bc5..9429cbabc75 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -684,6 +684,7 @@ model LiteLLM_SpendLogs { @@index([session_id]) @@index([litellm_call_id]) @@index([api_key, startTime]) + @@index([session_id, api_key, startTime(sort: Desc), request_id]) } model LiteLLM_BudgetWindowSpend { diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 85996430bc5..9429cbabc75 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -684,6 +684,7 @@ model LiteLLM_SpendLogs { @@index([session_id]) @@index([litellm_call_id]) @@index([api_key, startTime]) + @@index([session_id, api_key, startTime(sort: Desc), request_id]) } model LiteLLM_BudgetWindowSpend { diff --git a/schema.prisma b/schema.prisma index 85996430bc5..9429cbabc75 100644 --- a/schema.prisma +++ b/schema.prisma @@ -684,6 +684,7 @@ model LiteLLM_SpendLogs { @@index([session_id]) @@index([litellm_call_id]) @@index([api_key, startTime]) + @@index([session_id, api_key, startTime(sort: Desc), request_id]) } model LiteLLM_BudgetWindowSpend { From 8683f50a2790c43efb11761b984916893f09dbc6 Mon Sep 17 00:00:00 2001 From: shrey-berri <333930225+shrey-berri@users.noreply.github.com> Date: Sat, 26 Sep 2026 17:59:43 +0000 Subject: [PATCH 02/10] fix(spend): stop the session view page from grouping the whole date window The session-grouped /spend/logs/ui page grouped every row in the date window before taking 50 sessions, so its cost grew with the window instead of the page. A newest-first page now walks the window newest first and keeps a row only when no newer row of its session exists, stopping after page_size + 1 sessions. A row joins a session either by session_id or, when it has none, by its request_id, and the probe checks both. The tiebreak is written as an OR so the planner cannot derive a startTime range from it and drive the probe off the (api_key, startTime) index under custom plans. The walk reads at most 10000 rows. Ascending pages, OFFSET jumps and walks that run out of rows before filling the page use the previous GROUP BY page, so no request shape gets slower than before. The representative row of each session is one indexed lookup per session. --- .../spend_management_endpoints.py | 138 ++++++++++++++---- 1 file changed, 111 insertions(+), 27 deletions(-) diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 822b827f985..8fcaa137bb1 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -67,6 +67,7 @@ SPEND_LOGS_PAGINATION_COUNT_CAP: Final = 10000 _SESSION_KEY_EXPR: Final = "COALESCE(NULLIF(session_id, ''), request_id)" _SESSION_GROUP_KEY_SQL: Final = f"{_SESSION_KEY_EXPR}, api_key" +_SESSION_HEAD_WALK_ROWS: Final = 10000 _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')" @@ -2979,16 +2980,18 @@ async def _fetch_session_representatives( ) -> list[dict[str, object]]: # mutable-ok: _build_ui_spend_logs_response writes session counts onto each row """Fetch the 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} + 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 "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}, call_type IN {_MCP_CALL_TYPES_SQL}, "startTime" DESC - ) AS session_representatives + AND (request_id = requested.session_key OR session_id = requested.session_key) + AND api_key = requested.key + AND {_SESSION_KEY_EXPR} = requested.session_key + ORDER BY call_type IN {_MCP_CALL_TYPES_SQL}, "startTime" DESC + LIMIT 1 + ) AS rep """ rep_rows: Final[Sequence[dict[str, object]]] = await _query_raw( # mutable-ok: rows are enriched in place prisma_client, @@ -3003,6 +3006,79 @@ async def _fetch_session_representatives( return [rep_by_key[key] for key in session_keys if key in rep_by_key] # mutable-ok: rows are enriched in place +def _session_head_walk_sql(where_clause: str, keyset_clause: str, limit_index: int) -> str: + newer_than_head: Final = f"""newer.api_key = head.api_key + AND {where_clause} + AND (newer."startTime" > head."startTime" + OR (newer."startTime" = head."startTime" AND newer.request_id > head.request_id))""" + return f""" + SELECT head.session_key, head.api_key, head."startTime"::text AS last_activity + FROM ( + SELECT {_SESSION_KEY_EXPR} AS session_key, request_id, api_key, "startTime" + FROM "LiteLLM_SpendLogs" + WHERE {where_clause} + {keyset_clause} + ORDER BY "startTime" DESC, {_SESSION_KEY_EXPR} DESC, api_key DESC + LIMIT {_SESSION_HEAD_WALK_ROWS} + ) AS head + LEFT JOIN LATERAL ( + SELECT 1 AS hit + FROM ( + (SELECT 1 + FROM "LiteLLM_SpendLogs" AS newer + WHERE newer.session_id = head.session_key + AND {newer_than_head} + LIMIT 1) + UNION ALL + (SELECT 1 + FROM "LiteLLM_SpendLogs" AS newer + WHERE newer.request_id = head.session_key + AND NULLIF(newer.session_id, '') IS NULL + AND {newer_than_head}) + ) AS newer_rows + LIMIT 1 + ) AS newer_hit ON TRUE + WHERE newer_hit.hit IS NULL + ORDER BY head."startTime" DESC, head.session_key DESC, head.api_key DESC + LIMIT ${limit_index} + """ + + +async def _window_fits_the_head_walk( + prisma_client: "PrismaClient", where_clause: str, keyset_clause: str, params: Sequence[object] +) -> bool: + count_query: Final = f""" + SELECT COUNT(*) AS total_count + FROM ( + SELECT 1 + FROM "LiteLLM_SpendLogs" + WHERE {where_clause} + {keyset_clause} + LIMIT ${len(params) + 1} + ) AS walked_rows + """ + count_rows: Final[Sequence[_SpendLogsCountRow]] = await _query_raw( + prisma_client, count_query, *params, _SESSION_HEAD_WALK_ROWS + ) + return bool(count_rows) and int(count_rows[0]["total_count"]) < _SESSION_HEAD_WALK_ROWS + + +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, @@ -3048,7 +3124,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 walks the + window newest first and keeps each session's newest row, so it stops after + ``page_size + 1`` sessions instead of grouping the whole window; ascending + and OFFSET pages, and walks that run past ``_SESSION_HEAD_WALK_ROWS`` rows + without filling 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`` @@ -3067,11 +3147,10 @@ async def _ui_session_grouped_spend_logs( direction: Final = "DESC" if sort_desc else "ASC" cursor: Final = _parse_session_cursor(session_cursor) + cursor_bounds: Final = f"(${next_param_index}::timestamp, ${next_param_index + 1}, ${next_param_index + 2})" + keyset_clause: Final = f'AND ("startTime", {_SESSION_KEY_EXPR}, api_key) < {cursor_bounds}' if cursor else "" having_clause: Final = ( - f'HAVING (MAX("startTime"), {_SESSION_GROUP_KEY_SQL}) {cmp_op} ' - f"(${next_param_index}::timestamp, ${next_param_index + 1}, ${next_param_index + 2})" - if cursor - else "" + f'HAVING (MAX("startTime"), {_SESSION_GROUP_KEY_SQL}) {cmp_op} {cursor_bounds}' if cursor else "" ) cursor_params: Final[tuple[object, ...]] = cursor if cursor else () limit_index: Final = next_param_index + len(cursor_params) @@ -3080,21 +3159,26 @@ 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} - """ + page_params: Final = (*sql_params, *cursor_params, page_limit + 1) + walks_heads: Final = sort_desc and offset == 0 and page_limit > 0 + walked_rows: Final[Sequence[_SessionPageRow]] = ( + await _query_raw(prisma_client, _session_head_walk_sql(where_clause, keyset_clause, limit_index), *page_params) + if walks_heads + else () + ) + walk_is_complete: Final = walks_heads and ( + len(walked_rows) > page_limit + or await _window_fits_the_head_walk(prisma_client, where_clause, keyset_clause, (*sql_params, *cursor_params)) + ) 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) + walked_rows + if walk_is_complete or page_limit <= 0 + else await _query_raw( + prisma_client, + _grouped_session_page_sql(where_clause, having_clause, direction, limit_index, offset_clause), + *page_params, + *offset_params, + ) ) has_more: Final = len(page_rows) > page_limit From 7f7627b85938df148910aec1674b39d336fec4f8 Mon Sep 17 00:00:00 2001 From: shrey-berri <333930225+shrey-berri@users.noreply.github.com> Date: Sat, 26 Sep 2026 17:59:43 +0000 Subject: [PATCH 03/10] test(spend): cover session view identity, ordering and read bounds Integration tests on real Postgres seed a 40000 row window and check page order, the exact session total, and that page 1, a cursor page, a filtered page and a page served under the generic plan all stay under a tuple bound. A small seed checks a session whose newest row has no session_id, one session_id under two api_keys, two rows at one startTime, the MCP preference of the representative, cursor walks, and ascending and OFFSET pages. A unit test checks that a head walk cut short falls back to the grouped page. --- .../spend/test_session_view_page_cost.py | 399 ++++++++++++++++++ .../test_spend_management_endpoints.py | 89 +++- .../test_spend_query_optimization.py | 16 +- 3 files changed, 480 insertions(+), 24 deletions(-) create mode 100644 tests/integration/spend/test_session_view_page_cost.py 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..6856e9b04af --- /dev/null +++ b/tests/integration/spend/test_session_view_page_cost.py @@ -0,0 +1,399 @@ +"""The session-grouped /spend/logs/ui page lists each session once and reads the window at most once. + +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. +""" + +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 +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" + + +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.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}" + ) + + +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-llm", f"{marker}-mcp", at(80)), + _SeedRow(f"{marker}-mcp-tool", f"{marker}-mcp", at(90), call_type="call_mcp_tool"), + ) + + +@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}-mcp-llm", "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, represented by its 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, and one session_id under two " + "api_keys is two sessions" + ) + assert page["total"] == 6 + 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"] == 6, "the capped count must count one session_id under two api_keys as two sessions" + assert walked == _newest_first(small_seed) + assert [page["has_more"] for page in cursor_pages] == [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"] == 6 + assert _sessions(offset_page) == _newest_first(small_seed)[2:4] + assert offset_page["has_more"] is True 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 c6a4173b583..9bf0540b18e 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 @@ -13,6 +13,7 @@ from fastapi.testclient import TestClient import litellm import litellm.proxy.proxy_server as ps +from litellm.proxy.spend_tracking.spend_management_endpoints import _SESSION_HEAD_WALK_ROWS def _default_date_range(): @@ -7207,7 +7208,8 @@ 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_grouped_mock_prisma(session_page_rows, session_total, representative_rows): @@ -7216,7 +7218,7 @@ def _session_grouped_mock_prisma(session_page_rows, session_total, representativ async def mock_query_raw(sql_query, *params): if "COUNT(*) AS total_count" in sql_query: return [{"total_count": session_total}] - if "DISTINCT ON" in sql_query: + if "CROSS JOIN LATERAL" in sql_query or "DISTINCT ON" in sql_query: return representative_rows if "COALESCE(SUM(spend)" in sql_query: return [ @@ -7268,7 +7270,7 @@ def _session_grouped_paginating_prisma(sessions, counted_total=None): async def mock_query_raw(sql_query, *params): if "COUNT(*) AS total_count" in sql_query: return [{"total_count": min(len(sessions) if counted_total is None else counted_total, params[-1])}] - if "DISTINCT ON" in sql_query: + if "CROSS JOIN LATERAL" 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 [] @@ -7331,21 +7333,14 @@ async def test_ui_view_spend_logs_group_by_session_first_page(client, monkeypatc 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}, call_type IN ('call_mcp_tool', 'list_mcp_tools'), \"startTime\" DESC" - in rep_call[0] - ), "the session representative must prefer the newest non-MCP call" assert rep_call[-2] == ["sess-1", "req-solo"] assert rep_call[-1] == ["hashed-key", "hashed-key"] finally: @@ -7354,7 +7349,7 @@ async def test_ui_view_spend_logs_group_by_session_first_page(client, monkeypatc @pytest.mark.asyncio async def test_ui_view_spend_logs_group_by_session_cursor_page(client, monkeypatch): - """A session_cursor becomes a HAVING keyset predicate instead of an OFFSET.""" + """A session_cursor becomes a WHERE keyset predicate instead of an OFFSET.""" page_rows = [_session_page_row("sess-2", "2026-08-29 07:00:00")] reps = [_session_representative_row("req-2", "sess-2")] mock_prisma = _session_grouped_mock_prisma(page_rows, 3, reps) @@ -7387,14 +7382,82 @@ async def test_ui_view_spend_logs_group_by_session_cursor_page(client, monkeypat page_query_call = mock_prisma.db.query_raw.await_args_list[0] page_query_sql = page_query_call.args[0] - assert f'HAVING (MAX("startTime"), {SESSION_GROUP_KEY_SQL}) <' in page_query_sql + assert "HAVING" not 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)" + ), "cursor params must line up with (startTime, session key, api_key)" + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +def _session_walk_mock_prisma(walked_rows, walk_window_rows, grouped_rows): + """Mock prisma serving the head walk, its window-size check and the grouped fallback page separately.""" + + async def mock_query_raw(sql_query, *params): + if "AS walked_rows" in sql_query: + return [{"total_count": walk_window_rows}] + if "COUNT(*) AS total_count" in sql_query: + return [{"total_count": 3}] + if "CROSS JOIN LATERAL" 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 [] + if "LEFT JOIN LATERAL" in sql_query: + return walked_rows + 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( + ("walk_window_rows", "expected_request_ids", "expected_has_more"), + [ + pytest.param( + _SESSION_HEAD_WALK_ROWS, ["req-grouped-1", "req-grouped-2"], True, id="walk-cut-short-groups-the-window" + ), + pytest.param(_SESSION_HEAD_WALK_ROWS - 1, ["req-walked-1"], False, id="walk-covered-the-window"), + ], +) +async def test_ui_view_spend_logs_group_by_session_falls_back_when_the_walk_is_cut_short( + client, monkeypatch, walk_window_rows, expected_request_ids, expected_has_more +): + """A head walk that found fewer sessions than the page needs is only trusted when it read the whole window.""" + mock_prisma = _session_walk_mock_prisma( + walked_rows=[_session_page_row("walked-1", "2026-08-29 10:00:00")], + walk_window_rows=walk_window_rows, + 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) 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 54e5a6d5385..449196b45c2 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 @@ -513,9 +513,8 @@ 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). @@ -526,7 +525,6 @@ async def test_spend_logs_ui_group_by_session_paginates_sessions(monkeypatch): 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) @@ -539,7 +537,7 @@ async def test_spend_logs_ui_group_by_session_paginates_sessions(monkeypatch): async def mock_query_raw(sql_query, *params): if "COUNT(*) AS total_count" in sql_query: return [{"total_count": 60}] - if "DISTINCT ON" in sql_query: + if "CROSS JOIN LATERAL" in sql_query or "DISTINCT ON" in sql_query: return representative_rows return session_rows @@ -571,22 +569,18 @@ async def test_spend_logs_ui_group_by_session_paginates_sessions(monkeypatch): emitted = [call[0] for call in mock_prisma.db.query_raw.call_args_list] page_sql = emitted[0][0] - assert f"GROUP BY {group_key}" in page_sql, f"page must select sessions, not calls. SQL was:\n{page_sql}" assert "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 IN ('call_mcp_tool', 'list_mcp_tools'), \"startTime\" DESC" in rep_sql, ( - "the session representative must prefer the newest non-MCP call" - ) + assert emitted[2][-2] == [row["session_key"] for row in session_rows[:50]] + assert emitted[2][-1] == ["k"] * 50 assert "COUNT(*) OVER ()" not in rep_sql assert [row["request_id"] for row in response["data"]] == ["req-1", "req-2"] From 1f3f512e504330a9d3a5e6e37806e750c4021db2 Mon Sep 17 00:00:00 2001 From: shrey-berri <333930225+shrey-berri@users.noreply.github.com> Date: Sat, 26 Sep 2026 19:39:37 +0000 Subject: [PATCH 04/10] perf(spend): find session heads by walking newest first without a new index Reading the window newest first, the first row seen of each session is that session's newest row, so the per-row newer-row probe and the index it needed are unnecessary. The page walks the existing (startTime, request_id) index in batches sized from how quickly new sessions appear, keeps the first row of each session in SQL, and stops once page_size + 1 sessions are settled. Filters are evaluated on each batch rather than steering the index choice, and cursor pages drop sessions that already have a row at or above the cursor with one lookup per candidate session. A walk that would read more than 20000 rows falls back to the grouped page. Drops the session head index and its migration. --- .../migration.sql | 2 - .../litellm_proxy_extras/schema.prisma | 1 - litellm/proxy/schema.prisma | 1 - .../spend_management_endpoints.py | 279 +++++++++++++----- schema.prisma | 1 - 5 files changed, 205 insertions(+), 79 deletions(-) delete mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260924204413_add_spend_logs_session_head_index/migration.sql diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260924204413_add_spend_logs_session_head_index/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260924204413_add_spend_logs_session_head_index/migration.sql deleted file mode 100644 index da26297b59c..00000000000 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260924204413_add_spend_logs_session_head_index/migration.sql +++ /dev/null @@ -1,2 +0,0 @@ --- CreateIndex (CONCURRENTLY): keep this a single statement so Prisma Migrate applies it outside a transaction -CREATE INDEX CONCURRENTLY IF NOT EXISTS "LiteLLM_SpendLogs_session_id_api_key_startTime_request_id_idx" ON "LiteLLM_SpendLogs"("session_id", "api_key", "startTime" DESC, "request_id"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 9429cbabc75..85996430bc5 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -684,7 +684,6 @@ model LiteLLM_SpendLogs { @@index([session_id]) @@index([litellm_call_id]) @@index([api_key, startTime]) - @@index([session_id, api_key, startTime(sort: Desc), request_id]) } model LiteLLM_BudgetWindowSpend { diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 9429cbabc75..85996430bc5 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -684,7 +684,6 @@ model LiteLLM_SpendLogs { @@index([session_id]) @@index([litellm_call_id]) @@index([api_key, startTime]) - @@index([session_id, api_key, startTime(sort: Desc), request_id]) } model LiteLLM_BudgetWindowSpend { diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 8fcaa137bb1..c140b39a3a3 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -67,7 +67,8 @@ SPEND_LOGS_PAGINATION_COUNT_CAP: Final = 10000 _SESSION_KEY_EXPR: Final = "COALESCE(NULLIF(session_id, ''), request_id)" _SESSION_GROUP_KEY_SQL: Final = f"{_SESSION_KEY_EXPR}, api_key" -_SESSION_HEAD_WALK_ROWS: Final = 10000 +_SESSION_WALK_FIRST_BATCH_ROWS: Final = 1000 +_SESSION_WALK_MAX_ROWS: Final = 20000 _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')" @@ -2740,6 +2741,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( @@ -2849,6 +2851,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, @@ -3006,61 +3009,181 @@ async def _fetch_session_representatives( return [rep_by_key[key] for key in session_keys if key in rep_by_key] # mutable-ok: rows are enriched in place -def _session_head_walk_sql(where_clause: str, keyset_clause: str, limit_index: int) -> str: - newer_than_head: Final = f"""newer.api_key = head.api_key - AND {where_clause} - AND (newer."startTime" > head."startTime" - OR (newer."startTime" = head."startTime" AND newer.request_id > head.request_id))""" - return f""" - SELECT head.session_key, head.api_key, head."startTime"::text AS last_activity - FROM ( - SELECT {_SESSION_KEY_EXPR} AS session_key, request_id, api_key, "startTime" - FROM "LiteLLM_SpendLogs" - WHERE {where_clause} - {keyset_clause} - ORDER BY "startTime" DESC, {_SESSION_KEY_EXPR} DESC, api_key DESC - LIMIT {_SESSION_HEAD_WALK_ROWS} - ) AS head - LEFT JOIN LATERAL ( - SELECT 1 AS hit +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 1 - FROM "LiteLLM_SpendLogs" AS newer - WHERE newer.session_id = head.session_key - AND {newer_than_head} - LIMIT 1) - UNION ALL - (SELECT 1 - FROM "LiteLLM_SpendLogs" AS newer - WHERE newer.request_id = head.session_key - AND NULLIF(newer.session_id, '') IS NULL - AND {newer_than_head}) - ) AS newer_rows - LIMIT 1 - ) AS newer_hit ON TRUE - WHERE newer_hit.hit IS NULL - ORDER BY head."startTime" DESC, head.session_key DESC, head.api_key DESC - LIMIT ${limit_index} - """ + 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 "LiteLLM_SpendLogs" + WHERE {self.where_clause} + AND api_key = requested.key + AND (session_id = requested.session_key + OR (request_id = requested.session_key AND NULLIF(session_id, '') IS NULL)) + 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], # mutable-ok: prisma serializes array params from a list + [row["api_key"] for row in candidates], # mutable-ok: prisma serializes array params from a list + ) + return frozenset((row["session_key"], row["api_key"]) for row in shown) -async def _window_fits_the_head_walk( - prisma_client: "PrismaClient", where_clause: str, keyset_clause: str, params: Sequence[object] -) -> bool: - count_query: Final = f""" - SELECT COUNT(*) AS total_count - FROM ( - SELECT 1 - FROM "LiteLLM_SpendLogs" - WHERE {where_clause} - {keyset_clause} - LIMIT ${len(params) + 1} - ) AS walked_rows - """ - count_rows: Final[Sequence[_SpendLogsCountRow]] = await _query_raw( - prisma_client, count_query, *params, _SESSION_HEAD_WALK_ROWS - ) - return bool(count_rows) and int(count_rows[0]["total_count"]) < _SESSION_HEAD_WALK_ROWS + 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, + ) + 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( @@ -3107,6 +3230,7 @@ async def _count_grouped_sessions( 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, @@ -3124,11 +3248,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 newest-first page walks the - window newest first and keeps each session's newest row, so it stops after - ``page_size + 1`` sessions instead of grouping the whole window; ascending - and OFFSET pages, and walks that run past ``_SESSION_HEAD_WALK_ROWS`` rows - without filling the page, group the window instead. 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 every ``_SESSION_WALK_BATCH_ROWS`` batch 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`` @@ -3147,10 +3271,11 @@ async def _ui_session_grouped_spend_logs( direction: Final = "DESC" if sort_desc else "ASC" cursor: Final = _parse_session_cursor(session_cursor) - cursor_bounds: Final = f"(${next_param_index}::timestamp, ${next_param_index + 1}, ${next_param_index + 2})" - keyset_clause: Final = f'AND ("startTime", {_SESSION_KEY_EXPR}, api_key) < {cursor_bounds}' if cursor else "" having_clause: Final = ( - f'HAVING (MAX("startTime"), {_SESSION_GROUP_KEY_SQL}) {cmp_op} {cursor_bounds}' if cursor else "" + f'HAVING (MAX("startTime"), {_SESSION_GROUP_KEY_SQL}) {cmp_op} ' + f"(${next_param_index}::timestamp, ${next_param_index + 1}, ${next_param_index + 2})" + if cursor + else "" ) cursor_params: Final[tuple[object, ...]] = cursor if cursor else () limit_index: Final = next_param_index + len(cursor_params) @@ -3159,24 +3284,30 @@ 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_params: Final = (*sql_params, *cursor_params, page_limit + 1) - walks_heads: Final = sort_desc and offset == 0 and page_limit > 0 - walked_rows: Final[Sequence[_SessionPageRow]] = ( - await _query_raw(prisma_client, _session_head_walk_sql(where_clause, keyset_clause, limit_index), *page_params) - if walks_heads - else () - ) - walk_is_complete: Final = walks_heads and ( - len(walked_rows) > page_limit - or await _window_fits_the_head_walk(prisma_client, where_clause, keyset_clause, (*sql_params, *cursor_params)) + walked_rows: Final = ( + await _SessionHeadWalk( + prisma_client=prisma_client, + window_clause=" AND ".join(window_conditions) if window_conditions else "TRUE", + 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]] = ( - walked_rows - if walk_is_complete or page_limit <= 0 + () + if page_limit <= 0 + 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), - *page_params, + *sql_params, + *cursor_params, + page_limit + 1, *offset_params, ) ) diff --git a/schema.prisma b/schema.prisma index 9429cbabc75..85996430bc5 100644 --- a/schema.prisma +++ b/schema.prisma @@ -684,7 +684,6 @@ model LiteLLM_SpendLogs { @@index([session_id]) @@index([litellm_call_id]) @@index([api_key, startTime]) - @@index([session_id, api_key, startTime(sort: Desc), request_id]) } model LiteLLM_BudgetWindowSpend { From 9adfe67c3435b35ecbd9511e88a6f59e9fcfe03f Mon Sep 17 00:00:00 2001 From: shrey-berri <333930225+shrey-berri@users.noreply.github.com> Date: Sat, 26 Sep 2026 19:39:37 +0000 Subject: [PATCH 05/10] test(spend): cover the cursor check, representative guard, filters and ties Adds a session with rows on both sides of the cursor, a row whose request_id equals another session's key, and a filtered page to the small integration seed, plus a unit test that a walk reads past a tied timestamp before settling. --- .../spend/test_session_view_page_cost.py | 29 +++-- .../test_spend_management_endpoints.py | 112 +++++++++++++++--- .../test_spend_query_optimization.py | 7 +- 3 files changed, 123 insertions(+), 25 deletions(-) diff --git a/tests/integration/spend/test_session_view_page_cost.py b/tests/integration/spend/test_session_view_page_cost.py index 6856e9b04af..b4988466791 100644 --- a/tests/integration/spend/test_session_view_page_cost.py +++ b/tests/integration/spend/test_session_view_page_cost.py @@ -304,8 +304,10 @@ def _small_rows(marker: str) -> tuple[_SeedRow, ...]: _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}-shared", f"{marker}-other", at(95), api_key=f"{marker}-key-b"), ) @@ -346,6 +348,7 @@ def _tie_as_one(request_id: str) -> str: def _newest_first(marker: str) -> list[tuple[str, str]]: return [ + (f"{marker}-shared", f"{marker}-key-b"), (f"{marker}-mcp-llm", "K"), (f"{marker}-tie", "K"), (f"{marker}-shared-b", f"{marker}-key-b"), @@ -360,10 +363,10 @@ def test_session_view_lists_each_session_once(gateway: Gateway, small_seed: str) assert _sessions(page) == _newest_first(small_seed), ( "each session once, represented by its 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, and one session_id under two " - "api_keys is two sessions" + "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"] == 6 + assert page["total"] == 7 assert page["has_more"] is False @@ -382,9 +385,11 @@ def test_session_view_cursor_walk_counts_sessions_per_api_key(gateway: Gateway, ) walked: Final = [session for page in cursor_pages for session in _sessions(page)] - assert first["total"] == 6, "the capped count must count one session_id under two api_keys as two sessions" - assert walked == _newest_first(small_seed) - assert [page["has_more"] for page in cursor_pages] == [True, True, False] + 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( @@ -394,6 +399,16 @@ def test_session_view_ascending_and_offset_pages_match_the_newest_first_order( 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"] == 6 + 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 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 9bf0540b18e..22a8c1c58f2 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 @@ -13,7 +13,6 @@ from fastapi.testclient import TestClient import litellm import litellm.proxy.proxy_server as ps -from litellm.proxy.spend_tracking.spend_management_endpoints import _SESSION_HEAD_WALK_ROWS def _default_date_range(): @@ -7212,10 +7211,32 @@ 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_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 "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: @@ -7268,6 +7289,10 @@ 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 "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: @@ -7335,7 +7360,6 @@ async def test_ui_view_spend_logs_group_by_session_first_page(client, monkeypatc page_query_sql = emitted[0][0] 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 "LIMIT" in count_sql and "FROM (" in count_sql, "the grouped count must stay bounded" @@ -7394,20 +7418,18 @@ async def test_ui_view_spend_logs_group_by_session_cursor_page(client, monkeypat app.dependency_overrides.pop(ps.user_api_key_auth, None) -def _session_walk_mock_prisma(walked_rows, walk_window_rows, grouped_rows): - """Mock prisma serving the head walk, its window-size check and the grouped fallback page separately.""" +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 "AS walked_rows" in sql_query: - return [{"total_count": walk_window_rows}] + if "WITH batch AS MATERIALIZED" in sql_query: + return _session_walk_batch(walked_rows, min(window_rows_per_batch, params[-1])) if "COUNT(*) AS total_count" in sql_query: return [{"total_count": 3}] if "CROSS JOIN LATERAL" 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 [] - if "LEFT JOIN LATERAL" in sql_query: - return walked_rows return grouped_rows mock_prisma = MagicMock() @@ -7418,21 +7440,19 @@ def _session_walk_mock_prisma(walked_rows, walk_window_rows, grouped_rows): @pytest.mark.asyncio @pytest.mark.parametrize( - ("walk_window_rows", "expected_request_ids", "expected_has_more"), + ("window_rows_per_batch", "expected_request_ids", "expected_has_more"), [ - pytest.param( - _SESSION_HEAD_WALK_ROWS, ["req-grouped-1", "req-grouped-2"], True, id="walk-cut-short-groups-the-window" - ), - pytest.param(_SESSION_HEAD_WALK_ROWS - 1, ["req-walked-1"], False, id="walk-covered-the-window"), + 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_is_cut_short( - client, monkeypatch, walk_window_rows, expected_request_ids, expected_has_more +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 head walk that found fewer sessions than the page needs is only trusted when it read the whole window.""" + """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")], - walk_window_rows=walk_window_rows, + 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"), @@ -7462,6 +7482,64 @@ async def test_ui_view_spend_logs_group_by_session_falls_back_when_the_walk_is_c 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 "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("c", tied), _walk_row("c", "2026-08-29 09:00:00", is_edge=True, batch_rows=1)] + if "COUNT(*) AS total_count" in sql_query: + return [{"total_count": 3}] + if "CROSS JOIN LATERAL" 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-c"] + assert data["next_session_cursor"] == f"{tied}|hashed-key|c" + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + @pytest.mark.asyncio async def test_ui_view_spend_logs_group_by_session_jumps_to_page_without_cursor(client, monkeypatch): """page > 1 with no session_cursor (the UI's last-page jump) serves the sessions that page starts at.""" 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 449196b45c2..7b358015ec9 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 @@ -535,6 +535,12 @@ 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 "COUNT(*) AS total_count" in sql_query: return [{"total_count": 60}] if "CROSS JOIN LATERAL" in sql_query or "DISTINCT ON" in sql_query: @@ -570,7 +576,6 @@ async def test_spend_logs_ui_group_by_session_paginates_sessions(monkeypatch): emitted = [call[0] for call in mock_prisma.db.query_raw.call_args_list] page_sql = emitted[0][0] assert "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] From 8729832b8a66077cbc489f548c14fc17c2e5bcdc Mon Sep 17 00:00:00 2001 From: shrey-berri <333930225+shrey-berri@users.noreply.github.com> Date: Sat, 26 Sep 2026 20:18:19 +0000 Subject: [PATCH 06/10] perf(spend): keep the per-session lookups off the window scan under generic plans After five executions Postgres switches the representative and cursor lookups to a generic plan, which ANDed each session's rows with a scan of the whole date window on the startTime index. With 50 sessions of 150 rows that took about 500 ms instead of about 10 ms. Find the session's rows by its key first behind OFFSET 0, then apply the window and filters to them. --- .../spend_management_endpoints.py | 20 ++++++++++++------- 1 file changed, 13 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index c140b39a3a3..368a43e4a19 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -2974,6 +2974,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, @@ -2987,10 +2997,8 @@ async def _fetch_session_representatives( 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 "LiteLLM_SpendLogs" + FROM ({_session_rows_sql("requested.session_key", "requested.key")}) AS "LiteLLM_SpendLogs" WHERE {where_clause} - AND (request_id = requested.session_key OR session_id = requested.session_key) - AND api_key = requested.key AND {_SESSION_KEY_EXPR} = requested.session_key ORDER BY call_type IN {_MCP_CALL_TYPES_SQL}, "startTime" DESC LIMIT 1 @@ -3116,11 +3124,9 @@ class _SessionHeadWalk: FROM unnest(${len(params) + 1}::text[], ${len(params) + 2}::text[]) AS requested(session_key, key) CROSS JOIN LATERAL ( SELECT 1 - FROM "LiteLLM_SpendLogs" + FROM ({_session_rows_sql("requested.session_key", "requested.key")}) AS "LiteLLM_SpendLogs" WHERE {self.where_clause} - AND api_key = requested.key - AND (session_id = requested.session_key - OR (request_id = requested.session_key AND NULLIF(session_id, '') IS NULL)) + AND {_SESSION_KEY_EXPR} = requested.session_key AND ("startTime", requested.session_key, requested.key) >= {self._cursor_bounds()} LIMIT 1 ) AS newer_row 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 07/10] 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 From b7755a48accef77df346a9197ac546df33adbc97 Mon Sep 17 00:00:00 2001 From: shrey kharbanda Date: Fri, 23 Oct 2026 15:08:38 -0700 Subject: [PATCH 08/10] fix(spend): preserve database ordering across tied sessions --- .../spend_management_endpoints.py | 13 ++++++ .../spend/test_session_view_page_cost.py | 43 ++++++++++++++++++- .../test_spend_management_endpoints.py | 2 + 3 files changed, 57 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index bd96d11f39e..620a4f1be83 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -3276,6 +3276,19 @@ class _SessionHeadWalk: rows_read=total_read, size=next_size, ) + if len(frozenset(row["sort_time"] for row in all_heads)) < len(all_heads): + 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], # mutable-ok: prisma serializes array params from a list + [row["api_key"] for row in all_heads], # mutable-ok: prisma serializes array params from a list + [row["last_activity"] for row in all_heads], # mutable-ok: prisma serializes array params from a list + 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 ) diff --git a/tests/integration/spend/test_session_view_page_cost.py b/tests/integration/spend/test_session_view_page_cost.py index 9986698ae04..7e63dca53d6 100644 --- a/tests/integration/spend/test_session_view_page_cost.py +++ b/tests/integration/spend/test_session_view_page_cost.py @@ -17,7 +17,7 @@ from typing import Final import psycopg import pytest from integration._support.client import Gateway, string_value -from integration._support.database import ROWS +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 @@ -469,3 +469,44 @@ def test_session_view_filtered_page_only_lists_matching_sessions(gateway: Gatewa (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/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index dbad3834459..f5e46e23b21 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 @@ -7303,6 +7303,8 @@ async def test_ui_view_spend_logs_group_by_session_walk_finishes_a_tied_timestam 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 ("c", "b")] if "WITH batch AS MATERIALIZED" in sql_query: if '("startTime", request_id) <' not in sql_query: return [ From f3e2af348603e2ed6ac646eddbbeb9892044f485 Mon Sep 17 00:00:00 2001 From: shrey kharbanda Date: Fri, 23 Oct 2026 15:38:15 -0700 Subject: [PATCH 09/10] perf(spend): bound session enrichment reads on cold pages --- .../spend_management_endpoints.py | 20 +++++++++++-------- .../test_spend_management_endpoints.py | 8 ++++---- .../test_spend_query_optimization.py | 2 +- 3 files changed, 17 insertions(+), 13 deletions(-) diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 9f6b034804c..e1a41c7e2ab 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -3266,8 +3266,8 @@ class _SessionHeadWalk: self.prisma_client, query, *params, - [row["session_key"] for row in candidates], # mutable-ok: prisma serializes array params from a list - [row["api_key"] for row in candidates], # mutable-ok: prisma serializes array params from a list + [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) @@ -3321,9 +3321,9 @@ class _SessionHeadWalk: 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], # mutable-ok: prisma serializes array params from a list - [row["api_key"] for row in all_heads], # mutable-ok: prisma serializes array params from a list - [row["last_activity"] for row in all_heads], # mutable-ok: prisma serializes array params from a list + [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) @@ -5031,9 +5031,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/unit/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py index c0893ee95f9..cb79fcfbcb8 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py @@ -7212,7 +7212,7 @@ def _session_grouped_mock_prisma(session_page_rows, session_total, representativ 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: + if "AS rep" in sql_query or "DISTINCT ON" in sql_query: return representative_rows if "COALESCE(SUM(spend)" in sql_query: return [ @@ -7270,7 +7270,7 @@ def _session_grouped_paginating_prisma(sessions, counted_total=None): 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: + 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 [] @@ -7400,7 +7400,7 @@ def _session_walk_mock_prisma(walked_rows, window_rows_per_batch, grouped_rows): return _session_count_walk(3) if "COUNT(*) AS total_count" in sql_query: return [{"total_count": 3}] - if "CROSS JOIN LATERAL" in sql_query: + 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 [] @@ -7488,7 +7488,7 @@ async def test_ui_view_spend_logs_group_by_session_walk_finishes_a_tied_timestam return _session_count_walk(3) if "COUNT(*) AS total_count" in sql_query: return [{"total_count": 3}] - if "CROSS JOIN LATERAL" in sql_query: + if "AS rep" in sql_query: return [_session_representative_row(f"req-{session_key}", session_key) for session_key in params[-2]] return [] 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 72c882bcc74..5d6ae37f527 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_query_optimization.py +++ b/tests/unit/proxy/spend_tracking/test_spend_query_optimization.py @@ -539,7 +539,7 @@ async def test_spend_logs_ui_group_by_session_paginates_sessions(monkeypatch): 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: + if "AS rep" in sql_query or "DISTINCT ON" in sql_query: return representative_rows return session_rows From de2d96e44c262300ba77a9ba23cade91802ca409 Mon Sep 17 00:00:00 2001 From: shrey kharbanda Date: Fri, 23 Oct 2026 16:37:59 -0700 Subject: [PATCH 10/10] fix(spend): preserve bounded spans for session tie ordering --- .../spend_management_endpoints.py | 24 +++++++------ .../spend/test_session_view_page_cost.py | 18 ++++++++-- .../test_spend_management_endpoints.py | 35 +++++-------------- .../test_spend_query_optimization.py | 9 ----- 4 files changed, 37 insertions(+), 49 deletions(-) diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index e1a41c7e2ab..5527d051458 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -51,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, @@ -3315,17 +3316,18 @@ class _SessionHeadWalk: size=next_size, ) if len(frozenset(row["sort_time"] for row in all_heads)) < len(all_heads): - 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, - ) + 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 diff --git a/tests/integration/spend/test_session_view_page_cost.py b/tests/integration/spend/test_session_view_page_cost.py index 7e63dca53d6..056ccf5e4e2 100644 --- a/tests/integration/spend/test_session_view_page_cost.py +++ b/tests/integration/spend/test_session_view_page_cost.py @@ -62,6 +62,7 @@ class _SeedRow: start_time: str api_key: str = "K" call_type: str = "acompletion" + end_time: str | None = None def _hand_placed_rows(marker: str) -> tuple[_SeedRow, ...]: @@ -104,7 +105,15 @@ def _insert_rows(rows: tuple[_SeedRow, ...]) -> None: '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.start_time, KEY_ALIAS) + ( + 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 ], ) @@ -364,6 +373,8 @@ def _small_rows(marker: str) -> tuple[_SeedRow, ...]: _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"), ) @@ -406,7 +417,7 @@ def _tie_as_one(request_id: str) -> str: def _newest_first(marker: str) -> list[tuple[str, str]]: return [ (f"{marker}-shared", f"{marker}-key-b"), - (f"{marker}-mcp-llm", "K"), + (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"), @@ -419,7 +430,8 @@ def test_session_view_lists_each_session_once(gateway: Gateway, small_seed: str) page: Final = _small_page(gateway, sort_order="desc", page_size="50", page="1") assert _sessions(page) == _newest_first(small_seed), ( - "each session once, represented by its newest non-MCP row: a row without a session_id newer than the rows " + "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" ) 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 cb79fcfbcb8..b83640a8cb1 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py @@ -7330,22 +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 "OFFSET" not in page_query_sql - assert "HAVING" not in page_query_sql - - rep_call = emitted[2] - 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 WHERE 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) @@ -7375,17 +7366,6 @@ 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 "HAVING" not 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 (startTime, session key, api_key)" finally: app.dependency_overrides.pop(ps.user_api_key_auth, None) @@ -7475,15 +7455,18 @@ async def test_ui_view_spend_logs_group_by_session_walk_finishes_a_tied_timestam 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 ("c", "b")] + 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("B", tied), _walk_row("a", tied), _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)] + 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: @@ -7512,8 +7495,8 @@ async def test_ui_view_spend_logs_group_by_session_walk_finishes_a_tied_timestam ) assert response.status_code == 200, response.text data = response.json() - assert [row["request_id"] for row in data["data"]] == ["req-c"] - assert data["next_session_cursor"] == f"{tied}|hashed-key|c" + 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) 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 5d6ae37f527..aabce025a10 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_query_optimization.py +++ b/tests/unit/proxy/spend_tracking/test_spend_query_optimization.py @@ -570,15 +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 "OFFSET" not in page_sql, "the startTime page must be keyset-selected, not offset-selected" - - rep_sql = emitted[2][0] - assert emitted[2][-2] == [row["session_key"] for row in session_rows[:50]] - assert emitted[2][-1] == ["k"] * 50 - 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