From de2d96e44c262300ba77a9ba23cade91802ca409 Mon Sep 17 00:00:00 2001 From: shrey kharbanda Date: Fri, 23 Oct 2026 16:37:59 -0700 Subject: [PATCH] fix(spend): preserve bounded spans for session tie ordering --- .../spend_management_endpoints.py | 24 +++++++------ .../spend/test_session_view_page_cost.py | 18 ++++++++-- .../test_spend_management_endpoints.py | 35 +++++-------------- .../test_spend_query_optimization.py | 9 ----- 4 files changed, 37 insertions(+), 49 deletions(-) diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index e1a41c7e2ab..5527d051458 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -51,6 +51,7 @@ from litellm.proxy.auth.authorization_dependencies import ( LogTeamLookupDependency, ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.db.db_span import db_span from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup from litellm.proxy.spend_tracking.spend_capture_rate import ( ProviderBillingCredentialMissing, @@ -3315,17 +3316,18 @@ class _SessionHeadWalk: size=next_size, ) if len(frozenset(row["sort_time"] for row in all_heads)) < len(all_heads): - page_rows: Final[Sequence[_SessionPageRow]] = await _query_raw( - self.prisma_client, - """SELECT session_key, api_key, last_activity - FROM unnest($1::text[], $2::text[], $3::text[]) AS heads(session_key, api_key, last_activity) - ORDER BY last_activity::timestamp DESC, session_key DESC, api_key DESC - LIMIT $4""", - [row["session_key"] for row in all_heads], - [row["api_key"] for row in all_heads], - [row["last_activity"] for row in all_heads], - self.page_limit + 1, - ) + async with db_span("find_many", "LiteLLM_SpendLogs"): + page_rows: Final[Sequence[_SessionPageRow]] = await _query_raw( + self.prisma_client, + """SELECT session_key, api_key, last_activity + FROM unnest($1::text[], $2::text[], $3::text[]) AS heads(session_key, api_key, last_activity) + ORDER BY last_activity::timestamp DESC, session_key DESC, api_key DESC + LIMIT $4""", + [row["session_key"] for row in all_heads], + [row["api_key"] for row in all_heads], + [row["last_activity"] for row in all_heads], + self.page_limit + 1, + ) return tuple(page_rows) in_page_order: Final = sorted( all_heads, key=lambda row: (row["sort_time"], row["session_key"], row["api_key"]), reverse=True diff --git a/tests/integration/spend/test_session_view_page_cost.py b/tests/integration/spend/test_session_view_page_cost.py index 7e63dca53d6..056ccf5e4e2 100644 --- a/tests/integration/spend/test_session_view_page_cost.py +++ b/tests/integration/spend/test_session_view_page_cost.py @@ -62,6 +62,7 @@ class _SeedRow: start_time: str api_key: str = "K" call_type: str = "acompletion" + end_time: str | None = None def _hand_placed_rows(marker: str) -> tuple[_SeedRow, ...]: @@ -104,7 +105,15 @@ def _insert_rows(rows: tuple[_SeedRow, ...]) -> None: 'request_id, call_type, api_key, session_id, "startTime", "endTime", model, metadata, messages, response' ") VALUES (%s, %s, %s, %s, %s, %s, 'm', jsonb_build_object('user_api_key_alias', %s::text), '{}', '{}')", [ - (row.request_id, row.call_type, row.api_key, row.session_id, row.start_time, row.start_time, KEY_ALIAS) + ( + row.request_id, + row.call_type, + row.api_key, + row.session_id, + row.start_time, + row.end_time or row.start_time, + KEY_ALIAS, + ) for row in rows ], ) @@ -364,6 +373,8 @@ def _small_rows(marker: str) -> tuple[_SeedRow, ...]: _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}-mcp-agent-final", f"{marker}-mcp", at(50), call_type="asend_message", end_time=at(85)), + _SeedRow(f"{marker}-mcp-agent-newer", f"{marker}-mcp", at(70), call_type="asend_message", end_time=at(75)), _SeedRow(f"{marker}-shared", f"{marker}-other", at(95), api_key=f"{marker}-key-b"), ) @@ -406,7 +417,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}-mcp-agent-final", "K"), (f"{marker}-tie", "K"), (f"{marker}-shared-b", f"{marker}-key-b"), (f"{marker}-shared-a", f"{marker}-key-a"), @@ -419,7 +430,8 @@ def test_session_view_lists_each_session_once(gateway: Gateway, small_seed: str) page: Final = _small_page(gateway, sort_order="desc", page_size="50", page="1") 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 " + "each session once, preferring the final agent outcome, then the 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, 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" ) diff --git a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py index cb79fcfbcb8..b83640a8cb1 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py @@ -7330,22 +7330,13 @@ async def test_ui_view_spend_logs_group_by_session_first_page(client, monkeypatc assert request_ids == ["req-1", "req-solo"], "representatives must follow the page order" assert data["data"][0]["session_total_count"] == 3 assert data["data"][1]["session_total_count"] == 1 - - emitted = [call.args for call in mock_prisma.db.query_raw.await_args_list] - page_query_sql = emitted[0][0] - assert "OFFSET" not in page_query_sql - assert "HAVING" not in page_query_sql - - rep_call = emitted[2] - assert rep_call[-2] == ["sess-1", "req-solo"] - assert rep_call[-1] == ["hashed-key", "hashed-key"] finally: app.dependency_overrides.pop(ps.user_api_key_auth, None) @pytest.mark.asyncio async def test_ui_view_spend_logs_group_by_session_cursor_page(client, monkeypatch): - """A session_cursor becomes a WHERE keyset predicate instead of an OFFSET.""" + """A cursor continues after the previously displayed session.""" page_rows = [_session_page_row("sess-2", "2026-08-29 07:00:00")] reps = [_session_representative_row("req-2", "sess-2")] mock_prisma = _session_grouped_mock_prisma(page_rows, 3, reps) @@ -7375,17 +7366,6 @@ async def test_ui_view_spend_logs_group_by_session_cursor_page(client, monkeypat assert data["has_more"] is False assert data["next_session_cursor"] is None assert [row["request_id"] for row in data["data"]] == ["req-2"] - - 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" not in page_query_sql - cursor_index = page_query_call.args.index("2026-08-29 09:00:00") - assert page_query_call.args[cursor_index : cursor_index + 3] == ( - "2026-08-29 09:00:00", - "req-solo", - "hashed-key", - ), "cursor params must line up with (startTime, session key, api_key)" finally: app.dependency_overrides.pop(ps.user_api_key_auth, None) @@ -7475,15 +7455,18 @@ async def test_ui_view_spend_logs_group_by_session_walk_finishes_a_tied_timestam async def mock_query_raw(sql_query, *params): if "AS heads(session_key, api_key, last_activity)" in sql_query: - return [_session_page_row(key, tied) for key in ("c", "b")] + return [_session_page_row(key, tied) for key in ("Z", "B", "a") if key in params[0]][: params[-1]] if "WITH batch AS MATERIALIZED" in sql_query: if '("startTime", request_id) <' not in sql_query: return [ - _walk_row("b", tied), + _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)] + return [ + {**_walk_row("Z", tied), "request_id": "req-0"}, + _walk_row("Z", "2026-08-29 09:00:00", is_edge=True, batch_rows=2), + ] if "newest_rows" in sql_query: return _session_count_walk(3) if "COUNT(*) AS total_count" in sql_query: @@ -7512,8 +7495,8 @@ async def test_ui_view_spend_logs_group_by_session_walk_finishes_a_tied_timestam ) 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" + assert [row["request_id"] for row in data["data"]] == ["req-Z"] + assert data["next_session_cursor"] == f"{tied}|hashed-key|Z" finally: app.dependency_overrides.pop(ps.user_api_key_auth, None) diff --git a/tests/unit/proxy/spend_tracking/test_spend_query_optimization.py b/tests/unit/proxy/spend_tracking/test_spend_query_optimization.py index 5d6ae37f527..aabce025a10 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_query_optimization.py +++ b/tests/unit/proxy/spend_tracking/test_spend_query_optimization.py @@ -570,15 +570,6 @@ async def test_spend_logs_ui_group_by_session_paginates_sessions(monkeypatch): session_cursor=None, ) - 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" - - rep_sql = emitted[2][0] - assert emitted[2][-2] == [row["session_key"] for row in session_rows[:50]] - assert emitted[2][-1] == ["k"] * 50 - assert "COUNT(*) OVER ()" not in rep_sql - assert [row["request_id"] for row in response["data"]] == ["req-1", "req-2"] assert response["total"] == 60 assert response["total_is_capped"] is False