diff --git a/tests/integration/spend/test_session_view_page_cost.py b/tests/integration/spend/test_session_view_page_cost.py index 6856e9b04af..b4988466791 100644 --- a/tests/integration/spend/test_session_view_page_cost.py +++ b/tests/integration/spend/test_session_view_page_cost.py @@ -304,8 +304,10 @@ def _small_rows(marker: str) -> tuple[_SeedRow, ...]: _SeedRow(f"{marker}-shared-b", f"{marker}-shared", at(60), api_key=f"{marker}-key-b"), _SeedRow(f"{marker}-tie-1", f"{marker}-tie", at(70)), _SeedRow(f"{marker}-tie-2", f"{marker}-tie", at(70)), + _SeedRow(f"{marker}-mcp-old", f"{marker}-mcp", at(45)), _SeedRow(f"{marker}-mcp-llm", f"{marker}-mcp", at(80)), _SeedRow(f"{marker}-mcp-tool", f"{marker}-mcp", at(90), call_type="call_mcp_tool"), + _SeedRow(f"{marker}-shared", f"{marker}-other", at(95), api_key=f"{marker}-key-b"), ) @@ -346,6 +348,7 @@ def _tie_as_one(request_id: str) -> str: def _newest_first(marker: str) -> list[tuple[str, str]]: return [ + (f"{marker}-shared", f"{marker}-key-b"), (f"{marker}-mcp-llm", "K"), (f"{marker}-tie", "K"), (f"{marker}-shared-b", f"{marker}-key-b"), @@ -360,10 +363,10 @@ def test_session_view_lists_each_session_once(gateway: Gateway, small_seed: str) assert _sessions(page) == _newest_first(small_seed), ( "each session once, represented by its newest non-MCP row: a row without a session_id newer than the rows " - "adopting its request_id joins them, two rows at one startTime are one session, and one session_id under two " - "api_keys is two sessions" + "adopting its request_id joins them, two rows at one startTime are one session, one session_id under two " + "api_keys is two sessions, and a row whose request_id equals another session's key stays in its own session" ) - assert page["total"] == 6 + assert page["total"] == 7 assert page["has_more"] is False @@ -382,9 +385,11 @@ def test_session_view_cursor_walk_counts_sessions_per_api_key(gateway: Gateway, ) walked: Final = [session for page in cursor_pages for session in _sessions(page)] - assert first["total"] == 6, "the capped count must count one session_id under two api_keys as two sessions" - assert walked == _newest_first(small_seed) - assert [page["has_more"] for page in cursor_pages] == [True, True, False] + assert first["total"] == 7, "the capped count must count one session_id under two api_keys as two sessions" + assert walked == _newest_first(small_seed), ( + "a session whose older rows fall below the cursor was already shown and must not come back on a later page" + ) + assert [page["has_more"] for page in cursor_pages] == [True, True, True, False] def test_session_view_ascending_and_offset_pages_match_the_newest_first_order( @@ -394,6 +399,16 @@ def test_session_view_ascending_and_offset_pages_match_the_newest_first_order( offset_page: Final = _small_page(gateway, sort_order="desc", page_size="2", page="2") assert _sessions(ascending) == _newest_first(small_seed)[::-1] - assert ascending["total"] == 6 + assert ascending["total"] == 7 assert _sessions(offset_page) == _newest_first(small_seed)[2:4] assert offset_page["has_more"] is True + + +def test_session_view_filtered_page_only_lists_matching_sessions(gateway: Gateway, small_seed: str) -> None: + page: Final = _small_page(gateway, sort_order="desc", page_size="50", page="1", api_key=f"{small_seed}-key-b") + + assert _sessions(page) == [ + (f"{small_seed}-shared", f"{small_seed}-key-b"), + (f"{small_seed}-shared-b", f"{small_seed}-key-b"), + ] + assert page["total"] == 2 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 9bf0540b18e..22a8c1c58f2 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 @@ -13,7 +13,6 @@ from fastapi.testclient import TestClient import litellm import litellm.proxy.proxy_server as ps -from litellm.proxy.spend_tracking.spend_management_endpoints import _SESSION_HEAD_WALK_ROWS def _default_date_range(): @@ -7212,10 +7211,32 @@ SESSION_KEY_EXPR = "COALESCE(NULLIF(session_id, ''), request_id)" SESSION_GROUP_KEY_SQL = f"{SESSION_KEY_EXPR}, api_key" +def _session_walk_batch(page_rows, rows_read): + """Walk-batch rows for ``page_rows`` (newest first): one first row per session plus the batch's oldest row.""" + first_rows = [ + { + "session_key": row["session_key"], + "api_key": row["api_key"], + "last_activity": row["last_activity"], + "sort_time": row["last_activity"], + "request_id": f"req-{row['session_key']}", + "is_edge": False, + "batch_rows": 0, + } + for row in page_rows + ] + edge = [{**first_rows[-1], "is_edge": True, "batch_rows": rows_read}] if first_rows else [] + return first_rows + edge + + def _session_grouped_mock_prisma(session_page_rows, session_total, representative_rows): """Mock prisma dispatching the raw queries the keyset grouped path emits.""" async def mock_query_raw(sql_query, *params): + if "AS newer_row" in sql_query: + return [] + if "WITH batch AS MATERIALIZED" in sql_query: + return _session_walk_batch(session_page_rows, 0) if "COUNT(*) AS total_count" in sql_query: return [{"total_count": session_total}] if "CROSS JOIN LATERAL" in sql_query or "DISTINCT ON" in sql_query: @@ -7268,6 +7289,10 @@ def _session_grouped_paginating_prisma(sessions, counted_total=None): """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 "AS newer_row" in sql_query: + return [] + if "WITH batch AS MATERIALIZED" in sql_query: + return _session_walk_batch([_session_page_row(key, activity) for key, activity in sessions], 0) if "COUNT(*) AS total_count" in sql_query: return [{"total_count": min(len(sessions) if counted_total is None else counted_total, params[-1])}] if "CROSS JOIN LATERAL" in sql_query or "DISTINCT ON" in sql_query: @@ -7335,7 +7360,6 @@ async def test_ui_view_spend_logs_group_by_session_first_page(client, monkeypatc page_query_sql = emitted[0][0] assert "OFFSET" not in page_query_sql assert "HAVING" not in page_query_sql - assert emitted[0][-1] == 3, "page query fetches page_size + 1 sessions to detect has_more" count_sql = emitted[1][0] assert "LIMIT" in count_sql and "FROM (" in count_sql, "the grouped count must stay bounded" @@ -7394,20 +7418,18 @@ async def test_ui_view_spend_logs_group_by_session_cursor_page(client, monkeypat app.dependency_overrides.pop(ps.user_api_key_auth, None) -def _session_walk_mock_prisma(walked_rows, walk_window_rows, grouped_rows): - """Mock prisma serving the head walk, its window-size check and the grouped fallback page separately.""" +def _session_walk_mock_prisma(walked_rows, window_rows_per_batch, grouped_rows): + """Mock prisma whose walk batches keep returning ``walked_rows`` and report ``window_rows_per_batch`` rows read.""" async def mock_query_raw(sql_query, *params): - if "AS walked_rows" in sql_query: - return [{"total_count": walk_window_rows}] + if "WITH batch AS MATERIALIZED" in sql_query: + return _session_walk_batch(walked_rows, min(window_rows_per_batch, params[-1])) if "COUNT(*) AS total_count" in sql_query: return [{"total_count": 3}] if "CROSS JOIN LATERAL" 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 [] - if "LEFT JOIN LATERAL" in sql_query: - return walked_rows return grouped_rows mock_prisma = MagicMock() @@ -7418,21 +7440,19 @@ def _session_walk_mock_prisma(walked_rows, walk_window_rows, grouped_rows): @pytest.mark.asyncio @pytest.mark.parametrize( - ("walk_window_rows", "expected_request_ids", "expected_has_more"), + ("window_rows_per_batch", "expected_request_ids", "expected_has_more"), [ - pytest.param( - _SESSION_HEAD_WALK_ROWS, ["req-grouped-1", "req-grouped-2"], True, id="walk-cut-short-groups-the-window" - ), - pytest.param(_SESSION_HEAD_WALK_ROWS - 1, ["req-walked-1"], False, id="walk-covered-the-window"), + pytest.param(10**9, ["req-grouped-1", "req-grouped-2"], True, id="walk-that-cannot-settle-groups-the-window"), + pytest.param(10, ["req-walked-1"], False, id="walk-that-read-the-whole-window-is-trusted"), ], ) -async def test_ui_view_spend_logs_group_by_session_falls_back_when_the_walk_is_cut_short( - client, monkeypatch, walk_window_rows, expected_request_ids, expected_has_more +async def test_ui_view_spend_logs_group_by_session_falls_back_when_the_walk_cannot_settle( + client, monkeypatch, window_rows_per_batch, expected_request_ids, expected_has_more ): - """A head walk that found fewer sessions than the page needs is only trusted when it read the whole window.""" + """A walk that keeps finding too few sessions serves the grouped page; one that ran out of rows is exact.""" mock_prisma = _session_walk_mock_prisma( walked_rows=[_session_page_row("walked-1", "2026-08-29 10:00:00")], - walk_window_rows=walk_window_rows, + window_rows_per_batch=window_rows_per_batch, grouped_rows=[ _session_page_row("grouped-1", "2026-08-29 10:00:00"), _session_page_row("grouped-2", "2026-08-29 09:00:00"), @@ -7462,6 +7482,64 @@ async def test_ui_view_spend_logs_group_by_session_falls_back_when_the_walk_is_c app.dependency_overrides.pop(ps.user_api_key_auth, None) +def _walk_row(session_key, sort_time, is_edge=False, batch_rows=0): + return { + "session_key": session_key, + "api_key": "hashed-key", + "last_activity": sort_time, + "sort_time": sort_time, + "request_id": f"req-{session_key}", + "is_edge": is_edge, + "batch_rows": batch_rows, + } + + +@pytest.mark.asyncio +async def test_ui_view_spend_logs_group_by_session_walk_finishes_a_tied_timestamp_before_settling(client, monkeypatch): + """Sessions tied on their newest startTime order by session key, so the walk reads past the tie before settling.""" + tied = "2026-08-29 10:00:00" + + async def mock_query_raw(sql_query, *params): + if "WITH batch AS MATERIALIZED" in sql_query: + if '("startTime", request_id) <' not in sql_query: + return [ + _walk_row("b", tied), + _walk_row("a", tied), + _walk_row("a", tied, is_edge=True, batch_rows=params[-1]), + ] + return [_walk_row("c", tied), _walk_row("c", "2026-08-29 09:00:00", is_edge=True, batch_rows=1)] + if "COUNT(*) AS total_count" in sql_query: + return [{"total_count": 3}] + if "CROSS JOIN LATERAL" in sql_query: + return [_session_representative_row(f"req-{session_key}", session_key) for session_key in params[-2]] + return [] + + mock_prisma = MagicMock() + mock_prisma.db = MagicMock() + mock_prisma.db.query_raw = AsyncMock(side_effect=mock_query_raw) + 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() + response = client.get( + "/spend/logs/ui", + params={"start_date": start_date, "end_date": end_date, "group_by_session": "true", "page_size": 1}, + headers={"Authorization": "Bearer sk-test"}, + ) + assert response.status_code == 200, response.text + data = response.json() + assert [row["request_id"] for row in data["data"]] == ["req-c"] + assert data["next_session_cursor"] == f"{tied}|hashed-key|c" + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + @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) serves the sessions that page starts at.""" diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_query_optimization.py b/tests/test_litellm/proxy/spend_tracking/test_spend_query_optimization.py index 449196b45c2..7b358015ec9 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_query_optimization.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_query_optimization.py @@ -535,6 +535,12 @@ async def test_spend_logs_ui_group_by_session_paginates_sessions(monkeypatch): ] async def mock_query_raw(sql_query, *params): + if "WITH batch AS MATERIALIZED" in sql_query: + first_rows = [ + {**row, "sort_time": row["last_activity"], "request_id": row["session_key"], "is_edge": False} + for row in session_rows + ] + return [*first_rows, {**first_rows[-1], "is_edge": True, "batch_rows": 0}] if "COUNT(*) AS total_count" in sql_query: return [{"total_count": 60}] if "CROSS JOIN LATERAL" in sql_query or "DISTINCT ON" in sql_query: @@ -570,7 +576,6 @@ async def test_spend_logs_ui_group_by_session_paginates_sessions(monkeypatch): emitted = [call[0] for call in mock_prisma.db.query_raw.call_args_list] page_sql = emitted[0][0] assert "OFFSET" not in page_sql, "the startTime page must be keyset-selected, not offset-selected" - assert emitted[0][-1] == 51, "the page query fetches page_size + 1 sessions to detect has_more" count_call = emitted[1] count_sql = count_call[0]