fix(spend): fold logs-tab total into the page query to avoid a separate COUNT(*)

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:

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):