fix(proxy): bound the cursorless grouped log page offset

A page starting at or past SPEND_LOGS_PAGINATION_COUNT_CAP lies outside the
total the client is given, so it now returns no rows without running the page
query and the grouped top-N sort bound stays capped.

Rewrites the offset test to page a fake session store instead of asserting on
the generated SQL, and covers the last page inside the cap next to the first
page past it.

Claude-Session: https://claude.ai/code/session_01ESi9JwaXDww1vP3Qsrr4Mz
This commit is contained in:
ryan-crabbe-berri 2026-09-11 14:16:43 -07:00
parent 857b9ad7d3
commit 4295bf823a
2 changed files with 89 additions and 14 deletions

View file

@ -2947,7 +2947,9 @@ async def _ui_session_grouped_spend_logs(
page depth does not degrade the query plan. A request for ``page > 1``
without a cursor (the UI jumping straight to the last page, or back to a
page it never walked through) falls back to ``OFFSET (page - 1) *
page_size``, bounded by the capped total. Each session is represented
page_size``; a page starting at or past ``SPEND_LOGS_PAGINATION_COUNT_CAP``
lies outside the capped total the client is given, so it returns no rows
without running the query and the sort bound stays capped. Each session is represented
by its newest non-MCP row, enriched by ``_build_ui_spend_logs_response``
exactly like the flat listing, and the response carries
``next_session_cursor`` / ``has_more`` while ``total`` counts sessions
@ -2966,7 +2968,9 @@ async def _ui_session_grouped_spend_logs(
)
cursor_params: Final[tuple[object, ...]] = cursor if cursor else ()
limit_index: Final = next_param_index + len(cursor_params)
offset_params: Final[tuple[int, ...]] = ((page - 1) * page_size,) if cursor is None and page > 1 else ()
offset: Final = (page - 1) * page_size if cursor is None else 0
beyond_capped_window: Final = offset >= SPEND_LOGS_PAGINATION_COUNT_CAP
offset_params: Final[tuple[int, ...]] = (offset,) if offset and not beyond_capped_window else ()
offset_clause: Final = f"OFFSET ${limit_index + 1}" if offset_params else ""
page_query: Final = f"""
@ -2980,8 +2984,10 @@ async def _ui_session_grouped_spend_logs(
ORDER BY MAX("startTime") {direction}, {_SESSION_KEY_EXPR} {direction}, api_key {direction}
LIMIT ${limit_index} {offset_clause}
"""
page_rows: Final[Sequence[_SessionPageRow]] = await _query_raw(
prisma_client, page_query, *sql_params, *cursor_params, page_size + 1, *offset_params
page_rows: Final[Sequence[_SessionPageRow]] = (
()
if beyond_capped_window
else await _query_raw(prisma_client, page_query, *sql_params, *cursor_params, page_size + 1, *offset_params)
)
has_more: Final = len(page_rows) > page_size

View file

@ -6629,6 +6629,30 @@ def _session_page_row(session_key, last_activity):
return {"session_key": session_key, "api_key": "hashed-key", "last_activity": last_activity}
def _session_grouped_paginating_prisma(sessions):
"""Mock prisma serving the grouped page query out of ``sessions``, honoring the LIMIT and OFFSET it asks for."""
async def mock_query_raw(sql_query, *params):
if "COUNT(*) AS total_count" in sql_query:
return [{"total_count": min(len(sessions), params[-1])}]
if "DISTINCT ON" in sql_query:
return [_session_representative_row(f"req-{session_key}", session_key) for session_key in params[-2]]
if "COALESCE(SUM(spend)" in sql_query:
return []
bounds = re.search(r"LIMIT \$(\d+)(?: OFFSET \$(\d+))?", sql_query)
limit = params[int(bounds.group(1)) - 1]
offset = params[int(bounds.group(2)) - 1] if bounds.group(2) else 0
return [
_session_page_row(session_key, last_activity)
for session_key, last_activity in sessions[offset : offset + limit]
]
mock_prisma = MagicMock()
mock_prisma.db = MagicMock()
mock_prisma.db.query_raw = AsyncMock(side_effect=mock_query_raw)
return mock_prisma
@pytest.mark.asyncio
async def test_ui_view_spend_logs_group_by_session_first_page(client, monkeypatch):
"""One row per (session, api_key), session-count total, and a keyset cursor for the next page."""
@ -6744,10 +6768,9 @@ async def test_ui_view_spend_logs_group_by_session_cursor_page(client, monkeypat
@pytest.mark.asyncio
async def test_ui_view_spend_logs_group_by_session_jumps_to_page_without_cursor(client, monkeypatch):
"""page > 1 with no session_cursor (the UI's last-page jump) skips (page - 1) * page_size sessions by OFFSET."""
page_rows = [_session_page_row("sess-3", "2026-08-29 06:00:00")]
reps = [_session_representative_row("req-3", "sess-3")]
mock_prisma = _session_grouped_mock_prisma(page_rows, 60, reps)
"""page > 1 with no session_cursor (the UI's last-page jump) serves the sessions that page starts at."""
sessions = tuple((f"sess-{index:02d}", f"2026-08-29 10:{59 - index:02d}:00") for index in range(60))
mock_prisma = _session_grouped_paginating_prisma(sessions)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)
monkeypatch.setattr(
"litellm.proxy.spend_tracking.spend_management_endpoints._is_admin_view_safe",
@ -6772,14 +6795,60 @@ async def test_ui_view_spend_logs_group_by_session_jumps_to_page_without_cursor(
assert response.status_code == 200, response.text
data = response.json()
assert data["page"] == 3
assert data["total"] == 60
assert data["has_more"] is False
assert [row["request_id"] for row in data["data"]] == ["req-3"]
assert data["next_session_cursor"] is None
assert [row["request_id"] for row in data["data"]] == [f"req-sess-{index:02d}" for index in range(50, 60)]
finally:
app.dependency_overrides.pop(ps.user_api_key_auth, None)
page_query_call = mock_prisma.db.query_raw.await_args_list[0]
page_query_sql = page_query_call.args[0]
assert "HAVING" not in page_query_sql
assert "OFFSET" in page_query_sql
assert page_query_call.args[-2:] == (26, 50), "LIMIT page_size + 1 then OFFSET (page - 1) * page_size"
@pytest.mark.asyncio
async def test_ui_view_spend_logs_group_by_session_page_past_count_cap_is_empty(client, monkeypatch):
"""The last page inside the capped total still lists sessions; the page after it is empty and costs no query."""
cap = spend_management_endpoints.SPEND_LOGS_PAGINATION_COUNT_CAP
sessions = tuple((f"sess-{index:06d}", "2026-08-29 10:00:00") for index in range(cap + 50))
mock_prisma = _session_grouped_paginating_prisma(sessions)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)
monkeypatch.setattr(
"litellm.proxy.spend_tracking.spend_management_endpoints._is_admin_view_safe",
lambda user_api_key_dict: True,
)
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user"
)
try:
start_date, end_date = _default_date_range()
params = {
"start_date": start_date,
"end_date": end_date,
"group_by_session": "true",
"page_size": 25,
}
last_page = client.get(
"/spend/logs/ui",
params={**params, "page": cap // 25},
headers={"Authorization": "Bearer sk-test"},
)
assert last_page.status_code == 200, last_page.text
last_page_data = last_page.json()
assert last_page_data["total"] == cap
assert last_page_data["total_is_capped"] is True
assert last_page_data["data"][0]["request_id"] == f"req-sess-{cap - 25:06d}"
assert len(last_page_data["data"]) == 25
mock_prisma.db.query_raw.reset_mock()
past_cap = client.get(
"/spend/logs/ui",
params={**params, "page": cap // 25 + 1},
headers={"Authorization": "Bearer sk-test"},
)
assert past_cap.status_code == 200, past_cap.text
past_cap_data = past_cap.json()
assert past_cap_data["data"] == []
assert past_cap_data["has_more"] is False
assert past_cap_data["total"] == cap
assert mock_prisma.db.query_raw.await_count == 1, "only the bounded count query runs past the capped window"
finally:
app.dependency_overrides.pop(ps.user_api_key_auth, None)