From cf34162b4493ef04d62110d2300b47cbebc18a11 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Fri, 26 Jun 2026 10:03:31 +0300 Subject: [PATCH] fix(spend): fold logs-tab total into the page query to avoid a separate COUNT(*) The spend-logs UI list endpoint (/spend/logs/ui, /spend/logs/v2) ran a standalone SELECT COUNT(*) before the page query to compute total_pages. On sharded engines like YugabyteDB a COUNT(*) is a distributed RPC that contacts every tablet leader and aggregates partial results regardless of row count, so it hits the distributed RPC timeout and the logs tab 500s even on a one-minute window with a couple of rows. The startTime range cannot prune tablets because rows hash to tablets on request_id, not startTime. Fold the count into the same scan as the page data with COUNT(*) OVER () and read total off the returned rows, dropping the helper column before serialisation. One distributed scan per page load instead of two; the response shape is unchanged. An empty page carries no count row, in which case the total is zero. Resolves LIT-4027 --- .../spend_management_endpoints.py | 31 +++- .../test_spend_management_endpoints.py | 124 ++++++++++++-- .../test_spend_query_optimization.py | 151 ++++++++++++++++++ 3 files changed, 284 insertions(+), 22 deletions(-) diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 29fc6d7c30fa..400b1a7111f4 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -2010,11 +2010,6 @@ def parse_date(date_str: str) -> datetime: order_column = sort_by order_direction = (sort_order or "desc").lower() - # Get total count of records - total_records = await SpendLogsRepository(prisma_client).table.count( - where=where_conditions, - ) - # Build raw SQL to fetch paginated data WITHOUT heavy columns # (messages, response, proxy_server_request can be hundreds of KB per row). # These are only needed in the detail endpoint /spend/logs/ui/{request_id}. @@ -2128,7 +2123,8 @@ def parse_date(date_str: str) -> datetime: cache_hit, cache_key, request_tags, team_id, organization_id, end_user, requester_ip_address, session_id, status, mcp_namespaced_tool_name, agent_id, - COALESCE(request_duration_ms, (EXTRACT(EPOCH FROM ("endTime" - "startTime")) * 1000)::INTEGER) AS request_duration_ms + COALESCE(request_duration_ms, (EXTRACT(EPOCH FROM ("endTime" - "startTime")) * 1000)::INTEGER) AS request_duration_ms, + COUNT(*) OVER () AS total_count FROM "LiteLLM_SpendLogs" WHERE {" AND ".join(sql_conditions)} ORDER BY {_order_expr} {_sql_dir}{_nulls_clause} @@ -2138,13 +2134,34 @@ def parse_date(date_str: str) -> datetime: data = await prisma_client.db.query_raw(sql_query, *sql_params) + # `COUNT(*) OVER ()` folds the total-match count into the same scan as the + # page data; a standalone `COUNT(*)` is a distributed RPC on sharded + # engines like YugabyteDB that contacts every tablet and times out + # regardless of row count (LIT-4027). The hot path (page 1 and in-range + # pages) always carries the count on its rows, so the count round trip is + # gone there. Only an out-of-range page overshoots the last row and comes + # back empty; fall back to a direct count there so total/total_pages stay + # accurate rather than collapsing to zero. + if data: + total_records = int(data[0]["total_count"]) + elif page > 1: + total_records = int( + await SpendLogsRepository(prisma_client).table.count( + where=where_conditions, + ) + ) + else: + total_records = 0 + # query_raw returns the JSONB `metadata` column as a string (the Prisma # serialiser bypasses the model-layer JSON hydration we get on the ORM # path). The UI reads `metadata.status` / `metadata.error_information` # as object fields, so failure rows looked like successes (#29674). - # Re-hydrate to dict here. + # Re-hydrate to dict here. Also drop the window-function `total_count` + # helper column so it does not leak into the serialised rows. for row in data: if isinstance(row, dict): + row.pop("total_count", None) md = row.get("metadata") if isinstance(md, str): try: 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 6d0b4509b452..4ddca2d4a9da 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 @@ -2,6 +2,7 @@ import datetime import json import os +import re import sys from datetime import timezone @@ -60,31 +61,115 @@ def _filter_logs_by_date_range(logs, where): return filtered +def _reconstruct_ui_where_from_sql(sql_query, params): + """ + Rebuild the Prisma-style ``where`` dict the filter_fns below expect from the + raw SQL + params the endpoint emits. + + ``ui_view_spend_logs`` folds the total into the page query via + ``COUNT(*) OVER ()`` and no longer issues a separate ``count(where=...)`` + call, so the mock derives the active filter from the one query it sees + instead of from the (now absent) count call. + """ + where: dict = {} + clause = re.search(r"WHERE (.*) ORDER BY", sql_query, re.DOTALL) + if clause is None: + return where + + def _iso(value): + return value.isoformat() if hasattr(value, "isoformat") else str(value) + + eq_cols = { + "team_id": "team_id", + '"user"': "user", + "api_key": "api_key", + "request_id": "request_id", + "model": "model", + "model_id": "model_id", + "model_group": "model_group", + "end_user": "end_user", + } + date_bounds: dict = {} + metadata_conds: list = [] + for cond in (c.strip() for c in clause.group(1).split(" AND ")): + gte = re.search(r'"startTime" >= \(\$(\d+)', cond) + lte = re.search(r'"startTime" <= \(\$(\d+)', cond) + alias = re.search(r"user_api_key_alias' LIKE \$(\d+)", cond) + code = re.search(r"error_code' = \$(\d+)", cond) + msg = re.search(r"error_message' LIKE \$(\d+)", cond) + status = re.fullmatch(r"status = \$(\d+)", cond) + if gte: + date_bounds["gte"] = _iso(params[int(gte.group(1)) - 1]) + elif lte: + date_bounds["lte"] = _iso(params[int(lte.group(1)) - 1]) + elif "OR team_id = ANY" in cond: + where["OR"] = where.get("OR", []) + [{"multi_team": True}] + elif "status = 'success'" in cond: + where["OR"] = where.get("OR", []) + [{"status": "success"}] + elif status: + where["status"] = {"equals": params[int(status.group(1)) - 1]} + elif alias: + metadata_conds.append( + { + "path": ["user_api_key_alias"], + "string_contains": str(params[int(alias.group(1)) - 1]).strip("%"), + } + ) + elif code: + metadata_conds.append( + { + "path": ["error_information", "error_code"], + "equals": params[int(code.group(1)) - 1], + } + ) + elif msg: + metadata_conds.append( + { + "path": ["error_information", "error_message"], + "string_contains": str(params[int(msg.group(1)) - 1]).strip("%"), + } + ) + else: + for sql_col, key in eq_cols.items(): + eq = re.fullmatch(rf"{re.escape(sql_col)} = \$(\d+)", cond) + if eq: + where[key] = params[int(eq.group(1)) - 1] + break + + if date_bounds: + where["startTime"] = date_bounds + if len(metadata_conds) == 1: + where["metadata"] = metadata_conds[0] + elif len(metadata_conds) > 1: + where["AND"] = [{"metadata": cond} for cond in metadata_conds] + return where + + def make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_fn, team_lookup_fn=None): """ Create a MockPrismaClient for /spend/logs/ui endpoint tests. Args: mock_spend_logs: List of mock spend log dicts. - filter_fn: Callable[[dict], list] - receives where_conditions from count(), - returns the filtered list of logs for that query. + filter_fn: Callable[[dict], list] - receives the reconstructed + where_conditions, returns the filtered list of logs. team_lookup_fn: Optional async callable for team RBAC (find_unique). If provided, adds litellm_teamtable to db. """ - filtered_holder = [] class MockDB: async def count(self, *args, **kwargs): - where = kwargs.get("where", {}) - filtered = filter_fn(where) - filtered_holder.clear() - filtered_holder.extend(filtered) - return len(filtered) + return len(filter_fn(kwargs.get("where", {}))) async def query_raw(self, sql_query, *params): + filtered = filter_fn(_reconstruct_ui_where_from_sql(sql_query, params)) page_size = params[-2] if len(params) >= 2 else 50 skip = params[-1] if len(params) >= 1 else 0 - return filtered_holder[skip : skip + page_size] + total = len(filtered) + return [ + {**row, "total_count": total} + for row in filtered[skip : skip + page_size] + ] class MockPrismaClient: def __init__(self): @@ -608,7 +693,10 @@ async def mock_query_raw(sql_query, *params): sorted_logs = _sort_logs(base_logs, order) page_size = params[-2] if len(params) >= 2 else 50 skip = params[-1] if len(params) >= 1 else 0 - return sorted_logs[skip : skip + page_size] + return [ + {**row, "total_count": len(base_logs)} + for row in sorted_logs[skip : skip + page_size] + ] class MockPrismaClient: def __init__(self): @@ -748,7 +836,10 @@ async def mock_query_raw(sql_query, *params): ) page_size = params[-2] if len(params) >= 2 else 50 skip = params[-1] if len(params) >= 1 else 0 - return sorted_logs[skip : skip + page_size] + return [ + {**row, "total_count": len(base_logs)} + for row in sorted_logs[skip : skip + page_size] + ] class MockPrismaClient: def __init__(self): @@ -846,7 +937,10 @@ async def mock_query_raw(sql_query, *params): ) page_size = params[-2] if len(params) >= 2 else 50 skip = params[-1] if len(params) >= 1 else 0 - return sorted_logs[skip : skip + page_size] + return [ + {**row, "total_count": len(base_logs)} + for row in sorted_logs[skip : skip + page_size] + ] class MockPrismaClient: def __init__(self): @@ -957,7 +1051,7 @@ async def mock_query_raw(sql_query, *params): page_size = params[-2] if len(params) >= 2 else 50 skip = params[-1] if len(params) >= 1 else 0 return [ - {k: v for k, v in row.items() if k != "_ttft_ms"} + {**{k: v for k, v in row.items() if k != "_ttft_ms"}, "total_count": len(base_logs)} for row in sorted_logs[skip : skip + page_size] ] @@ -3668,7 +3762,7 @@ async def mock_count(*args, **kwargs): return 1 async def mock_query_raw(sql_query, *params): - return [raw_row] + return [{**raw_row, "total_count": 1}] class MockPrismaClient: def __init__(self): @@ -3754,7 +3848,7 @@ async def mock_count(*args, **kwargs): return 1 async def mock_query_raw(sql_query, *params): - return [raw_row] + return [{**raw_row, "total_count": 1}] class MockPrismaClient: def __init__(self): 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 fbac71e63720..5b793adbbb46 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 @@ -222,3 +222,154 @@ async def test_spend_logs_ui_wraps_params_in_at_time_zone_utc(monkeypatch): "/spend/logs/ui must wrap both `startTime` bounds with " f"`AT TIME ZONE 'UTC'`. SQL was:\n{sql}" ) + + +@pytest.mark.asyncio +async def test_spend_logs_ui_folds_count_into_window_function(monkeypatch): + """ + /spend/logs/ui must not issue a separate `COUNT(*)` round trip to compute + the total. On sharded engines like YugabyteDB a standalone `COUNT(*)` is a + distributed RPC that contacts every tablet and times out regardless of row + count, so the logs tab 500s (LIT-4027). The total is folded into the page + query via `COUNT(*) OVER ()` and read off the returned rows instead. + """ + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.spend_tracking.spend_management_endpoints import ( + ui_view_spend_logs, + ) + + rows = [ + {"request_id": "req-1", "metadata": "{}", "session_id": None, "total_count": 137}, + {"request_id": "req-2", "metadata": "{}", "session_id": None, "total_count": 137}, + ] + + mock_prisma = MagicMock() + mock_prisma.db = MagicMock() + mock_prisma.db.query_raw = AsyncMock(return_value=rows) + mock_prisma.db.litellm_spendlogs = MagicMock() + mock_prisma.db.litellm_spendlogs.count = AsyncMock(return_value=0) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin") + mock_request = MagicMock() + mock_request.url.path = "/spend/logs/ui" + + response = await ui_view_spend_logs( + request=mock_request, + api_key=None, + user_id=None, + request_id=None, + start_date="2026-02-16 00:00:00", + end_date="2026-02-16 23:59:59", + page=1, + page_size=50, + sort_by="startTime", + sort_order="desc", + user_api_key_dict=auth, + ) + + mock_prisma.db.litellm_spendlogs.count.assert_not_called() + + sql = mock_prisma.db.query_raw.call_args[0][0] + assert "COUNT(*) OVER ()" in sql, ( + "the page query must carry a window-function count so a separate " + f"distributed COUNT(*) is avoided. SQL was:\n{sql}" + ) + + assert response["total"] == 137 + assert response["total_pages"] == (137 + 50 - 1) // 50 + + for row in response["data"]: + assert "total_count" not in row, ( + "the window-function helper column must be stripped before " + "serialising rows" + ) + + +@pytest.mark.asyncio +async def test_spend_logs_ui_empty_page_reports_zero_total(monkeypatch): + """ + When a page matches no rows the window-function count row is absent, so the + total must fall back to zero without issuing a separate `COUNT(*)`. + """ + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.spend_tracking.spend_management_endpoints import ( + ui_view_spend_logs, + ) + + mock_prisma = MagicMock() + mock_prisma.db = MagicMock() + mock_prisma.db.query_raw = AsyncMock(return_value=[]) + mock_prisma.db.litellm_spendlogs = MagicMock() + mock_prisma.db.litellm_spendlogs.count = AsyncMock(return_value=0) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin") + mock_request = MagicMock() + mock_request.url.path = "/spend/logs/ui" + + response = await ui_view_spend_logs( + request=mock_request, + api_key=None, + user_id=None, + request_id=None, + start_date="2026-02-16 00:00:00", + end_date="2026-02-16 23:59:59", + page=1, + page_size=50, + sort_by="startTime", + sort_order="desc", + user_api_key_dict=auth, + ) + + mock_prisma.db.litellm_spendlogs.count.assert_not_called() + assert response["total"] == 0 + assert response["total_pages"] == 0 + assert response["data"] == [] + + +@pytest.mark.asyncio +async def test_spend_logs_ui_out_of_range_page_falls_back_to_count(monkeypatch): + """ + An out-of-range page (offset past the last matching row) returns no rows, so + the window-function count is unavailable. The total must not collapse to zero + there; it falls back to a direct count so total/total_pages stay accurate. + This fallback only fires off the hot path (page > 1 with an empty result), so + the YugabyteDB timeout the fix removes from page 1 stays removed. + """ + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.spend_tracking.spend_management_endpoints import ( + ui_view_spend_logs, + ) + + mock_prisma = MagicMock() + mock_prisma.db = MagicMock() + mock_prisma.db.query_raw = AsyncMock(return_value=[]) + mock_prisma.db.litellm_spendlogs = MagicMock() + mock_prisma.db.litellm_spendlogs.count = AsyncMock(return_value=7) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin") + mock_request = MagicMock() + mock_request.url.path = "/spend/logs/ui" + + response = await ui_view_spend_logs( + request=mock_request, + api_key=None, + user_id=None, + request_id=None, + start_date="2026-02-16 00:00:00", + end_date="2026-02-16 23:59:59", + page=99, + page_size=2, + sort_by="startTime", + sort_order="desc", + user_api_key_dict=auth, + ) + + mock_prisma.db.litellm_spendlogs.count.assert_called_once() + assert response["total"] == 7 + assert response["total_pages"] == (7 + 2 - 1) // 2