fix(spend): preserve bounded spans for session tie ordering

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

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"
     )

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)
 

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