From 3001970e366e2ed39a3c0931b22acdc6b44f1f50 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Sat, 19 Sep 2026 14:02:16 -0400 Subject: [PATCH] fix(proxy): type the concurrent flat-page count as a task and unpin SQL emission order in the offset test Co-Authored-By: Claude Fable 5 --- .../proxy/spend_tracking/spend_management_endpoints.py | 10 +++++----- .../spend_tracking/test_spend_management_endpoints.py | 5 +++-- 2 files changed, 8 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 63e5cc4672a..0da193fd904 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -3,7 +3,7 @@ import asyncio import collections import json import os -from collections.abc import Awaitable, Mapping, Sequence +from collections.abc import Mapping, Sequence from dataclasses import dataclass from datetime import date, datetime, timedelta, timezone from itertools import groupby @@ -2880,11 +2880,11 @@ async def ui_view_spend_logs( LIMIT ${p} OFFSET ${p + 1} """ ) - count_task: Final[Awaitable[Sequence[_SpendLogsCountRow] | None]] = _query_raw_or_none( - prisma_client, count_query, *sql_params, SPEND_LOGS_PAGINATION_COUNT_CAP + 1 + 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_task: Final = prisma_client.db.query_raw(sql_query, *sql_params, page_size, skip) - count_rows, data = await asyncio.gather(count_task, data_task) + 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 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 344dcabe8a8..445ea205881 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 @@ -7717,8 +7717,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)