scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-03 07:23:33 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 166

Merge de2d96e44c into 23ccff69a1

Browse source
This commit is contained in:
devin-ai-integration[bot] 2026-10-23 23:38:12 +00:00 • committed by GitHub
commit cecdd7e194
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 1113 additions and 91 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

359
litellm/proxy/spend_tracking/spend_management_endpoints.py
View file

@ -3,6 +3,7 @@ import collections
import json
import os
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from datetime import date, datetime, timedelta, timezone
from functools import partial
from itertools import groupby
@ -50,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,
@ -88,6 +90,10 @@ _DECODED_JSON: Final = TypeAdapter(object)
_SESSION_KEY_EXPR: Final = "COALESCE(NULLIF(session_id, ''), request_id)"
_SESSION_GROUP_KEY_SQL: Final = f"{_SESSION_KEY_EXPR}, api_key"
_SESSION_WALK_FIRST_BATCH_ROWS: Final = 1000
_SESSION_WALK_MAX_ROWS: Final = 20000
_SESSION_COUNT_FIRST_BATCH_ROWS: Final = 50000
_SESSION_COUNT_MAX_ROWS: Final = 200000
_MCP_CALL_TYPES_SQL: Final = "('call_mcp_tool', 'list_mcp_tools')"
_AGENT_CALL_TYPE_SQL: Final = "'asend_message'"
_SESSION_REPRESENTATIVE_ORDER_SQL: Final = (
@ -184,6 +190,11 @@ class _SpendLogsCountRow(TypedDict):
total_count: int
class _SessionCountBatchRow(TypedDict):
rows_read: ReadOnly[int]
sessions: ReadOnly[int]
class _PgClassRow(TypedDict):
relname: str
relkind: str
@ -2847,6 +2858,7 @@ async def ui_view_spend_logs(
sql_conditions.append(f"\"startTime\" <= (${p}::timestamptz AT TIME ZONE 'UTC')")
sql_params.append(end_date_obj)
p += 1
window_conditions: Final = tuple(sql_conditions)
if search is not None and start_date_obj is not None and end_date_obj is not None:
search_condition: Final = _build_spend_log_search_condition(
@ -2971,6 +2983,7 @@ async def ui_view_spend_logs(
):
return await _ui_session_grouped_spend_logs(
prisma_client=prisma_client,
window_conditions=window_conditions,
sql_conditions=sql_conditions,
sql_params=sql_params,
next_param_index=p,
@ -3093,6 +3106,16 @@ def _parse_session_cursor(session_cursor: str | None) -> tuple[str, str, str] |
return (last_activity, session_key, api_key)
def _session_rows_sql(session_key: str, api_key: str) -> str:
"""Every row that can belong to one session, found by its key alone. The OFFSET 0 keeps the caller's window
and filters out of this scan, so a generic plan cannot combine it with a scan of the whole window."""
return f"""SELECT *
FROM "LiteLLM_SpendLogs"
WHERE (request_id = {session_key} OR session_id = {session_key})
AND api_key = {api_key}
OFFSET 0"""
async def _fetch_session_representatives(
prisma_client: "PrismaClient",
where_clause: str,
@ -3102,16 +3125,16 @@ async def _fetch_session_representatives(
) -> list[dict[str, object]]: # mutable-ok: _build_ui_spend_logs_response writes session counts onto each row
"""Fetch the final agent outcome, or newest non-MCP row, of each ``(session_key, api_key)`` session, in ``session_keys`` order."""
rep_query: Final = f"""
SELECT * FROM (
SELECT DISTINCT ON ({_SESSION_GROUP_KEY_SQL})
{_SPEND_LOG_LIST_COLUMNS}
FROM "LiteLLM_SpendLogs"
SELECT rep.*
FROM unnest(${next_param_index}::text[], ${next_param_index + 1}::text[]) AS requested(session_key, key)
CROSS JOIN LATERAL (
SELECT {_SPEND_LOG_LIST_COLUMNS}
FROM ({_session_rows_sql("requested.session_key", "requested.key")}) AS "LiteLLM_SpendLogs"
WHERE {where_clause}
AND ({_SESSION_GROUP_KEY_SQL}) IN (
SELECT * FROM unnest(${next_param_index}::text[], ${next_param_index + 1}::text[])
)
ORDER BY {_SESSION_GROUP_KEY_SQL}, {_SESSION_REPRESENTATIVE_ORDER_SQL}
) AS session_representatives
AND {_SESSION_KEY_EXPR} = requested.session_key
ORDER BY {_SESSION_REPRESENTATIVE_ORDER_SQL}
LIMIT 1
) AS rep
"""
rep_rows: Final[Sequence[dict[str, object]]] = await _query_raw( # mutable-ok: rows are enriched in place
prisma_client,
@ -3126,6 +3149,211 @@ async def _fetch_session_representatives(
return [rep_by_key[key] for key in session_keys if key in rep_by_key]
class _SessionWalkRow(TypedDict):
session_key: ReadOnly[str]
api_key: ReadOnly[str]
last_activity: ReadOnly[str]
sort_time: ReadOnly[str]
request_id: ReadOnly[str]
is_edge: ReadOnly[bool]
batch_rows: ReadOnly[int]
class _SessionKeyRow(TypedDict):
session_key: ReadOnly[str]
api_key: ReadOnly[str]
@dataclass(frozen=True, slots=True)
class _SessionWalkBatch:
first_rows: tuple[_SessionWalkRow, ...]
edge: _SessionWalkRow | None
rows_read: int
@dataclass(frozen=True, slots=True)
class _SessionHeadWalk:
"""Walks the window newest first; the first row it meets of each session is that session's newest row."""
prisma_client: "PrismaClient"
window_clause: str
where_clause: str
sql_params: tuple[object, ...]
next_param_index: int
cursor: tuple[str, str, str] | None
page_limit: int
def _cursor_bounds(self) -> str:
return f"(${self.next_param_index}::timestamp, ${self.next_param_index + 1}, ${self.next_param_index + 2})"
def _params(self) -> tuple[object, ...]:
return (*self.sql_params, *(self.cursor or ()))
async def _batch(self, after: _SessionWalkRow | None, size: int) -> _SessionWalkBatch:
params: Final = self._params()
keyset_clause: Final = (
f'AND ("startTime", {_SESSION_KEY_EXPR}, api_key) < {self._cursor_bounds()}' if self.cursor else ""
)
below_cursor_clause: Final = f'AND "startTime" <= ${self.next_param_index}::timestamp' if self.cursor else ""
after_clause: Final = (
f'AND ("startTime", request_id) < (${len(params) + 1}::timestamp, ${len(params) + 2})' if after else ""
)
after_params: Final = (after["last_activity"], after["request_id"]) if after else ()
walk_columns: Final = (
"session_key, api_key, start_time::text AS last_activity, "
"to_char(start_time, 'YYYY-MM-DD HH24:MI:SS.US') AS sort_time, request_id"
)
query: Final = f"""
WITH batch AS MATERIALIZED (
SELECT {_SESSION_KEY_EXPR} AS session_key,
api_key,
"startTime" AS start_time,
request_id,
COALESCE(({self.where_clause}) {keyset_clause}, FALSE) AS in_page
FROM (
SELECT *
FROM "LiteLLM_SpendLogs"
WHERE {self.window_clause}
{below_cursor_clause}
{after_clause}
ORDER BY "startTime" DESC, request_id DESC
LIMIT ${len(params) + len(after_params) + 1}
) AS "LiteLLM_SpendLogs"
)
SELECT {walk_columns}, FALSE AS is_edge, 0 AS batch_rows
FROM (
SELECT DISTINCT ON (session_key, api_key) *
FROM batch
WHERE in_page
ORDER BY session_key, api_key, start_time DESC, request_id DESC
) AS first_rows
UNION ALL
SELECT {walk_columns}, TRUE AS is_edge, (SELECT COUNT(*) FROM batch)::int AS batch_rows
FROM (SELECT * FROM batch ORDER BY start_time, request_id LIMIT 1) AS edge
"""
rows: Final[Sequence[_SessionWalkRow]] = await _query_raw(
self.prisma_client, query, *params, *after_params, size
)
edges: Final = tuple(row for row in rows if row["is_edge"])
return _SessionWalkBatch(
first_rows=tuple(
sorted(
(row for row in rows if not row["is_edge"]),
key=lambda row: (row["sort_time"], row["request_id"]),
reverse=True,
)
),
edge=edges[0] if edges else None,
rows_read=edges[0]["batch_rows"] if edges else 0,
)
async def _shown_before_the_cursor(self, candidates: Sequence[_SessionWalkRow]) -> frozenset[tuple[str, str]]:
if not self.cursor or not candidates:
return frozenset()
params: Final = self._params()
query: Final = f"""
SELECT requested.session_key, requested.key AS api_key
FROM unnest(${len(params) + 1}::text[], ${len(params) + 2}::text[]) AS requested(session_key, key)
CROSS JOIN LATERAL (
SELECT 1
FROM ({_session_rows_sql("requested.session_key", "requested.key")}) AS "LiteLLM_SpendLogs"
WHERE {self.where_clause}
AND {_SESSION_KEY_EXPR} = requested.session_key
AND ("startTime", requested.session_key, requested.key) >= {self._cursor_bounds()}
LIMIT 1
) AS newer_row
"""
shown: Final[Sequence[_SessionKeyRow]] = await _query_raw(
self.prisma_client,
query,
*params,
[row["session_key"] for row in candidates],
[row["api_key"] for row in candidates],
)
return frozenset((row["session_key"], row["api_key"]) for row in shown)
async def page(self) -> tuple[_SessionPageRow, ...] | None:
"""The first ``page_limit + 1`` sessions in page order, or None when the walk gives up before it can tell."""
return await self._walk(
after=None, heads=(), seen=frozenset(), rows_read=0, size=_SESSION_WALK_FIRST_BATCH_ROWS
)
async def _walk(
self,
after: _SessionWalkRow | None,
heads: tuple[_SessionWalkRow, ...],
seen: frozenset[tuple[str, str]],
rows_read: int,
size: int,
) -> tuple[_SessionPageRow, ...] | None:
batch: Final = await self._batch(after, size)
first_seen: Final = tuple(row for row in batch.first_rows if (row["session_key"], row["api_key"]) not in seen)
shown_before: Final = await self._shown_before_the_cursor(first_seen)
all_heads: Final = heads + tuple(
row for row in first_seen if (row["session_key"], row["api_key"]) not in shown_before
)
window_exhausted: Final = batch.edge is None or batch.rows_read < size
page_is_settled: Final = (
len(all_heads) > self.page_limit
and batch.edge is not None
and batch.edge["sort_time"] < all_heads[self.page_limit]["sort_time"]
)
total_read: Final = rows_read + batch.rows_read
sessions_missing: Final = self.page_limit + 1 - len(all_heads)
next_size: Final = (
max(_SESSION_WALK_FIRST_BATCH_ROWS, sessions_missing * total_read * 5 // (4 * len(all_heads)))
if all_heads
else _SESSION_WALK_MAX_ROWS
)
if not (window_exhausted or page_is_settled):
if total_read + next_size > _SESSION_WALK_MAX_ROWS:
return None
return await self._walk(
after=batch.edge,
heads=all_heads,
seen=seen | frozenset((row["session_key"], row["api_key"]) for row in first_seen),
rows_read=total_read,
size=next_size,
)
if len(frozenset(row["sort_time"] for row in all_heads)) < len(all_heads):
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
)
return tuple(
_SessionPageRow(session_key=row["session_key"], api_key=row["api_key"], last_activity=row["last_activity"])
for row in in_page_order[: self.page_limit + 1]
)
def _grouped_session_page_sql(
where_clause: str, having_clause: str, direction: str, limit_index: int, offset_clause: str
) -> str:
return f"""
SELECT {_SESSION_KEY_EXPR} AS session_key,
api_key,
MAX("startTime")::text AS last_activity
FROM "LiteLLM_SpendLogs"
WHERE {where_clause}
GROUP BY {_SESSION_GROUP_KEY_SQL}
{having_clause}
ORDER BY MAX("startTime") {direction}, {_SESSION_KEY_EXPR} {direction}, api_key {direction}
LIMIT ${limit_index} {offset_clause}
"""
async def _count_grouped_sessions(
prisma_client: "PrismaClient",
where_clause: str,
@ -3152,8 +3380,55 @@ async def _count_grouped_sessions(
)
async def _count_sessions_newest_first(
prisma_client: "PrismaClient",
window_clause: str,
where_clause: str,
sql_params: Sequence[object],
next_param_index: int,
size: int = _SESSION_COUNT_FIRST_BATCH_ROWS,
) -> tuple[int, bool] | None:
"""Count sessions in the newest ``size`` rows of the window, stopping as soon as the count cap is passed or the
window runs out. None when the sessions are too long to settle within ``_SESSION_COUNT_MAX_ROWS`` rows."""
count_query: Final = f"""
WITH newest_rows AS MATERIALIZED (
SELECT {_SESSION_KEY_EXPR} AS session_key,
api_key,
COALESCE(({where_clause}), FALSE) AS in_page
FROM (
SELECT *
FROM "LiteLLM_SpendLogs"
WHERE {window_clause}
ORDER BY "startTime" DESC, request_id DESC
LIMIT ${next_param_index}
) AS "LiteLLM_SpendLogs"
)
SELECT (SELECT COUNT(*) FROM newest_rows)::int AS rows_read,
(SELECT COUNT(*) FROM (SELECT DISTINCT session_key, api_key FROM newest_rows WHERE in_page) AS s)::int
AS sessions
"""
rows: Final[Sequence[_SessionCountBatchRow]] = await _query_raw(prisma_client, count_query, *sql_params, size)
rows_read: Final = rows[0]["rows_read"] if rows else 0
sessions: Final = rows[0]["sessions"] if rows else 0
if sessions > SPEND_LOGS_PAGINATION_COUNT_CAP:
return (SPEND_LOGS_PAGINATION_COUNT_CAP, True)
if rows_read < size:
return (sessions, False)
next_size: Final = (
max(2 * size, (SPEND_LOGS_PAGINATION_COUNT_CAP + 1) * rows_read * 5 // (4 * sessions))
if sessions
else _SESSION_COUNT_MAX_ROWS + 1
)
if next_size > _SESSION_COUNT_MAX_ROWS:
return None
return await _count_sessions_newest_first(
prisma_client, window_clause, where_clause, sql_params, next_param_index, next_size
)
async def _ui_session_grouped_spend_logs(
prisma_client: "PrismaClient",
window_conditions: Sequence[str],
sql_conditions: Sequence[str],
sql_params: Sequence[object],
next_param_index: int,
@ -3171,7 +3446,11 @@ async def _ui_session_grouped_spend_logs(
next ``page_size`` sessions ordered by ``(MAX(startTime), session_key,
api_key)``, resumed from the ``session_cursor`` keyset
``'<last_activity>|<api_key>|<session_key>'`` instead of an OFFSET, so
page depth does not degrade the query plan. A request for ``page > 1``
page depth does not degrade the query plan. A newest-first page comes from
``_SessionHeadWalk``, which stops after ``page_size + 1`` sessions instead
of grouping the whole window; ascending and OFFSET pages, and walks that
read ``_SESSION_WALK_MAX_ROWS`` rows without settling the page, group the
window instead. A request for ``page > 1``
without a cursor (the UI jumping straight to the last page, or back to a
page it never walked through) falls back to ``OFFSET (page - 1) *
page_size``, trimmed to the end of the ``SPEND_LOGS_PAGINATION_COUNT_CAP``
@ -3182,10 +3461,14 @@ async def _ui_session_grouped_spend_logs(
``next_session_cursor`` / ``has_more`` while ``total`` counts sessions
(capped like the flat total). A page that runs out of sessions while still
holding some is itself the end of the list, so its ``total`` is
``offset + len(page)`` and the grouped count query is skipped; a page that
``offset + len(page)`` and the count is skipped; a page that
starts past the end says nothing about the total, so that one is counted.
The count walks the window newest first like the page and stops once it
passes the cap, grouping the whole window only when the sessions are too
long for that walk to settle within ``_SESSION_COUNT_MAX_ROWS`` rows.
"""
where_clause: Final = " AND ".join(sql_conditions) if sql_conditions else "TRUE"
window_clause: Final = " AND ".join(window_conditions) if window_conditions else "TRUE"
cmp_op: Final = "<" if sort_desc else ">"
direction: Final = "DESC" if sort_desc else "ASC"
@ -3203,21 +3486,32 @@ async def _ui_session_grouped_spend_logs(
offset_params: Final[tuple[int, ...]] = (offset,) if offset and page_limit > 0 else ()
offset_clause: Final = f"OFFSET ${limit_index + 1}" if offset_params else ""
page_query: Final = f"""
SELECT {_SESSION_KEY_EXPR} AS session_key,
api_key,
MAX("startTime")::text AS last_activity
FROM "LiteLLM_SpendLogs"
WHERE {where_clause}
GROUP BY {_SESSION_GROUP_KEY_SQL}
{having_clause}
ORDER BY MAX("startTime") {direction}, {_SESSION_KEY_EXPR} {direction}, api_key {direction}
LIMIT ${limit_index} {offset_clause}
"""
walked_rows: Final = (
await _SessionHeadWalk(
prisma_client=prisma_client,
window_clause=window_clause,
where_clause=where_clause,
sql_params=tuple(sql_params),
next_param_index=next_param_index,
cursor=cursor,
page_limit=page_limit,
).page()
if sort_desc and offset == 0 and page_limit > 0
else None
)
page_rows: Final[Sequence[_SessionPageRow]] = (
()
if page_limit <= 0
else await _query_raw(prisma_client, page_query, *sql_params, *cursor_params, page_limit + 1, *offset_params)
else walked_rows
if walked_rows is not None
else await _query_raw(
prisma_client,
_grouped_session_page_sql(where_clause, having_clause, direction, limit_index, offset_clause),
*sql_params,
*cursor_params,
page_limit + 1,
*offset_params,
)
)
has_more: Final = len(page_rows) > page_limit
@ -3230,9 +3524,18 @@ async def _ui_session_grouped_spend_logs(
page_starts_inside_the_list: Final = offset == 0 or len(page_rows) > 0
page_ends_the_list: Final = cursor is None and page_limit > 0 and not has_more and page_starts_inside_the_list
counted_newest_first: Final = (
None
if page_ends_the_list
else await _count_sessions_newest_first(
prisma_client, window_clause, where_clause, sql_params, next_param_index
)
)
total_records, total_is_capped = (
(offset + len(page_rows), False)
if page_ends_the_list
else counted_newest_first
if counted_newest_first is not None
else await _count_grouped_sessions(prisma_client, where_clause, sql_params, next_param_index)
)
@ -4730,9 +5033,13 @@ async def _build_ui_spend_logs_response(
COALESCE(SUM(prompt_tokens), 0)::bigint AS session_total_prompt_tokens,
COALESCE(SUM(completion_tokens), 0)::bigint AS session_total_completion_tokens,
COALESCE(SUM(total_tokens), 0)::bigint AS session_total_tokens
FROM "LiteLLM_SpendLogs"
WHERE session_id = ANY($1::text[])
AND api_key = ANY($2::text[])
FROM unnest($1::text[]) AS requested(session_key)
CROSS JOIN LATERAL (
SELECT * FROM "LiteLLM_SpendLogs"
WHERE session_id = requested.session_key
AND api_key = ANY($2::text[])
OFFSET 0
) AS session_rows
GROUP BY session_id, api_key
) s
LEFT JOIN LATERAL (

524
tests/integration/spend/test_session_view_page_cost.py Normal file
View file

@ -0,0 +1,524 @@
"""The session-grouped /spend/logs/ui page lists each session once and reads a bounded slice of the window.
The page stops after page_size + 1 sessions and the session count stops once it
passes the count cap, so a window with fewer sessions than the cap is read at
most once. Tuple reads come from pg_stat, and the ordering, representative and
total assertions pin what the page returns.
"""
import os
import time
import uuid
from collections.abc import Callable, Iterator, Mapping
from dataclasses import dataclass
from datetime import datetime, timedelta
from typing import Final
import psycopg
import pytest
from integration._support.client import Gateway, string_value
from integration._support.database import ROWS, read_rows
from pydantic import JsonValue
from litellm.proxy.spend_tracking.spend_management_endpoints import SPEND_LOGS_PAGINATION_COUNT_CAP
SEED_ROWS: Final = 40_000
LONG_SESSION_ROWS: Final = 5_000
SESSION_LEN: Final = 5
BULK_ROWS: Final = SEED_ROWS - LONG_SESSION_ROWS
PAGE_SIZE: Final = 50
MAX_TUPLES_PER_PAGE: Final = SEED_ROWS * 2
_BASE: Final = datetime(2001, 1, 1)
M1_NEWEST_SECOND: Final = SEED_ROWS + 8
M2_NEWEST_SECOND: Final = SEED_ROWS + 5
BLOCKER_NEWEST_SECOND: Final = SEED_ROWS + 3
WINDOW_START: Final = _BASE.strftime("%Y-%m-%d %H:%M:%S")
WINDOW_END: Final = (_BASE + timedelta(seconds=M1_NEWEST_SECOND)).strftime("%Y-%m-%d %H:%M:%S")
_BULK_SESSIONS: Final = len({n // SESSION_LEN for n in range(1, BULK_ROWS + 1)})
EXPECTED_SESSIONS: Final = _BULK_SESSIONS + 5
_STATS_SETTLE_S: Final = 4.0
_STATS_STABLE_POLLS: Final = 4
_STATS_TIMEOUT_S: Final = 30.0
_QUIET_SETTLE_S: Final = 12.0
_QUIET_STABLE_POLLS: Final = 5
_QUIET_TIMEOUT_S: Final = 60.0
def _at(seconds: int) -> str:
return (_BASE + timedelta(seconds=seconds)).strftime("%Y-%m-%d %H:%M:%S")
KEY_ALIAS: Final = "intg-sesswin-alias"
_SMALL_BASE: Final = datetime(2002, 6, 1)
SMALL_WINDOW_START: Final = _SMALL_BASE.strftime("%Y-%m-%d %H:%M:%S")
SMALL_WINDOW_END: Final = (_SMALL_BASE + timedelta(seconds=100)).strftime("%Y-%m-%d %H:%M:%S")
@dataclass(frozen=True, slots=True)
class _SeedRow:
request_id: str
session_id: str | None
start_time: str
api_key: str = "K"
call_type: str = "acompletion"
end_time: str | None = None
def _hand_placed_rows(marker: str) -> tuple[_SeedRow, ...]:
prefix: Final = f"intg-sesswin-{marker}-"
blocker_request_id: Final = f"{prefix}x"
return (
_SeedRow(blocker_request_id, None, _at(SEED_ROWS + 1)),
_SeedRow(f"{prefix}x1", blocker_request_id, _at(SEED_ROWS + 2)),
_SeedRow(f"{prefix}x2", blocker_request_id, _at(BLOCKER_NEWEST_SECOND)),
_SeedRow(f"{prefix}m2a", f"{marker}-m2", _at(SEED_ROWS + 4)),
_SeedRow(f"{prefix}m2b", f"{marker}-m2", _at(M2_NEWEST_SECOND)),
_SeedRow(f"{prefix}m1a", f"{marker}-m1", _at(SEED_ROWS + 6)),
_SeedRow(f"{prefix}m1b", f"{marker}-m1", _at(SEED_ROWS + 7)),
_SeedRow(f"{prefix}m1c", f"{marker}-m1", _at(M1_NEWEST_SECOND)),
_SeedRow(f"intg-straddle-{marker}-in", f"{marker}-straddle", _at(1000)),
_SeedRow(f"intg-straddle-{marker}-out", f"{marker}-straddle", _at(SEED_ROWS + 2_000_000)),
)
def _seed(marker: str) -> None:
with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as connection:
connection.execute(
'INSERT INTO "LiteLLM_SpendLogs" ('
'request_id, call_type, api_key, session_id, "startTime", "endTime", model, metadata, messages, response'
") "
"SELECT %s || n, 'acompletion', 'K', "
"CASE WHEN n > %s THEN %s || '-long' ELSE %s || '-s-' || (n / %s) END, "
"ts, ts, 'e2e-sesswin', jsonb_build_object('user_api_key_alias', %s::text), '{}'::jsonb, '{}'::jsonb "
"FROM generate_series(1, %s) AS g(n) "
"CROSS JOIN LATERAL (SELECT %s::timestamp + (n || ' seconds')::interval AS ts) AS t",
(f"intg-sesswin-{marker}-", BULK_ROWS, marker, marker, SESSION_LEN, KEY_ALIAS, SEED_ROWS, WINDOW_START),
)
_insert_rows(_hand_placed_rows(marker))
def _insert_rows(rows: tuple[_SeedRow, ...]) -> None:
with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as connection:
connection.cursor().executemany(
'INSERT INTO "LiteLLM_SpendLogs" ('
'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.end_time or row.start_time,
KEY_ALIAS,
)
for row in rows
],
)
connection.execute('VACUUM ANALYZE "LiteLLM_SpendLogs"')
def _settle_autovacuum() -> None:
with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as connection:
connection.execute('VACUUM ANALYZE "LiteLLM_SpendLogs"')
connection.execute('ALTER TABLE "LiteLLM_SpendLogs" SET (autovacuum_enabled = false)')
def _cleanup(marker: str) -> None:
with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as connection:
try:
connection.execute(
'DELETE FROM "LiteLLM_SpendLogs" WHERE request_id LIKE %s OR request_id LIKE %s',
(f"intg-sesswin-{marker}-%", f"intg-straddle-{marker}-%"),
)
connection.execute('VACUUM "LiteLLM_SpendLogs"')
finally:
connection.execute('ALTER TABLE "LiteLLM_SpendLogs" RESET (autovacuum_enabled)')
_wait_for_stats_quiet()
def _wait_for_stats_quiet() -> None:
_wait_until_stable(
_tuples_read,
settle_s=_STATS_SETTLE_S,
stable_polls=_STATS_STABLE_POLLS,
timeout_s=_STATS_TIMEOUT_S,
)
def _tuples_read() -> int:
"""Heap tuples Postgres reports read on LiteLLM_SpendLogs: seq_tup_read
plus idx_tup_fetch (heap rows fetched through index scans). Index entries
traversed do not count, so bulk INSERTs never inflate the reading."""
with psycopg.connect(os.environ["DATABASE_URL"]) as connection:
row = connection.execute(
"SELECT t.seq_tup_read, t.idx_tup_fetch FROM pg_stat_user_tables t WHERE t.relname = 'LiteLLM_SpendLogs'",
).fetchone()
assert row is not None
return int(row[0] or 0) + int(row[1] or 0)
def _wait_until_stable(read: Callable[[], int], *, settle_s: float, stable_polls: int, timeout_s: float) -> int:
"""pg_stat readings once the collector has gone quiet. Each pooled backend
flushes its stats independently, so the counter climbs in delayed jumps:
wait at least ``settle_s`` seconds, then take the reading once it is
unchanged for ``stable_polls`` consecutive polls."""
deadline: Final = time.monotonic() + timeout_s
min_wait: Final = time.monotonic() + settle_s
stable: int = 0
last: int = read()
while time.monotonic() < deadline:
current: int = read()
stable = stable + 1 if current == last else 0
last = current
if stable >= stable_polls and time.monotonic() > min_wait:
return current
time.sleep(0.5)
raise AssertionError("pg_stat counters never settled")
def _tuples_read_quiet() -> int:
return _wait_until_stable(
_tuples_read,
settle_s=_QUIET_SETTLE_S,
stable_polls=_QUIET_STABLE_POLLS,
timeout_s=_QUIET_TIMEOUT_S,
)
def _data_rows(page: Mapping[str, JsonValue]) -> list[dict[str, JsonValue]]:
return ROWS.validate_python(page["data"])
def _session_page(gateway: Gateway, **extra: str) -> dict[str, JsonValue]:
return gateway.get(
"/spend/logs/ui",
params={
"group_by_session": "true",
"start_date": WINDOW_START,
"end_date": WINDOW_END,
"sort_by": "startTime",
"sort_order": "desc",
"page_size": str(PAGE_SIZE),
**extra,
},
)
@pytest.fixture
def seeded(gateway: Gateway) -> Iterator[str]:
marker: Final = uuid.uuid4().hex
with gateway.scenario() as scenario:
scenario.cleanups.callback(_cleanup, marker)
_settle_autovacuum()
_seed(marker)
_wait_for_stats_quiet()
yield marker
def test_session_view_first_page_reads_a_bounded_slice(gateway: Gateway, seeded: str) -> None:
prefix: Final = f"intg-sesswin-{seeded}-"
before: Final = _tuples_read_quiet()
page: Final = _session_page(gateway, page="1")
tuples_read: Final = _tuples_read_quiet() - before
assert page["has_more"] is True
assert page["total"] == EXPECTED_SESSIONS
assert EXPECTED_SESSIONS < SPEND_LOGS_PAGINATION_COUNT_CAP
assert page["total_is_capped"] is False
data: Final = _data_rows(page)
assert len(data) == PAGE_SIZE
assert data[0]["request_id"] == f"{prefix}m1c", "multi-row session must be represented by its newest row"
assert data[1]["request_id"] == f"{prefix}m2b"
assert data[2]["request_id"] == f"{prefix}x2", (
"a NULL session_id row and the rows adopting its request_id as session_id form one session"
)
blocker_keys: Final = {f"{prefix}x", f"{prefix}x1", f"{prefix}x2"}
assert sum(1 for row in data if row["request_id"] in blocker_keys) == 1
assert data[3]["request_id"] == f"{prefix}{SEED_ROWS}", "the long session must be represented by its newest row"
expected_fivers: Final = [BULK_ROWS, BULK_ROWS - 1, BULK_ROWS - SESSION_LEN - 1]
assert [row["request_id"] for row in data[4:7]] == [f"{prefix}{n}" for n in expected_fivers], (
"each 5-row session must be represented by its newest row, newest session first"
)
start_times: Final = [string_value(row["startTime"]) for row in data]
assert start_times == sorted(start_times, reverse=True)
assert tuples_read < MAX_TUPLES_PER_PAGE, (
f"first page read {tuples_read} tuples for a window of {SEED_ROWS} rows and page_size {PAGE_SIZE}; "
"the page must stop after page_size + 1 sessions, leaving the capped count as the only full pass"
)
def test_session_view_cursor_page_reads_a_bounded_slice(gateway: Gateway, seeded: str) -> None:
first: Final = _session_page(gateway, page="1")
cursor: Final = first["next_session_cursor"]
assert isinstance(cursor, str)
before: Final = _tuples_read_quiet()
second: Final = _session_page(gateway, page="2", session_cursor=cursor)
tuples_read: Final = _tuples_read_quiet() - before
second_data: Final = _data_rows(second)
assert len(second_data) == PAGE_SIZE
oldest_first: Final = string_value(_data_rows(first)[-1]["startTime"])
assert all(string_value(row["startTime"]) < oldest_first for row in second_data)
assert tuples_read < MAX_TUPLES_PER_PAGE, (
f"cursor page read {tuples_read} tuples for a window of {SEED_ROWS} rows and page_size {PAGE_SIZE}; "
"the page must stop after page_size + 1 sessions, leaving the capped count as the only full pass"
)
def test_session_view_repeated_page_stays_bounded_under_the_generic_plan(gateway: Gateway, seeded: str) -> None:
"""The proxy runs these as prepared statements and Postgres switches to a generic plan after five executions, so a plan that only looks good with literal values shows up here."""
for _ in range(19):
_session_page(gateway, page="1")
before: Final = _tuples_read_quiet()
_session_page(gateway, page="1")
tuples_read: Final = _tuples_read_quiet() - before
assert tuples_read < MAX_TUPLES_PER_PAGE, (
f"the 20th page read {tuples_read} tuples for a window of {SEED_ROWS} rows and page_size {PAGE_SIZE}"
)
def test_session_view_uses_newest_row_inside_the_window(gateway: Gateway, seeded: str) -> None:
page: Final = _session_page(gateway, page="1", session_id=f"{seeded}-straddle")
assert page["total"] == 1
data: Final = _data_rows(page)
assert [row["request_id"] for row in data] == [f"intg-straddle-{seeded}-in"]
def test_session_view_filtered_first_page_stays_bounded(gateway: Gateway, seeded: str) -> None:
"""A filter on a column no index covers is where the planner once drove every per-row probe off the wrong index."""
before: Final = _tuples_read_quiet()
page: Final = _session_page(gateway, page="1", key_alias=KEY_ALIAS)
tuples_read: Final = _tuples_read_quiet() - before
assert page["total"] == EXPECTED_SESSIONS
assert _data_rows(page)[0]["request_id"] == f"intg-sesswin-{seeded}-m1c"
assert tuples_read < MAX_TUPLES_PER_PAGE, (
f"filtered first page read {tuples_read} tuples for a window of {SEED_ROWS} rows and page_size {PAGE_SIZE}"
)
MANY_SESSIONS_ROWS: Final = 150_000
MANY_SESSIONS_LEN: Final = 2
_MANY_BASE: Final = datetime(2003, 3, 1)
MANY_WINDOW_START: Final = _MANY_BASE.strftime("%Y-%m-%d %H:%M:%S")
MANY_WINDOW_END: Final = (_MANY_BASE + timedelta(seconds=MANY_SESSIONS_ROWS + 1)).strftime("%Y-%m-%d %H:%M:%S")
@pytest.fixture
def many_sessions(gateway: Gateway) -> Iterator[str]:
marker: Final = uuid.uuid4().hex
with gateway.scenario() as scenario:
scenario.cleanups.callback(_cleanup, marker)
_settle_autovacuum()
with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as connection:
connection.execute(
'INSERT INTO "LiteLLM_SpendLogs" ('
'request_id, call_type, api_key, session_id, "startTime", "endTime", model, metadata, messages, '
"response) "
"SELECT %s || n, 'acompletion', 'K', %s || '-s-' || (n / %s), ts, ts, 'e2e-sesswin', '{}'::jsonb, "
"'{}'::jsonb, '{}'::jsonb "
"FROM generate_series(1, %s) AS g(n) "
"CROSS JOIN LATERAL (SELECT %s::timestamp + (n || ' seconds')::interval AS ts) AS t",
(f"intg-sesswin-{marker}-", marker, MANY_SESSIONS_LEN, MANY_SESSIONS_ROWS, MANY_WINDOW_START),
)
connection.execute('VACUUM ANALYZE "LiteLLM_SpendLogs"')
_wait_for_stats_quiet()
yield marker
def test_session_view_count_stops_at_the_cap_on_a_window_of_many_sessions(gateway: Gateway, many_sessions: str) -> None:
"""A window holding far more sessions than the count cap is counted from its newest rows, not grouped whole."""
assert MANY_SESSIONS_ROWS // MANY_SESSIONS_LEN > SPEND_LOGS_PAGINATION_COUNT_CAP
before: Final = _tuples_read_quiet()
page: Final = gateway.get(
"/spend/logs/ui",
params={
"group_by_session": "true",
"start_date": MANY_WINDOW_START,
"end_date": MANY_WINDOW_END,
"sort_by": "startTime",
"sort_order": "desc",
"page_size": str(PAGE_SIZE),
"page": "1",
},
)
tuples_read: Final = _tuples_read_quiet() - before
assert page["total"] == SPEND_LOGS_PAGINATION_COUNT_CAP
assert page["total_is_capped"] is True
assert _data_rows(page)[0]["request_id"] == f"intg-sesswin-{many_sessions}-{MANY_SESSIONS_ROWS}"
assert tuples_read < MANY_SESSIONS_ROWS // 2, (
f"the page read {tuples_read} tuples of a {MANY_SESSIONS_ROWS} row window; the page and the capped count "
"must both stop long before the end of a window with more sessions than the cap"
)
def _small_rows(marker: str) -> tuple[_SeedRow, ...]:
def at(seconds: int) -> str:
return (_SMALL_BASE + timedelta(seconds=seconds)).strftime("%Y-%m-%d %H:%M:%S")
return (
_SeedRow(f"{marker}-adopter", f"{marker}-root", at(10)),
_SeedRow(f"{marker}-root", None, at(20)),
_SeedRow(f"{marker}-adopter2", f"{marker}-root2", at(30)),
_SeedRow(f"{marker}-root2", "", at(40)),
_SeedRow(f"{marker}-shared-a", f"{marker}-shared", at(50), api_key=f"{marker}-key-a"),
_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}-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"),
)
@pytest.fixture
def small_seed(gateway: Gateway) -> Iterator[str]:
marker: Final = f"intg-sesssmall-{uuid.uuid4().hex}"
with gateway.scenario() as scenario:
scenario.cleanups.callback(_delete_small_rows, marker)
_insert_rows(_small_rows(marker))
yield marker
def _delete_small_rows(marker: str) -> None:
with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as connection:
connection.execute('DELETE FROM "LiteLLM_SpendLogs" WHERE request_id LIKE %s', (f"{marker}-%",))
def _small_page(gateway: Gateway, **extra: str) -> dict[str, JsonValue]:
return gateway.get(
"/spend/logs/ui",
params={
"group_by_session": "true",
"start_date": SMALL_WINDOW_START,
"end_date": SMALL_WINDOW_END,
"sort_by": "startTime",
**extra,
},
)
def _sessions(page: Mapping[str, JsonValue]) -> list[tuple[str, str]]:
return [(_tie_as_one(string_value(row["request_id"])), string_value(row["api_key"])) for row in _data_rows(page)]
def _tie_as_one(request_id: str) -> str:
return request_id.removesuffix("-1").removesuffix("-2") if "-tie-" in request_id else request_id
def _newest_first(marker: str) -> list[tuple[str, str]]:
return [
(f"{marker}-shared", f"{marker}-key-b"),
(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"),
(f"{marker}-root2", "K"),
(f"{marker}-root", "K"),
]
def test_session_view_lists_each_session_once(gateway: Gateway, small_seed: str) -> None:
page: Final = _small_page(gateway, sort_order="desc", page_size="50", page="1")
assert _sessions(page) == _newest_first(small_seed), (
"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"
)
assert page["total"] == 7
assert page["has_more"] is False
def test_session_view_cursor_walk_counts_sessions_per_api_key(gateway: Gateway, small_seed: str) -> None:
first: Final = _small_page(gateway, sort_order="desc", page_size="2", page="1")
cursor_pages: Final = [first]
while cursor_pages[-1]["has_more"] is True:
cursor_pages.append(
_small_page(
gateway,
sort_order="desc",
page_size="2",
page=str(len(cursor_pages) + 1),
session_cursor=string_value(cursor_pages[-1]["next_session_cursor"]),
)
)
walked: Final = [session for page in cursor_pages for session in _sessions(page)]
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(
gateway: Gateway, small_seed: str
) -> None:
ascending: Final = _small_page(gateway, sort_order="asc", page_size="50", page="1")
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"] == 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
@pytest.mark.parametrize("identities", ((("a", "K"), ("B", "K")), (("shared", "a"), ("shared", "B"))))
def test_session_view_cursor_pages_preserve_database_order_for_tied_sessions(
gateway: Gateway, identities: tuple[tuple[str, str], tuple[str, str]]
) -> None:
marker: Final = f"intg-sesscollation-{uuid.uuid4().hex}"
rows: Final = tuple(
_SeedRow(f"{marker}-{index}", f"{marker}-{session}", SMALL_WINDOW_START, api_key=key)
for index, (session, key) in enumerate(identities)
)
with gateway.scenario() as scenario:
scenario.cleanups.callback(_delete_small_rows, marker)
_insert_rows(rows)
expected: Final = tuple(
(string_value(row["request_id"]), string_value(row["api_key"]))
for row in read_rows(
'SELECT request_id, api_key FROM "LiteLLM_SpendLogs" WHERE request_id LIKE %s '
'ORDER BY "startTime" DESC, session_id DESC, api_key DESC',
(f"{marker}-%",),
)
)
first: Final = _small_page(gateway, session_id=marker, page_size="1", page="1")
second: Final = _small_page(
gateway,
session_id=marker,
page_size="1",
page="2",
session_cursor=string_value(first["next_session_cursor"]),
)
walked: Final = tuple(
(string_value(row["request_id"]), string_value(row["api_key"]))
for row in _data_rows(first) + _data_rows(second)
)
assert walked == expected, "cursor pages must match database ordering without omissions or duplicates"
assert first["total"] == second["total"] == len(rows)
assert first["has_more"] is True
assert second["has_more"] is False

278
tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py
View file

@ -7173,16 +7173,46 @@ def test_parse_session_cursor():
)
SESSION_GROUP_KEY_SQL = "COALESCE(NULLIF(session_id, ''), request_id), api_key"
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_count_walk(sessions):
"""Count-walk result for a window that ran out of rows after holding ``sessions`` sessions."""
return [{"rows_read": 0, "sessions": sessions}]
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 "newest_rows" in sql_query:
return _session_count_walk(session_total)
if "COUNT(*) AS total_count" in sql_query:
return [{"total_count": session_total}]
if "DISTINCT ON" in sql_query:
if "AS rep" in sql_query or "DISTINCT ON" in sql_query:
return representative_rows
if "COALESCE(SUM(spend)" in sql_query:
return [
@ -7232,9 +7262,15 @@ 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 "newest_rows" in sql_query:
return _session_count_walk(len(sessions) if counted_total is None else counted_total)
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 "DISTINCT ON" in sql_query:
if "AS rep" in sql_query or "DISTINCT ON" 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 []
@ -7294,32 +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 f"GROUP BY {SESSION_GROUP_KEY_SQL}" in page_query_sql
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 f"GROUP BY {SESSION_GROUP_KEY_SQL}" in count_sql
assert "LIMIT" in count_sql and "FROM (" in count_sql, "the grouped count must stay bounded"
rep_call = emitted[2]
assert f"DISTINCT ON ({SESSION_GROUP_KEY_SQL})" in rep_call[0]
assert (
f"ORDER BY {SESSION_GROUP_KEY_SQL}, " + spend_management_endpoints._SESSION_REPRESENTATIVE_ORDER_SQL
) in rep_call[0]
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 HAVING 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)
@ -7349,21 +7366,214 @@ 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 f'HAVING (MAX("startTime"), {SESSION_GROUP_KEY_SQL}) <' 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 (MAX(startTime), session key, api_key)"
finally:
app.dependency_overrides.pop(ps.user_api_key_auth, None)
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 "WITH batch AS MATERIALIZED" in sql_query:
return _session_walk_batch(walked_rows, min(window_rows_per_batch, params[-1]))
if "newest_rows" in sql_query:
return _session_count_walk(3)
if "COUNT(*) AS total_count" in sql_query:
return [{"total_count": 3}]
if "AS rep" 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 []
return grouped_rows
mock_prisma = MagicMock()
mock_prisma.db = MagicMock()
mock_prisma.db.query_raw = AsyncMock(side_effect=mock_query_raw)
return mock_prisma
@pytest.mark.asyncio
@pytest.mark.parametrize(
("window_rows_per_batch", "expected_request_ids", "expected_has_more"),
[
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_cannot_settle(
client, monkeypatch, window_rows_per_batch, expected_request_ids, expected_has_more
):
"""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")],
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"),
_session_page_row("grouped-3", "2026-08-29 08:00:00"),
],
)
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": 2},
headers={"Authorization": "Bearer sk-test"},
)
assert response.status_code == 200, response.text
data = response.json()
assert [row["request_id"] for row in data["data"]] == expected_request_ids
assert data["has_more"] is expected_has_more
finally:
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 "AS heads(session_key, api_key, last_activity)" in sql_query:
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("a", tied),
_walk_row("a", tied, is_edge=True, batch_rows=params[-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:
return [{"total_count": 3}]
if "AS rep" 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-Z"]
assert data["next_session_cursor"] == f"{tied}|hashed-key|Z"
finally:
app.dependency_overrides.pop(ps.user_api_key_auth, None)
def _session_window(session_count, rows_per_session, newer_filtered_out_rows=0):
"""Session keys of a window's rows, newest first, each session's rows back to back; None marks a newer row the
request's filters exclude."""
return [None] * newer_filtered_out_rows + [
f"sess-{index}" for index in range(session_count) for _ in range(rows_per_session)
]
def _counting_prisma(window):
"""Mock prisma answering the count walk and the grouped count from ``window``, recording the rows each walk read."""
rows_read_per_walk = []
async def mock_query_raw(sql_query, *params):
if "newest_rows" in sql_query:
newest = window[: params[-1]]
rows_read_per_walk.append(len(newest))
return [{"rows_read": len(newest), "sessions": len(set(newest) - {None})}]
if "COUNT(*) AS total_count" in sql_query:
return [{"total_count": min(len(set(window) - {None}), params[-1])}]
raise AssertionError(f"unexpected query: {sql_query}")
mock_prisma = MagicMock()
mock_prisma.db = MagicMock()
mock_prisma.db.query_raw = AsyncMock(side_effect=mock_query_raw)
return mock_prisma, rows_read_per_walk
_COUNT_CAP = spend_management_endpoints.SPEND_LOGS_PAGINATION_COUNT_CAP
@pytest.mark.asyncio
@pytest.mark.parametrize(
("window", "expected_total", "walk_settles", "walk_reads_every_row"),
[
pytest.param(_session_window(30000, 5), (_COUNT_CAP, True), True, False, id="many-sessions-stop-at-the-cap"),
pytest.param(
_session_window(_COUNT_CAP + 2000, 5), (_COUNT_CAP, True), True, True, id="cap-passed-after-the-first-batch"
),
pytest.param(_session_window(_COUNT_CAP, 5), (_COUNT_CAP, False), True, True, id="exactly-the-cap-is-exact"),
pytest.param(_session_window(1000, 5), (1000, False), True, True, id="small-window-is-counted-exactly"),
pytest.param(_session_window(20, 20000), (20, False), False, False, id="long-sessions-fall-back-to-grouping"),
pytest.param(
_session_window(10, 5, newer_filtered_out_rows=400000),
(10, False),
False,
False,
id="filter-matching-no-recent-rows-falls-back-at-once",
),
],
)
async def test_session_count_matches_the_grouped_count_and_stops_early(
window, expected_total, walk_settles, walk_reads_every_row
):
"""The newest-first count gives the grouped count's answer, reading only part of a window of many sessions."""
mock_prisma, rows_read_per_walk = _counting_prisma(window)
counted = await spend_management_endpoints._count_sessions_newest_first(mock_prisma, "TRUE", "TRUE", (), 1)
total = (
counted
if counted is not None
else await spend_management_endpoints._count_grouped_sessions(mock_prisma, "TRUE", (), 1)
)
assert total == expected_total
assert (counted is not None) is walk_settles
assert (max(rows_read_per_walk) == len(window)) is walk_reads_every_row
assert max(rows_read_per_walk) <= spend_management_endpoints._SESSION_COUNT_MAX_ROWS
if not walk_settles:
assert rows_read_per_walk == [spend_management_endpoints._SESSION_COUNT_FIRST_BATCH_ROWS], (
"a walk that cannot reach the cap within its row budget must stop after its first batch"
)
@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."""

43
tests/unit/proxy/spend_tracking/test_spend_query_optimization.py
View file

@ -510,20 +510,15 @@ async def test_global_spend_report_team_group_forwards_team_id(monkeypatch):
async def test_spend_logs_ui_group_by_session_paginates_sessions(monkeypatch):
"""
With group_by_session=true, /spend/logs/ui must page and count SESSIONS,
not raw calls: the page is one representative row per session (DISTINCT
ON the session group key, preferring non-MCP calls, newest first), the
bounded count counts groups, and on the default startTime sort the page is
not raw calls: the page is one representative row per session, newest
session first, the count stays bounded, and on the default startTime sort the page is
selected by keyset (no OFFSET) so deep pages do not degrade. Otherwise the
UI collapses a server page of N calls into fewer visible rows while the
footer still claims N (issue #38060).
"""
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.spend_tracking.spend_management_endpoints import (
SPEND_LOGS_PAGINATION_COUNT_CAP,
ui_view_spend_logs,
)
from litellm.proxy.spend_tracking.spend_management_endpoints import ui_view_spend_logs
group_key = "COALESCE(NULLIF(session_id, ''), request_id), api_key"
session_rows = [
{"session_key": f"req-{index}", "api_key": "k", "last_activity": f"2026-02-16 10:{59 - index:02d}:00"}
for index in range(51)
@ -534,9 +529,17 @@ 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 "newest_rows" in sql_query:
return [{"rows_read": 0, "sessions": 60}]
if "COUNT(*) AS total_count" in sql_query:
return [{"total_count": 60}]
if "AS session_representatives" in sql_query:
if "AS rep" in sql_query or "DISTINCT ON" in sql_query:
return representative_rows
return session_rows
@ -567,28 +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 f"GROUP BY {group_key}" in page_sql, f"page must select sessions, not calls. SQL was:\n{page_sql}"
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]
assert f"GROUP BY {group_key}" in count_sql, f"grouped total must count sessions. SQL was:\n{count_sql}"
assert "COUNT(*) OVER ()" not in count_sql
assert "LIMIT" in count_sql and "FROM (" in count_sql, "the grouped count must stay bounded"
assert count_call[-1] == SPEND_LOGS_PAGINATION_COUNT_CAP + 1
rep_sql = emitted[2][0]
assert f"DISTINCT ON ({group_key})" in rep_sql, f"page must return one row per session. SQL was:\n{rep_sql}"
assert (
f"ORDER BY {group_key}, (call_type = 'asend_message') DESC, "
"CASE WHEN call_type = 'asend_message' THEN \"endTime\" END DESC NULLS LAST, "
"call_type IN ('call_mcp_tool', 'list_mcp_tools'), \"startTime\" DESC"
) in rep_sql, "the session representative must prefer the final agent outcome, then the newest non-MCP call"
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