scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-11 08:38:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 280

fix(proxy): open the Logs drawer by request_id without scanning for litellm_call_id (#50789)

Browse source 
* fix(proxy): open the Logs drawer by request_id without scanning for litellm_call_id

The detail lookup matched request_id OR litellm_call_id, which Postgres runs as a full table scan when the litellm_call_id index is missing or invalid. Read by request_id, the primary key, first and by litellm_call_id only on a miss, both under the spend log statement timeout.

* test(proxy): type the call ID index test helpers

* fix(proxy): limit the Logs drawer call ID lookup to start_date and answer 400 for a malformed date

* fix(proxy): let spend log reads wait for a busy connection like a plain query

* test(proxy): let the session pages read double take max_wait

---------

Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
tavren27 2026-11-10 09:03:21 -08:00 • committed by GitHub
parent 82ef50bacb
commit d70855ca82
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 445 additions and 64 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -132,11 +132,6 @@ class _SupportsModelDump(Protocol):
def model_dump(self) -> Mapping[str, object]: ...
class _SpendLogOwnerRow(TypedDict):
user: ReadOnly[str | None]
team_id: ReadOnly[str | None]
class _ActivityRow(TypedDict):
date: str
api_requests: int
@ -273,6 +268,60 @@ async def _query_spend_logs(prisma_client: PrismaClient, sql_query: str, *args:
)
_SpendLogIdColumn: TypeAlias = Literal["request_id", "litellm_call_id"]
@dataclass(frozen=True, slots=True)
class _SpendLogWindow:
start: datetime | None = None
end: datetime | None = None
_UNBOUNDED_SPEND_LOG_WINDOW: Final = _SpendLogWindow()
def _spend_log_id_filter(
id_column: _SpendLogIdColumn, lookup_id: str, window: _SpendLogWindow
) -> tuple[str, tuple[object, ...]]:
match id_column:
case "request_id":
return "request_id = $1", (lookup_id,)
case "litellm_call_id":
bounds: Final = tuple(
(operator, bound.isoformat())
for operator, bound in ((">=", window.start), ("<=", window.end))
if bound is not None
)
bound_sql: Final = "".join(
f""" AND "startTime" {operator} (${index}::timestamptz AT TIME ZONE 'UTC')"""
for index, (operator, _) in enumerate(bounds, start=2)
)
return f"litellm_call_id = $1{bound_sql}", (lookup_id, *(bound for _, bound in bounds))
def _spend_log_lookup_timeout_hint(id_column: _SpendLogIdColumn) -> str:
match id_column:
case "request_id":
return "Try again"
case "litellm_call_id":
return (
"Pass a narrower start_date and end_date, or ask an admin to check that the "
"LiteLLM_SpendLogs_litellm_call_id_idx index exists and is valid"
)
async def _query_spend_logs_by_id(
prisma_client: PrismaClient, id_column: _SpendLogIdColumn, sql_query: str, *args: object
) -> Sequence[Mapping[str, object]]:
return await query_raw_within_statement_timeout(
prisma_client,
sql_query,
*args,
what=f"Spend log lookup by {id_column}",
hint=_spend_log_lookup_timeout_hint(id_column),
)
@dataclass(frozen=True, slots=True)
class _SpendLogsReader:
prisma_client: PrismaClient
@ -399,9 +448,11 @@ def _request_id_or_call_id_clause(request_id: str) -> tuple[_RequestIdEquals, _L
return (request_id_clause, call_id_clause)
async def _find_spend_log_owners(prisma_client: PrismaClient, request_id: str) -> Sequence[_SpendLogOwnerRow]:
async def _find_spend_log_owners(
prisma_client: PrismaClient, id_column: _SpendLogIdColumn, lookup_id: str, window: _SpendLogWindow
) -> frozenset[tuple[str | None, str | None]]:
"""Read the distinct ``(user, team_id)`` owner pairs across every spend log row
identified by ``request_id`` or ``litellm_call_id``.
whose ``id_column`` is ``lookup_id``.
``litellm_call_id`` is populated from the client-settable ``x-litellm-call-id``
request header, so it is not guaranteed unique to one tenant: any number of rows
@ -409,13 +460,14 @@ async def _find_spend_log_owners(prisma_client: PrismaClient, request_id: str) -
carrying the caller's id could otherwise push the caller's own owner pair past a
row-sample cap and lock them out of their own lookup.
"""
sql_query: Final = """
SELECT DISTINCT "user", team_id
FROM "LiteLLM_SpendLogs"
WHERE request_id = $1 OR litellm_call_id = $1
"""
owners: Final[Sequence[_SpendLogOwnerRow] | None] = await _query_raw_or_none(prisma_client, sql_query, request_id)
return owners if owners is not None else ()
id_filter, id_params = _spend_log_id_filter(id_column, lookup_id, window)
rows: Final = await _query_spend_logs_by_id(
prisma_client,
id_column,
f'SELECT DISTINCT "user", team_id FROM "LiteLLM_SpendLogs" WHERE {id_filter}',
*id_params,
)
return frozenset(_fetched_row_owner(row) for row in rows)
async def _count_spend_logs(prisma_client: PrismaClient, where: Mapping[str, object]) -> int:
@ -3292,6 +3344,18 @@ async def _resolve_request_response_payload(
)
def _parse_spend_log_ui_date(value: str | None) -> datetime | None:
if value is None:
return None
try:
return datetime.strptime(value, "%Y-%m-%d %H:%M:%S").replace(tzinfo=timezone.utc)
except ValueError as error:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail={"error": f"Invalid date {value!r}, expected YYYY-MM-DD HH:MM:SS in UTC"},
) from error
@router.get(
"/spend/logs/ui/{request_id}",
tags=["Budget & Spend Tracking"],
@ -3319,6 +3383,7 @@ async def ui_view_request_response_for_request_id(
"""
from litellm.proxy.proxy_server import prisma_client
window: Final = _SpendLogWindow(start=_parse_spend_log_ui_date(start_date), end=_parse_spend_log_ui_date(end_date))
caller_is_admin: Final = is_admin_view_safe(user_api_key_dict=user_api_key_dict)
if not caller_is_admin:
if prisma_client is None:
@ -3335,16 +3400,10 @@ async def ui_view_request_response_for_request_id(
prisma_client=prisma_client,
user_api_key_dict=user_api_key_dict,
request_id=request_id,
window=window,
)
custom_loggers: Final = litellm.logging_callback_manager.get_active_additional_logging_utils_from_custom_logger()
start_date_obj: datetime | None = None
end_date_obj: datetime | None = None
if start_date is not None:
start_date_obj = datetime.strptime(start_date, "%Y-%m-%d %H:%M:%S").replace(tzinfo=timezone.utc)
if end_date is not None:
end_date_obj = datetime.strptime(end_date, "%Y-%m-%d %H:%M:%S").replace(tzinfo=timezone.utc)
spend_log_row: Final = (
None
if prisma_client is None
@ -3354,6 +3413,7 @@ async def ui_view_request_response_for_request_id(
request_id=request_id,
caller_is_admin=caller_is_admin,
log_team_lookup=log_team_lookup,
window=window,
)
)
stored_request_id: Final = _stored_request_id(spend_log_row, request_id)
@ -3361,8 +3421,8 @@ async def ui_view_request_response_for_request_id(
for custom_logger in custom_loggers:
payload = await custom_logger.get_request_response_payload(
request_id=stored_request_id,
start_time_utc=start_date_obj,
end_time_utc=end_date_obj,
start_time_utc=window.start,
end_time_utc=window.end,
)
if payload is not None:
if not caller_is_admin and prisma_client is not None:
@ -4877,6 +4937,7 @@ async def _assert_user_can_view_request_id(
prisma_client: PrismaClient,
user_api_key_dict: UserAPIKeyAuth,
request_id: str,
window: _SpendLogWindow = _UNBOUNDED_SPEND_LOG_WINDOW,
) -> None:
"""
Verify the requesting non-admin user is allowed to view at least one spend-log
@ -4889,11 +4950,24 @@ async def _assert_user_can_view_request_id(
so a missing row can't be used to read a payload out of cold storage via the
detail endpoint.
"""
owners: Final = await _find_spend_log_owners(prisma_client, request_id)
for owner in owners:
if await _user_can_view_spend_log_owner(prisma_client, user_api_key_dict, owner["user"], owner["team_id"]):
return
raise _spend_log_forbidden(request_id)
if not (
await _user_can_view_spend_log_by_id(prisma_client, user_api_key_dict, "request_id", request_id, window)
or await _user_can_view_spend_log_by_id(prisma_client, user_api_key_dict, "litellm_call_id", request_id, window)
):
raise _spend_log_forbidden(request_id)
async def _user_can_view_spend_log_by_id(
prisma_client: PrismaClient,
user_api_key_dict: UserAPIKeyAuth,
id_column: _SpendLogIdColumn,
lookup_id: str,
window: _SpendLogWindow,
) -> bool:
for user, team_id in await _find_spend_log_owners(prisma_client, id_column, lookup_id, window):
if await _user_can_view_spend_log_owner(prisma_client, user_api_key_dict, user, team_id):
return True
return False
async def _spend_log_read_scope(user_api_key_dict: UserAPIKeyAuth, log_team_lookup: LogTeamLookup) -> OwnedRows:
@ -4927,32 +5001,54 @@ def _read_scope_where(scope: ReadScope) -> Mapping[str, object]:
return grants[0] if len(grants) == 1 else {"OR": list(grants)}
def _spend_log_payload_query(request_id: str, scope: ReadScope) -> tuple[str, tuple[object, ...]]:
def _spend_log_payload_query(
id_column: _SpendLogIdColumn, lookup_id: str, scope: ReadScope, window: _SpendLogWindow
) -> tuple[str, tuple[object, ...]]:
"""
Fetch the one row an id lookup resolves to, preferring the exact ``request_id``
match over rows that merely carry the id as their client-set ``litellm_call_id``.
A non-admin viewer only ever gets rows they own or rows of a team they may view.
Fetch one row whose ``id_column`` is ``lookup_id``. A non-admin viewer only ever
gets rows they own or rows of a team they may view.
"""
scope_clause, scope_params = read_scope_sql(scope, 2)
id_filter, id_params = _spend_log_id_filter(id_column, lookup_id, window)
scope_clause, scope_params = read_scope_sql(scope, len(id_params) + 1)
scope_sql: Final = f" AND {scope_clause}" if scope_clause else ""
return (
f"""
SELECT request_id, messages, response, proxy_server_request, metadata, "user", team_id
FROM "LiteLLM_SpendLogs"
WHERE (request_id = $1 OR litellm_call_id = $1){scope_sql}
ORDER BY (request_id = $1) DESC
WHERE {id_filter}{scope_sql}
LIMIT 1
""",
(request_id, *scope_params),
(*id_params, *scope_params),
)
async def _query_spend_log_payload(
prisma_client: PrismaClient,
id_column: _SpendLogIdColumn,
lookup_id: str,
scope: ReadScope,
window: _SpendLogWindow,
) -> Sequence[Mapping[str, object]]:
sql_query, sql_params = _spend_log_payload_query(id_column, lookup_id, scope, window)
return await _query_spend_logs_by_id(prisma_client, id_column, sql_query, *sql_params)
async def _fetch_spend_log_payload_rows(
prisma_client: PrismaClient, lookup_id: str, scope: ReadScope, window: _SpendLogWindow
) -> Sequence[Mapping[str, object]]:
by_request_id: Final = await _query_spend_log_payload(prisma_client, "request_id", lookup_id, scope, window)
if by_request_id:
return by_request_id
return await _query_spend_log_payload(prisma_client, "litellm_call_id", lookup_id, scope, window)
async def _resolve_spend_log_payload_row(
prisma_client: PrismaClient,
user_api_key_dict: UserAPIKeyAuth,
request_id: str,
caller_is_admin: bool,
log_team_lookup: LogTeamLookup,
window: _SpendLogWindow,
) -> Mapping[str, object] | None:
"""
Resolve an id lookup to the caller's own spend-log row before any payload
@ -4962,8 +5058,7 @@ async def _resolve_spend_log_payload_row(
``request_id`` is the key that names the caller's own request.
"""
scope: Final = AllRows() if caller_is_admin else await _spend_log_read_scope(user_api_key_dict, log_team_lookup)
sql_query, sql_params = _spend_log_payload_query(request_id, scope)
rows: Final[Sequence[Mapping[str, object]] | None] = await _query_raw_or_none(prisma_client, sql_query, *sql_params)
rows: Final = await _fetch_spend_log_payload_rows(prisma_client, request_id, scope, window)
if not rows:
return None
if not caller_is_admin:

9
litellm/proxy/utils.py
View file

@ -5367,16 +5367,19 @@ async def query_raw_within_statement_timeout(
sql: str,
*params: object,
what: str = "Spend logs",
hint: str = "Pick a shorter date range",
timeout: timedelta = _SPEND_LOG_READ_STATEMENT_TIMEOUT,
) -> Sequence[Mapping[str, object]]:
"""One raw read on the read target that Postgres cancels at ``timeout``.
Prisma's HTTP client gives up on a statement after 30s anyway, so without
this the Admin UI's error is the same but Postgres keeps running the scan.
A cancelled read surfaces as a 503 asking for a narrower date range.
A cancelled read surfaces as a 503 carrying ``hint``.
"""
try:
async with prisma_client.read_db.tx(timeout=timeout + _STATEMENT_TIMEOUT_TRANSACTION_GRACE) as transaction:
async with prisma_client.read_db.tx(
max_wait=timeout, timeout=timeout + _STATEMENT_TIMEOUT_TRANSACTION_GRACE
) as transaction:
await transaction.execute_raw(f"SET LOCAL statement_timeout = {int(timeout / timedelta(milliseconds=1))}")
return await transaction.query_raw(sql, *params)
except Exception as e:
@ -5384,7 +5387,7 @@ async def query_raw_within_statement_timeout(
raise
raise HTTPException(
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
detail={"error": f"{what} took longer than {int(timeout.total_seconds())}s. Pick a shorter date range"},
detail={"error": f"{what} took longer than {int(timeout.total_seconds())}s. {hint}"},
) from e

4
tests/integration/spend/test_session_grouped_logs_pages.py
View file

@ -320,8 +320,8 @@ class _ExplainedReads:
self.rows_scanned: tuple[int, ...] = ()
@asynccontextmanager
async def tx(self, *, timeout: timedelta) -> AsyncIterator["_ExplainedReads"]:
async with self.database.tx(timeout=timeout) as transaction:
async def tx(self, *, max_wait: timedelta, timeout: timedelta) -> AsyncIterator["_ExplainedReads"]:
async with self.database.tx(max_wait=max_wait, timeout=timeout) as transaction:
reads: Final = _ExplainedReads(transaction)
try:
yield reads

126
tests/integration/spend/test_spend_log_read_scope.py
View file

@ -1,10 +1,12 @@
import asyncio
import json
import os
import uuid
from collections.abc import AsyncIterator
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from hashlib import sha256
from typing import Final
from typing import Final, Literal
from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
import psycopg
@ -18,7 +20,13 @@ from psycopg.types.json import Jsonb
from pydantic import TypeAdapter
from litellm.proxy.auth.authorization import AllRows, OwnedRows, ReadScope
from litellm.proxy.spend_tracking.spend_management_endpoints import _spend_log_payload_query, read_scope_sql
from litellm.proxy.db.prisma_client import PrismaWrapper
from litellm.proxy.spend_tracking.spend_management_endpoints import (
_spend_log_payload_query,
_SpendLogWindow,
read_scope_sql,
)
from litellm.proxy.utils import query_raw_within_statement_timeout
@dataclass(frozen=True, slots=True)
@ -146,17 +154,121 @@ async def test_ownership_sql_selects_allowed_rows_and_intersects_filters(
@pytest.mark.asyncio
@pytest.mark.parametrize(
("scope", "expected"),
[(AllRows(), ("foreign",)), (OwnedRows("caller"), ("own",)), (OwnedRows(None), ())],
("id_column", "scope", "expected"),
[
("request_id", AllRows(), ("foreign",)),
("request_id", OwnedRows("caller"), ()),
("litellm_call_id", AllRows(), ("own",)),
("litellm_call_id", OwnedRows("caller"), ("own",)),
("litellm_call_id", OwnedRows(None), ()),
],
)
async def test_payload_sql_filters_foreign_collisions_and_prefers_exact_ids_for_admins(
spend_database: Prisma, scope: ReadScope, expected: tuple[str, ...]
async def test_payload_sql_filters_foreign_collisions(
spend_database: Prisma,
id_column: Literal["request_id", "litellm_call_id"],
scope: ReadScope,
expected: tuple[str, ...],
) -> None:
query, params = _spend_log_payload_query("foreign", scope)
query, params = _spend_log_payload_query(id_column, "foreign", scope, _SpendLogWindow())
result: Final = await spend_database.query_raw(query, *params)
assert tuple(row.request_id for row in REQUEST_IDS.validate_python(result)) == expected
SEEDED_AT: Final = datetime(2026, 1, 1, tzinfo=timezone.utc)
@pytest.mark.asyncio
@pytest.mark.parametrize("scope", [AllRows(), OwnedRows("caller", ("first",))])
@pytest.mark.parametrize(
("window", "expected"),
[
(_SpendLogWindow(start=SEEDED_AT - timedelta(seconds=1)), ("own",)),
(_SpendLogWindow(start=SEEDED_AT, end=SEEDED_AT), ("own",)),
(_SpendLogWindow(start=SEEDED_AT + timedelta(seconds=1)), ()),
(_SpendLogWindow(end=SEEDED_AT - timedelta(seconds=1)), ()),
],
)
async def test_payload_sql_by_call_id_keeps_to_the_window(
spend_database: Prisma, scope: ReadScope, window: _SpendLogWindow, expected: tuple[str, ...]
) -> None:
query, params = _spend_log_payload_query("litellm_call_id", "foreign", scope, window)
result: Final = await spend_database.query_raw(query, *params)
assert tuple(row.request_id for row in REQUEST_IDS.validate_python(result)) == expected
async def _drop_call_id_indexes(spend_database: Prisma) -> None:
await spend_database.execute_raw(
"""
DO $$
DECLARE index_name text;
BEGIN
FOR index_name IN
SELECT indexname FROM pg_indexes
WHERE schemaname = current_schema()
AND tablename = 'LiteLLM_SpendLogs'
AND indexdef LIKE '%litellm_call_id%'
LOOP
EXECUTE format('DROP INDEX %I', index_name);
END LOOP;
END $$
"""
)
async def _plan_without_seq_scans(spend_database: Prisma, query: str, params: tuple[object, ...]) -> str:
async with spend_database.tx() as transaction:
await transaction.execute_raw("SET LOCAL enable_seqscan = off")
return json.dumps(await transaction.query_raw(f"EXPLAIN (FORMAT JSON) {query}", *params))
@pytest.mark.asyncio
@pytest.mark.parametrize("scope", [AllRows(), OwnedRows("caller", ("first",))])
@pytest.mark.parametrize(
("id_column", "lookup_id", "window"),
[
("request_id", "own", _SpendLogWindow()),
("litellm_call_id", "foreign", _SpendLogWindow(start=SEEDED_AT)),
],
)
async def test_payload_sql_reads_through_an_index_without_the_call_id_index(
spend_database: Prisma,
scope: ReadScope,
id_column: Literal["request_id", "litellm_call_id"],
lookup_id: str,
window: _SpendLogWindow,
) -> None:
await _drop_call_id_indexes(spend_database)
query, params = _spend_log_payload_query(id_column, lookup_id, scope, window)
plan: Final = await _plan_without_seq_scans(spend_database, query, params)
assert "Seq Scan" not in plan, plan
assert "Index Cond" in plan, plan
@dataclass(frozen=True, slots=True)
class _ReadTarget:
read_db: PrismaWrapper
@pytest.mark.asyncio
async def test_statement_timeout_read_waits_for_a_busy_connection() -> None:
parsed: Final = urlsplit(os.environ["DATABASE_URL"])
one_connection_url: Final = urlunsplit(
parsed._replace(query=urlencode({**dict(parse_qsl(parsed.query)), "connection_limit": 1, "pool_timeout": 60}))
)
database: Final = Prisma(datasource={"url": one_connection_url})
await database.connect()
try:
holder: Final = asyncio.create_task(database.query_raw("SELECT pg_sleep(3)::text AS slept"))
await asyncio.sleep(0.5)
rows: Final = await query_raw_within_statement_timeout(
_ReadTarget(PrismaWrapper(original_prisma=database, iam_token_db_auth=False)), "SELECT 1 AS ready"
)
await holder
finally:
await database.disconnect()
assert rows == [{"ready": 1}]
def _delete_session(session_id: str) -> None:
with psycopg.connect(os.environ["DATABASE_URL"]) as connection:
connection.execute('DELETE FROM "LiteLLM_SpendLogs" WHERE session_id = %s', (session_id,))

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

@ -6,10 +6,11 @@ import hashlib
import json
import re
import sqlite3
from collections.abc import AsyncIterator, Mapping, Sequence
from collections.abc import AsyncIterator, Awaitable, Mapping, Sequence
from datetime import timezone
from functools import partial
from types import SimpleNamespace
from typing import Final, Literal
from typing import Final, Literal, Protocol, TypeAlias
from unittest.mock import AsyncMock, MagicMock, patch
import prisma
@ -469,13 +470,15 @@ def test_can_user_view_spend_log_false_for_other_roles():
def _emulate_spend_log_owner_lookup(rows, sql_query, params):
"""Emulate the ownership lookup SQL over an in-memory spend-log corpus,
matching only the id columns its WHERE clause compares to the lookup id and
honoring DISTINCT and any literal LIMIT the query carries so a capped or
non-distinct query produces the truncated result it would in Postgres."""
lookup_id = params[0]
id_columns = [column for column in ("request_id", "litellm_call_id") if f"{column} = $1" in sql_query]
matches = [
{"user": row.get("user"), "team_id": row.get("team_id")}
for row in rows
if lookup_id in (row.get("request_id"), row.get("litellm_call_id"))
if any(row.get(column) == lookup_id for column in id_columns)
]
if "DISTINCT" in sql_query:
deduped = []
@ -2808,15 +2811,17 @@ async def test_ui_view_request_response_custom_logger_is_keyed_by_callers_own_re
that id as its request_id. The custom logger is asked for the caller's own stored
request_id, so the caller gets their payload rather than a 403 from the foreign
payload's owner check, and the foreign payload is never fetched."""
prisma = MagicMock(
db=MagicMock(
query_raw=AsyncMock(
side_effect=[
[{"user": "other_user", "team_id": None}, {"user": "caller_user", "team_id": None}],
[_payload_row("caller-req", "shared-id", "caller_user", "caller prompt")],
]
)
)
prisma = MagicMock()
statement_timeout_reads(
prisma,
AsyncMock(
side_effect=[
[{"user": "other_user", "team_id": None}],
[{"user": "caller_user", "team_id": None}],
[],
[_payload_row("caller-req", "shared-id", "caller_user", "caller prompt")],
]
),
)
cold_storage = {
"shared-id": {
@ -6649,7 +6654,8 @@ def test_ui_view_request_response_reads_from_cold_storage(client, monkeypatch):
async def _query_raw(_sql, *_args):
return [placeholder_row]
fake_prisma = SimpleNamespace(db=SimpleNamespace(query_raw=_query_raw))
fake_db = SimpleNamespace(query_raw=_query_raw)
fake_prisma = SimpleNamespace(db=fake_db, read_db=ReadTargetOf(fake_db))
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", fake_prisma)
cold_logger = _FakeColdStorageLogger(
@ -7736,6 +7742,8 @@ def _fake_prisma_with_owned_spend_log(owner_user_id, messages_json, response_jso
]
class _Prisma:
read_db = property(lambda self: ReadTargetOf(self.db))
def __init__(self):
self.db = _DB()
@ -7842,7 +7850,8 @@ def test_ui_view_request_response_internal_user_missing_row_forbidden(client, mo
from types import SimpleNamespace
fake_prisma = SimpleNamespace(db=_DB())
fake_db = _DB()
fake_prisma = SimpleNamespace(db=fake_db, read_db=ReadTargetOf(fake_db))
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", fake_prisma)
custom_logger = _RecordingAdditionalLoggingUtils(
@ -7871,6 +7880,168 @@ def test_ui_view_request_response_internal_user_missing_row_forbidden(client, mo
app.dependency_overrides.pop(ps.user_api_key_auth, None)
_FakeSpendLogRows: TypeAlias = Sequence[Mapping[str, object]]
class _FakeRawQuery(Protocol):
def __call__(self, statement_timeout_set: bool, sql_query: str, /, *params: object) -> Awaitable[_FakeSpendLogRows]: ...
def _started_at_or_after(row: Mapping[str, object], start: object) -> bool:
started: Final = row.get("startTime")
return (
isinstance(started, datetime.datetime)
and isinstance(start, str)
and started >= datetime.datetime.fromisoformat(start)
)
class _SpendLogsWithoutCallIdIndex:
"""A spend-log table missing LiteLLM_SpendLogs_litellm_call_id_idx. A read filtering on
litellm_call_id scans the whole table, so it fails the test unless it runs under a
statement timeout, which Postgres cancels it at. One that also bounds "startTime" from
below reads only the rows of that window, through the startTime index."""
def __init__(self, rows: _FakeSpendLogRows) -> None:
self.rows: Final = rows
self.db: Final = SimpleNamespace(query_raw=partial(self._query_raw, False))
@property
def read_db(self) -> SimpleNamespace:
return SimpleNamespace(tx=lambda **_kwargs: _StatementTimeoutTransaction(self._query_raw))
async def _query_raw(self, statement_timeout_set: bool, sql_query: str, *params: object) -> _FakeSpendLogRows:
where: Final = sql_query.split("WHERE", 1)[1]
by_call_id: Final = "litellm_call_id" in where
if by_call_id and not statement_timeout_set:
raise AssertionError("unbounded full scan of LiteLLM_SpendLogs")
if by_call_id and '"startTime"' not in where:
raise statement_timeout_error()
matches: Final = (
[row for row in self.rows if row["litellm_call_id"] == params[0] and _started_at_or_after(row, params[1])]
if by_call_id
else [row for row in self.rows if row["request_id"] == params[0]]
)
if 'SELECT DISTINCT "user", team_id' in sql_query:
return [{"user": row["user"], "team_id": row["team_id"]} for row in matches]
return matches
class _StatementTimeoutTransaction:
def __init__(self, query_raw: _FakeRawQuery) -> None:
self._query_raw: Final = query_raw
self._statement_timeout_set = False
async def execute_raw(self, sql: str) -> int:
self._statement_timeout_set = self._statement_timeout_set or sql.startswith("SET LOCAL statement_timeout")
return 0
async def query_raw(self, sql_query: str, *params: object) -> _FakeSpendLogRows:
return await self._query_raw(self._statement_timeout_set, sql_query, *params)
async def __aenter__(self) -> "_StatementTimeoutTransaction":
return self
async def __aexit__(self, *exc_info: object) -> None:
pass
@pytest.mark.parametrize("user_role", [LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.INTERNAL_USER])
def test_ui_view_request_response_opens_request_id_without_call_id_index(client, monkeypatch, user_role):
"""
Regression: the Logs drawer opens a row by its request_id. Matching litellm_call_id in
the same read made Postgres scan the whole table when its index was missing, which on a
large table ran until the load balancer gave up, so the row must resolve by request_id.
"""
monkeypatch.setattr(
"litellm.proxy.proxy_server.prisma_client",
_SpendLogsWithoutCallIdIndex([_payload_row("req-1", "call-1", "user_a", "drawer prompt")]),
)
monkeypatch.setattr(
litellm.logging_callback_manager, "get_active_additional_logging_utils_from_custom_logger", lambda: []
)
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(user_role=user_role, user_id="user_a")
try:
response = client.get("/spend/logs/ui/req-1", headers={"Authorization": "Bearer sk-test"})
assert response.status_code == 200, response.text
assert response.json()["messages"] == [{"role": "user", "content": "drawer prompt"}]
finally:
app.dependency_overrides.pop(ps.user_api_key_auth, None)
@pytest.mark.parametrize("user_role", [LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.INTERNAL_USER])
def test_ui_view_request_response_call_id_scan_times_out_with_503(client, monkeypatch, user_role):
"""An id only a litellm_call_id carries still needs the scan, which Postgres cancels at
the read statement timeout, so the caller gets a 503 naming the index, not a hang."""
monkeypatch.setattr(
"litellm.proxy.proxy_server.prisma_client",
_SpendLogsWithoutCallIdIndex([_payload_row("req-1", "call-1", "user_a", "drawer prompt")]),
)
monkeypatch.setattr(
litellm.logging_callback_manager, "get_active_additional_logging_utils_from_custom_logger", lambda: []
)
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(user_role=user_role, user_id="user_a")
try:
response = client.get("/spend/logs/ui/call-1", headers={"Authorization": "Bearer sk-test"})
assert response.status_code == 503, response.text
assert "LiteLLM_SpendLogs_litellm_call_id_idx" in response.json()["detail"]["error"]
finally:
app.dependency_overrides.pop(ps.user_api_key_auth, None)
@pytest.mark.parametrize(
("user_role", "start_date", "expected_status", "expected_messages"),
[
(LitellmUserRoles.PROXY_ADMIN, "2026-11-09 11:59:59", 200, [{"role": "user", "content": "drawer prompt"}]),
(LitellmUserRoles.INTERNAL_USER, "2026-11-09 11:59:59", 200, [{"role": "user", "content": "drawer prompt"}]),
(LitellmUserRoles.PROXY_ADMIN, "2026-11-09 12:00:01", 200, None),
(LitellmUserRoles.INTERNAL_USER, "2026-11-09 12:00:01", 403, None),
],
)
def test_ui_view_request_response_reads_call_id_within_start_date_without_call_id_index(
client, monkeypatch, user_role, start_date, expected_status, expected_messages
):
"""
Regression: the drawer passes start_date, so a call ID lookup reads only the rows from
then on instead of scanning the whole table, and a row from before it is not found.
"""
row: Final = {
**_payload_row("req-1", "call-1", "user_a", "drawer prompt"),
"startTime": datetime.datetime(2026, 11, 9, 12, tzinfo=timezone.utc),
}
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", _SpendLogsWithoutCallIdIndex([row]))
monkeypatch.setattr(
litellm.logging_callback_manager, "get_active_additional_logging_utils_from_custom_logger", lambda: []
)
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(user_role=user_role, user_id="user_a")
try:
response = client.get(
"/spend/logs/ui/call-1", params={"start_date": start_date}, headers={"Authorization": "Bearer sk-test"}
)
assert response.status_code == expected_status, response.text
if expected_status == 200:
assert (response.json() or {}).get("messages") == expected_messages, response.text
finally:
app.dependency_overrides.pop(ps.user_api_key_auth, None)
def test_ui_view_request_response_rejects_malformed_start_date_with_400(client, monkeypatch):
monkeypatch.setattr(
"litellm.proxy.proxy_server.prisma_client",
_SpendLogsWithoutCallIdIndex([_payload_row("req-1", "call-1", "user_a", "drawer prompt")]),
)
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN, user_id="user_a"
)
try:
response = client.get(
"/spend/logs/ui/req-1", params={"start_date": "2026-11-09"}, headers={"Authorization": "Bearer sk-test"}
)
assert response.status_code == 400, response.text
finally:
app.dependency_overrides.pop(ps.user_api_key_auth, None)
@pytest.mark.parametrize(
("parent_status", "child_status", "expected"),
[("failure", "success", "failure"), ("success", "failure", "success")],