From 5f4f8b433a2115de129bcabe6c19ca03bf848655 Mon Sep 17 00:00:00 2001
From: Tin Chi Lo <tin@berri.ai>
Date: Wed, 29 Jul 2026 20:29:46 -0700
Subject: [PATCH] fix(proxy): stamp session_total_count on
 /spend/logs/session/ui rows

The session-scoped endpoint built its response without the enrichment
the main /spend/logs/ui list applies, so rows it returned always had
session_total_count unset. The logs UI drawer uses that field as a
fallback signal to stay in session mode; selecting a row from inside
an already-open session (rather than from the main list) silently
dropped the drawer out of session view since the newly selected row
never carried a valid count. Every row from this endpoint shares the
requested session_id, so the endpoint's own total_records is already
the correct value for every row.
---
 .../spend_management_endpoints.py             |  6 +++
 .../test_spend_management_endpoints.py        | 51 +++++++++++++++++++
 2 files changed, 57 insertions(+)

diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py
index 42788227acc..afed2c95ef9 100644
--- a/litellm/proxy/spend_tracking/spend_management_endpoints.py
+++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py
@@ -3364,6 +3364,12 @@ async def ui_view_session_spend_logs(
 
         total_pages = (total_records + page_size - 1) // page_size
 
+        # Every row here shares session_id, so total_records IS each row's session_total_count.
+        # The UI drawer needs this to stay in session mode when a row selected from this endpoint
+        # (rather than /spend/logs/ui) becomes the displayed log.
+        for row in result:
+            row["session_total_count"] = total_records
+
         return {
             "data": result,
             "total": total_records,
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 71206687b5c..3e6dc006210 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
@@ -1582,6 +1582,57 @@ async def test_ui_view_session_spend_logs_pagination(client, monkeypatch):
         app.dependency_overrides.pop(ps.user_api_key_auth, None)
 
 
+@pytest.mark.asyncio
+async def test_ui_view_session_spend_logs_includes_session_total_count(client, monkeypatch):
+    """
+    Regression test: /spend/logs/session/ui must stamp session_total_count onto
+    every row it returns. The UI logs drawer falls back to this field to decide
+    whether the currently displayed log belongs to a multi-call session; before
+    this fix the endpoint never set it, so selecting a row from inside an open
+    session (rather than from /spend/logs/ui) silently dropped the drawer out of
+    session mode.
+    """
+    mock_spend_logs = [
+        {"id": "log1", "request_id": "req1", "session_id": "session-123", "startTime": "2024-01-01T00:00:00Z"},
+        {"id": "log2", "request_id": "req2", "session_id": "session-123", "startTime": "2024-01-02T00:00:00Z"},
+    ]
+
+    class MockDB:
+        async def count(self, *args, **kwargs):
+            return len(mock_spend_logs)
+
+        async def query_raw(self, sql_query, session_id, page_size, skip):
+            # page_size=1 deliberately returns fewer rows than the session's real
+            # total, so a naive `len(result)` fallback would get this wrong too.
+            return [mock_spend_logs[0]]
+
+    class MockPrismaClient:
+        def __init__(self):
+            self.db = MockDB()
+            self.db.litellm_spendlogs = self.db
+
+    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MockPrismaClient())
+
+    app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
+        user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user"
+    )
+
+    try:
+        response = client.get(
+            "/spend/logs/session/ui",
+            params={"session_id": "session-123", "page": 1, "page_size": 1},
+            headers={"Authorization": "Bearer sk-test"},
+        )
+
+        assert response.status_code == 200
+        data = response.json()
+        assert data["total"] == 2
+        assert len(data["data"]) == 1
+        assert data["data"][0]["session_total_count"] == 2
+    finally:
+        app.dependency_overrides.pop(ps.user_api_key_auth, None)
+
+
 @pytest.mark.asyncio
 async def test_ui_view_session_spend_logs_scopes_non_admin_to_own_logs(client, monkeypatch):
     own_log = {