mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
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:
parent
857b9ad7d3
commit
4295bf823a
2 changed files with 89 additions and 14 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue