From d70855ca826d5aa1eeab7adfef56d8524ca40bb7 Mon Sep 17 00:00:00 2001 From: tavren27 <348969968+tavren27@users.noreply.github.com> Date: Tue, 10 Nov 2026 09:03:21 -0800 Subject: [PATCH] fix(proxy): open the Logs drawer by request_id without scanning for litellm_call_id (#50789) * 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> --- .../spend_management_endpoints.py | 171 +++++++++++---- litellm/proxy/utils.py | 9 +- .../spend/test_session_grouped_logs_pages.py | 4 +- .../spend/test_spend_log_read_scope.py | 126 ++++++++++- .../test_spend_management_endpoints.py | 199 ++++++++++++++++-- 5 files changed, 445 insertions(+), 64 deletions(-) diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index f3dc672a072..7d9a0dcc971 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -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: diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 97c327cbe4e..f78e68220b1 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -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 diff --git a/tests/integration/spend/test_session_grouped_logs_pages.py b/tests/integration/spend/test_session_grouped_logs_pages.py index 7989cb6a7bc..9d457c63157 100644 --- a/tests/integration/spend/test_session_grouped_logs_pages.py +++ b/tests/integration/spend/test_session_grouped_logs_pages.py @@ -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 diff --git a/tests/integration/spend/test_spend_log_read_scope.py b/tests/integration/spend/test_spend_log_read_scope.py index 8a208d61aed..b6226d0664e 100644 --- a/tests/integration/spend/test_spend_log_read_scope.py +++ b/tests/integration/spend/test_spend_log_read_scope.py @@ -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,)) diff --git a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py index 27a5734eb45..4a418d87f22 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py @@ -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")],