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

[yassin-berriai]

Relevant issues

Resolves LIT-4027

Linear ticket

LIT-4027

Pre-Submission checklist

Please complete all items before asking a LiteLLM maintainer to review your PR

• I have added meaningful tests
• My PR passes all CI/CD checks (e.g., lint, format, unit tests)
• My PR's scope is as isolated as possible; it only solves 1 specific problem
• I have requested a Greptile review by commenting @greptileai and received a Confidence Score of at least 4/5 before requesting a maintainer review

Screenshots / Proof of Fix

The logs tab (/spend/logs/ui, /spend/logs/v2) ran a standalone SELECT COUNT(*) before the page query to compute total_pages. On a sharded engine 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 tab 500s even on a one-minute window with a couple of rows. The startTime range can't prune tablets because rows hash to tablets on request_id, not startTime.

Reproduced against a live proxy on real Postgres with log_statement=all, after three real gpt-4.1-mini chat completions populated LiteLLM_SpendLogs. The customer-visible behavior is the SQL the endpoint emits per page load; on YugabyteDB that standalone count is the statement that times out.

Before the fix, one logs-tab request emits a separate sub-select count statement (parse + bind + execute):

$ curl "http://127.0.0.1:4027/spend/logs/ui?start_date=2026-06-01%2000:00:00&end_date=2026-06-30%2023:59:59&page=1&page_size=2" -H "Authorization: Bearer $KEY"
total       : 3
total_pages : 2
page_size   : 2
rows        : 2
row0 has total_count key: False

# Postgres statement log for that request:
execute s35: SELECT COUNT(*) FROM (SELECT "public"."LiteLLM_SpendLogs"."request_id"
             FROM "public"."LiteLLM_SpendLogs"
             WHERE ("startTime" >= $1 AND "startTime" <= $2) OFFSET $3) AS "sub"

After the fix the standalone count is gone; the count rides along the single data scan via COUNT(*) OVER (), and the HTTP response is identical:

$ curl "http://127.0.0.1:4027/spend/logs/ui?start_date=2026-06-01%2000:00:00&end_date=2026-06-30%2023:59:59&page=1&page_size=2" -H "Authorization: Bearer $KEY"
total       : 3
total_pages : 2
page_size   : 2
rows        : 2
row0 has total_count key: False

# Postgres statement log for that request:
SELECT request_id, call_type, ..., COUNT(*) OVER () AS total_count
FROM "LiteLLM_SpendLogs"
WHERE "startTime" >= ($1::timestamptz AT TIME ZONE 'UTC') AND ...
ORDER BY "startTime" DESC LIMIT $n OFFSET $m

Statement tally across the two runs, isolated by a log marker:

BEFORE standalone-count statements: 3
AFTER  standalone-count statements: 0

The helper total_count column is stripped before serialisation, so it never leaks into the response rows (row0 has total_count key: False above).

Type

🐛 Bug Fix

Changes

litellm/proxy/spend_tracking/spend_management_endpoints.py: in ui_view_spend_logs, drop the SpendLogsRepository(...).table.count(...) round trip and add COUNT(*) OVER () AS total_count to the existing page query. total_records is read off the first returned row and the total_count column is popped from each row before it reaches _build_ui_spend_logs_response. One distributed scan per page load instead of two; the response fields (data, total, page, page_size, total_pages) are unchanged.

The hot path that the ticket reports (the first logs-tab load, page 1, and any in-range page) always returns rows, so the count rides along on those rows and the standalone count is gone. The one case with no count row is an out-of-range page whose offset overshoots the last matching row; there total would otherwise collapse to zero, so that single case falls back to a direct count to keep total/total_pages accurate. That fallback never fires on the reported path, so the YugabyteDB timeout stays removed.

This keeps the scope to the logs-tab list endpoint that the ticket reports. The session drill-down endpoint ui_view_session_spend_logs has the same table.count shape but filters on a single indexed session_id, so it isn't the reported hot path; folding its count the same way is a reasonable follow-up if it ever shows up on YugabyteDB.

Tests in tests/test_litellm/proxy/spend_tracking/test_spend_query_optimization.py assert that the endpoint issues no separate count call on the hot path, that the page SQL carries COUNT(*) OVER (), that total/total_pages are derived from the row count, that the helper column doesn't leak, that a page-1 empty result reports a zero total without a count call, and that an out-of-range page falls back to a direct count so the total stays accurate. They fail on the pre-fix code (the separate count call fires) and pass with the fix. The existing test_spend_management_endpoints.py mocks were updated to the new mechanism (the filter now derives from the single query and each row carries total_count)

[yassin-berriai]

@greptileai

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}

[greptile-apps]

Greptile Summary

Eliminates the standalone COUNT(*) round-trip from ui_view_spend_logs by folding the total into the existing page query via COUNT(*) OVER (). This directly fixes a YugabyteDB tablet-fanout timeout on the logs tab. On the hot path (page 1 and all in-range pages) the total rides along on the returned rows; out-of-range pages (page > 1, empty result) fall back to a direct count so total/total_pages stay accurate rather than collapsing to zero.

• spend_management_endpoints.py: Removes the pre-query table.count() call, appends COUNT(*) OVER () AS total_count to the SELECT, reads total_records from the first row, and pops the helper column before serialisation.
• Tests: Existing mocks updated to include total_count in each row; three new tests cover the hot path (no separate count), an empty page-1 result (no count, returns 0), and an out-of-range page (fallback count returns accurate total).

Confidence Score: 5/5

Safe to merge — the change is narrowly scoped to a single endpoint, the response contract is unchanged, and all edge cases (in-range rows, empty table, out-of-range page) have explicit regression tests.

The window-function substitution is mechanically correct: COUNT(*) OVER () always returns the full-match count on every row of the page result, so the hot path never fires a second round-trip. The out-of-range fallback (page > 1, empty data) reuses the same where_conditions variable that the old standalone count used, so filter parity is preserved. The total_count column is safely popped with a default before serialisation. Test coverage is thorough and the new tests fail on the pre-fix code as described.

No files require special attention.

Important Files Changed

Filename	Overview
litellm/proxy/spend_tracking/spend_management_endpoints.py	Removes the standalone COUNT() round-trip for ui_view_spend_logs and replaces it with a COUNT() OVER () window function folded into the page query. The fallback to a direct count is correctly scoped to page > 1 with an empty result set; page 1 with no rows correctly returns zero without a count call.
tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py	Existing mocks updated to inject total_count into each returned row, matching the new window-function mechanism. The _reconstruct_ui_where_from_sql helper infers the active filter from the SQL text rather than from the now-absent count() call. Coverage is maintained and assertions remain strict on total/total_pages values.
tests/test_litellm/proxy/spend_tracking/test_spend_query_optimization.py	Three new tests added: verifies no separate COUNT(*) call fires on the hot path, verifies page 1 with an empty table returns zero without a count call, and verifies an out-of-range page (page > 1, empty result) falls back to a direct count so total/total_pages stay accurate.

_{Reviews (4): Last reviewed commit: "fix(spend): fold logs-tab total into the..." | Re-trigger Greptile}

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[yassin-berriai]

@greptileai

[yassin-berriai]

Addressed the out-of-range page edge case from the last review. The hot path (page 1 and in-range pages) still issues a single scan with no separate count; only an out-of-range page whose offset overshoots the last row now falls back to a direct count so total/total_pages stay accurate instead of collapsing to zero. Added a regression test for that fallback. @greptileai