Merge 5ebcf8ece1 into 24a05b29a6

litellm-proxy-extras/litellm_proxy_extras/migrations/20260911000000_spend_logs_session_window_index/migration.sql

@@ -0,0 +1,14 @@
+-- CreateIndex (CONCURRENTLY)
+--
+-- Covering index for the session-grouped /spend/logs/ui listing. The page,
+-- count, and representative queries all GROUP BY
+-- COALESCE(NULLIF(session_id, ''), request_id), api_key over a startTime
+-- window; with only the plain startTime index they aggregate full heap rows
+-- (hundreds of KB each from messages/response), which is the 40s+ page load.
+-- This index makes those window scans index-only.
+--
+-- CREATE INDEX CONCURRENTLY cannot run inside a transaction, so this
+-- migration must stay a single statement (see
+-- 20260415120000_health_check_latest_per_model_index for the full
+-- disclaimer).
+CREATE INDEX CONCURRENTLY IF NOT EXISTS "LiteLLM_SpendLogs_session_window_idx" ON "LiteLLM_SpendLogs"("startTime", session_id, request_id, api_key, status);

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([startTime, session_id, request_id, api_key, status], map: "LiteLLM_SpendLogs_session_window_idx")
 }
 
 model LiteLLM_BudgetWindowSpend {

litellm/proxy/schema.prisma

@@ -684,6 +684,7 @@ model LiteLLM_SpendLogs {
   @@index([session_id])
   @@index([litellm_call_id])
   @@index([api_key, startTime])
+  @@index([startTime, session_id, request_id, api_key, status], map: "LiteLLM_SpendLogs_session_window_idx")
 }
 
 model LiteLLM_BudgetWindowSpend {

litellm/proxy/spend_tracking/spend_management_endpoints.py

@@ -1,4 +1,5 @@
 #### SPEND MANAGEMENT #####
+import asyncio
 import collections
 import json
 import os
@@ -2891,13 +2892,6 @@ async def ui_view_spend_logs(
                 LIMIT ${p}
             ) AS bounded_matches
         """
-        count_rows: Final[Sequence[_SpendLogsCountRow] | None] = await _query_raw_or_none(
-            prisma_client, count_query, *sql_params, SPEND_LOGS_PAGINATION_COUNT_CAP + 1
-        )
-        raw_total: Final = int(count_rows[0]["total_count"]) if count_rows else 0
-        total_is_capped: Final = raw_total > SPEND_LOGS_PAGINATION_COUNT_CAP
-        total_records: Final = SPEND_LOGS_PAGINATION_COUNT_CAP if total_is_capped else raw_total
-
         sql_query: Final = (
             f"""
                 SELECT * FROM (
@@ -2920,9 +2914,14 @@ async def ui_view_spend_logs(
             LIMIT ${p} OFFSET ${p + 1}
         """
         )
-        sql_params.extend([page_size, skip])
-
-        data: Final = await prisma_client.db.query_raw(sql_query, *sql_params)
+        count_task: Final[asyncio.Task[Sequence[_SpendLogsCountRow] | None]] = asyncio.create_task(
+            _query_raw_or_none(prisma_client, count_query, *sql_params, SPEND_LOGS_PAGINATION_COUNT_CAP + 1)
+        )
+        data: Final = await prisma_client.db.query_raw(sql_query, *sql_params, page_size, skip)
+        count_rows: Final = await count_task
+        raw_total: Final = int(count_rows[0]["total_count"]) if count_rows else 0
+        total_is_capped: Final = raw_total > SPEND_LOGS_PAGINATION_COUNT_CAP
+        total_records: Final = SPEND_LOGS_PAGINATION_COUNT_CAP if total_is_capped else raw_total
 
         if request_id is not None and not is_v2 and not is_admin_view:
             await _assert_user_owns_fetched_spend_rows(
@@ -2937,8 +2936,6 @@ async def ui_view_spend_logs(
         # Calculate total pages
         total_pages: Final = (total_records + page_size - 1) // page_size
 
-        verbose_proxy_logger.debug("data= %s", json.dumps(data, indent=4, default=str))
-
         return await _build_ui_spend_logs_response(
             prisma_client,
             data,
@@ -2977,15 +2974,25 @@ async def _fetch_session_representatives(
     next_param_index: int,
     session_keys: Sequence[tuple[str, str]],
 ) -> 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."""
+    """Fetch the newest non-MCP row of each ``(session_key, api_key)`` session, in ``session_keys`` order.
+
+    The SQL matches candidates on the plain ``session_id`` / ``request_id`` /
+    ``api_key`` columns so the ``(startTime, session_id, request_id, api_key)``
+    index applies; a row-tuple ``IN`` on the COALESCE group key forces a full
+    window scan. That predicate admits a superset of the requested pairs (the
+    cross product of session ids and api keys, plus rows whose request_id
+    collides with a requested key), so the exact ``(session_key, api_key)``
+    pairing is restored by the ``rep_by_key`` lookup below.
+    """
     rep_query: Final = f"""
         SELECT * FROM (
             SELECT DISTINCT ON ({_SESSION_GROUP_KEY_SQL})
                 {_SPEND_LOG_LIST_COLUMNS}
             FROM "LiteLLM_SpendLogs"
             WHERE {where_clause}
-              AND ({_SESSION_GROUP_KEY_SQL}) IN (
-                  SELECT * FROM unnest(${next_param_index}::text[], ${next_param_index + 1}::text[])
+              AND (
+                  (session_id = ANY(${next_param_index}::text[]) AND api_key = ANY(${next_param_index + 1}::text[]))
+                  OR request_id = ANY(${next_param_index}::text[])
               )
             ORDER BY {_SESSION_GROUP_KEY_SQL}, call_type IN {_MCP_CALL_TYPES_SQL}, "startTime" DESC
         ) AS session_representatives
@@ -3091,6 +3098,9 @@ async def _ui_session_grouped_spend_logs(
         ORDER BY MAX("startTime") {direction}, {_SESSION_KEY_EXPR} {direction}, api_key {direction}
         LIMIT ${limit_index} {offset_clause}
     """
+    count_task: Final[asyncio.Task[tuple[int, bool]]] = asyncio.create_task(
+        _count_grouped_sessions(prisma_client, where_clause, sql_params, next_param_index)
+    )
     page_rows: Final[Sequence[_SessionPageRow]] = (
         ()
         if page_limit <= 0
@@ -3107,11 +3117,9 @@ 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
-    total_records, total_is_capped = (
-        (offset + len(page_rows), False)
-        if page_ends_the_list
-        else await _count_grouped_sessions(prisma_client, where_clause, sql_params, next_param_index)
-    )
+    if page_ends_the_list:
+        count_task.cancel()
+    total_records, total_is_capped = (offset + len(page_rows), False) if page_ends_the_list else await count_task
 
     session_keys: Final = tuple((row["session_key"], row["api_key"]) for row in visible_rows)
     data: Final[list[dict[str, object]]] = (  # mutable-ok: _build_ui_spend_logs_response writes onto each row

schema.prisma

@@ -684,6 +684,7 @@ model LiteLLM_SpendLogs {
   @@index([session_id])
   @@index([litellm_call_id])
   @@index([api_key, startTime])
+  @@index([startTime, session_id, request_id, api_key, status], map: "LiteLLM_SpendLogs_session_window_idx")
 }
 
 model LiteLLM_BudgetWindowSpend {

tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py

@@ -7297,6 +7297,10 @@ async def test_ui_view_spend_logs_group_by_session_first_page(client, monkeypatc
     reps = [
         _session_representative_row("req-solo", None),
         _session_representative_row("req-1", "sess-1"),
+        # The plain-column predicate admits the cross product of session ids and
+        # api keys, so the endpoint must drop rows whose (session_key, api_key)
+        # pair was never requested.
+        _session_representative_row("req-cross-product", "sess-unrequested"),
     ]
     mock_prisma = _session_grouped_mock_prisma(page_rows, 3, reps)
     monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)
@@ -7346,6 +7350,12 @@ async def test_ui_view_spend_logs_group_by_session_first_page(client, monkeypatc
             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 "session_id = ANY(" in rep_call[0] and "OR request_id = ANY(" in rep_call[0], (
+            "the representative lookup must filter on the plain session_id/request_id columns so the "
+            "(startTime, session_id, request_id, api_key, status) index applies; a row-tuple IN on the "
+            "COALESCE group key forces a full window scan (the 40s+ logs page)"
+        )
+        assert f"({SESSION_GROUP_KEY_SQL}) IN" not in rep_call[0]
         assert rep_call[-2] == ["sess-1", "req-solo"]
         assert rep_call[-1] == ["hashed-key", "hashed-key"]
     finally:
@@ -7635,8 +7645,9 @@ async def test_ui_view_spend_logs_group_by_session_offset_for_non_starttime_sort
         assert "next_session_cursor" not in data
         emitted_sql = [call.args[0] for call in mock_prisma.db.query_raw.await_args_list]
         assert "HAVING" not in " ".join(emitted_sql)
-        assert f"DISTINCT ON ({SESSION_GROUP_KEY_SQL})" in emitted_sql[1]
-        assert "OFFSET" in emitted_sql[1]
+        page_sql = [sql for sql in emitted_sql if f"DISTINCT ON ({SESSION_GROUP_KEY_SQL})" in sql]
+        assert len(page_sql) == 1
+        assert "OFFSET" in page_sql[0]
     finally:
         app.dependency_overrides.pop(ps.user_api_key_auth, None)
 

tests/test_litellm/proxy/spend_tracking/test_spend_query_optimization.py

@@ -203,20 +203,31 @@ async def test_spend_logs_ui_wraps_params_in_at_time_zone_utc(monkeypatch):
 
 def _make_ui_spend_logs_mock(count_total, page_rows):
     """
-    Build a prisma mock whose first `query_raw` (the bounded count) returns
-    `count_total` and whose second `query_raw` (the page data) returns
-    `page_rows`.
+    Build a prisma mock whose bounded-count `query_raw` returns `count_total`
+    and whose page-data `query_raw` returns `page_rows`, dispatching on the
+    SQL so the two queries may run in either order (they run concurrently).
     """
+
+    async def mock_query_raw(sql_query, *params):
+        if "COUNT(*) AS total_count" in sql_query:
+            return [{"total_count": count_total}]
+        return page_rows
+
     mock_prisma = MagicMock()
     mock_prisma.db = MagicMock()
-    mock_prisma.db.query_raw = AsyncMock(
-        side_effect=[[{"total_count": count_total}], page_rows]
-    )
+    mock_prisma.db.query_raw = AsyncMock(side_effect=mock_query_raw)
     mock_prisma.db.litellm_spendlogs = MagicMock()
     mock_prisma.db.litellm_spendlogs.count = AsyncMock(return_value=0)
     return mock_prisma
 
 
+def _query_raw_call_matching(mock_prisma, sql_fragment):
+    """Return the single query_raw call whose SQL contains ``sql_fragment``."""
+    matches = [call for call in mock_prisma.db.query_raw.call_args_list if sql_fragment in call[0][0]]
+    assert len(matches) == 1, f"expected exactly one query containing {sql_fragment!r}, got {len(matches)}"
+    return matches[0]
+
+
 @pytest.mark.asyncio
 async def test_spend_logs_ui_uses_bounded_count_not_full_scan(monkeypatch):
     """
@@ -260,21 +271,17 @@ async def test_spend_logs_ui_uses_bounded_count_not_full_scan(monkeypatch):
 
     mock_prisma.db.litellm_spendlogs.count.assert_not_called()
 
-    count_call = mock_prisma.db.query_raw.call_args_list[0]
+    count_call = _query_raw_call_matching(mock_prisma, "COUNT(*) AS total_count")
     count_sql = count_call[0][0]
     assert "COUNT(*) OVER ()" not in count_sql
     assert "LIMIT" in count_sql and "FROM (" in count_sql, (
-        "the total must come from a bounded subquery count, not a full-window "
-        f"scan. SQL was:\n{count_sql}"
-    )
-    assert count_call[0][-1] == SPEND_LOGS_PAGINATION_COUNT_CAP + 1, (
-        "the bounded count must probe at most cap+1 rows"
+        f"the total must come from a bounded subquery count, not a full-window scan. SQL was:\n{count_sql}"
     )
+    assert count_call[0][-1] == SPEND_LOGS_PAGINATION_COUNT_CAP + 1, "the bounded count must probe at most cap+1 rows"
 
-    page_sql = mock_prisma.db.query_raw.call_args_list[1][0][0]
+    page_sql = _query_raw_call_matching(mock_prisma, "ORDER BY")[0][0]
     assert "COUNT(*) OVER ()" not in page_sql, (
-        "the page query must not carry a window count that forces a full-window "
-        f"scan. SQL was:\n{page_sql}"
+        f"the page query must not carry a window count that forces a full-window scan. SQL was:\n{page_sql}"
     )
     assert "GROUP BY" not in count_sql and "DISTINCT ON" not in page_sql, (
         "without group_by_session the endpoint must keep raw per-call pagination"
@@ -302,9 +309,7 @@ async def test_spend_logs_ui_caps_total_for_large_result_sets(monkeypatch):
     )
 
     page_rows = [{"request_id": "req-1", "metadata": "{}", "session_id": None}]
-    mock_prisma = _make_ui_spend_logs_mock(
-        count_total=SPEND_LOGS_PAGINATION_COUNT_CAP + 1, page_rows=page_rows
-    )
+    mock_prisma = _make_ui_spend_logs_mock(count_total=SPEND_LOGS_PAGINATION_COUNT_CAP + 1, page_rows=page_rows)
     monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)
 
     auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin")
@@ -343,13 +348,7 @@ async def test_spend_logs_ui_empty_page_reports_zero_total(monkeypatch):
         ui_view_spend_logs,
     )
 
-    # First query_raw call is the bounded count (0 matches), second is the empty
-    # page.
-    mock_prisma = MagicMock()
-    mock_prisma.db = MagicMock()
-    mock_prisma.db.query_raw = AsyncMock(side_effect=[[{"total_count": 0}], []])
-    mock_prisma.db.litellm_spendlogs = MagicMock()
-    mock_prisma.db.litellm_spendlogs.count = AsyncMock(return_value=0)
+    mock_prisma = _make_ui_spend_logs_mock(count_total=0, page_rows=[])
 
     monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)
 
@@ -391,13 +390,7 @@ async def test_spend_logs_ui_out_of_range_page_keeps_total(monkeypatch):
         ui_view_spend_logs,
     )
 
-    # First query_raw call is the bounded count (7 matches), second is the
-    # out-of-range page (empty).
-    mock_prisma = MagicMock()
-    mock_prisma.db = MagicMock()
-    mock_prisma.db.query_raw = AsyncMock(side_effect=[[{"total_count": 7}], []])
-    mock_prisma.db.litellm_spendlogs = MagicMock()
-    mock_prisma.db.litellm_spendlogs.count = AsyncMock(return_value=0)
+    mock_prisma = _make_ui_spend_logs_mock(count_total=7, page_rows=[])
 
     monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)
 
@@ -634,10 +627,10 @@ async def test_spend_logs_ui_group_by_session_offset_pages_for_other_sorts(monke
     )
 
     group_key = "COALESCE(NULLIF(session_id, ''), request_id), api_key"
-    count_sql = mock_prisma.db.query_raw.call_args_list[0][0][0]
+    count_sql = _query_raw_call_matching(mock_prisma, "COUNT(*) AS total_count")[0][0]
     assert f"GROUP BY {group_key}" in count_sql
 
-    page_call = mock_prisma.db.query_raw.call_args_list[1][0]
+    page_call = _query_raw_call_matching(mock_prisma, f"DISTINCT ON ({group_key})")[0]
     page_sql = page_call[0]
     assert f"DISTINCT ON ({group_key})" in page_sql
     assert "ORDER BY spend DESC" in page_sql
@@ -682,7 +675,7 @@ async def test_spend_logs_ui_request_id_lookup_with_grouping_returns_exact_row(m
         group_by_session=True,
     )
 
-    page_call = mock_prisma.db.query_raw.call_args_list[1]
+    page_call = _query_raw_call_matching(mock_prisma, "DISTINCT ON")
     assert "request_id = $" in page_call[0][0], "the request_id equality filter must survive grouping"
     assert "req-deep-link" in page_call[0]
     assert [row["request_id"] for row in response["data"]] == ["req-deep-link"]