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.
This commit is contained in:
Tin Chi Lo 2026-07-29 20:29:46 -07:00
parent 072a6eef50
commit 5f4f8b433a
2 changed files with 57 additions and 0 deletions

View file

@ -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,

View file

@ -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 = {