diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index f14f3d93b79..0704d2c2444 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -18,7 +18,7 @@ import traceback from collections.abc import Awaitable, Callable, Iterable, Iterator, Mapping, Sequence from collections.abc import Set as AbstractSet from dataclasses import dataclass -from datetime import datetime, timedelta, timezone +from datetime import date, datetime, timedelta, timezone from itertools import chain from types import MappingProxyType from typing import ( @@ -228,6 +228,7 @@ from litellm.proxy.management_helpers.utils import ( management_endpoint_wrapper, ) from litellm.proxy.spend_tracking.committed_spend import drop_committed_spend +from litellm.proxy.spend_tracking.spend_log_api_key_scope import spend_log_api_keys_for_teams from litellm.proxy.utils import ( PrismaClient, ProxyLogging, @@ -7463,9 +7464,11 @@ def aggregated_date_range_error(start_date: str | None, end_date: str | None) -> return None -def _team_user_spend_sql(*, team_count: int, restrict_to_user: bool) -> str: +def _team_user_spend_sql(*, team_count: int, restrict_to_user: bool, narrow_to_api_keys: bool) -> str: team_placeholders: Final = ", ".join(f"${i}" for i in range(3, 3 + team_count)) user_clause: Final = f' AND sl."user" = ${3 + team_count}' if restrict_to_user else "" + api_key_placeholder: Final = 3 + team_count + (1 if restrict_to_user else 0) + api_key_clause: Final = f" AND sl.api_key = ANY(${api_key_placeholder}::text[])" if narrow_to_api_keys else "" return f""" SELECT sl.team_id, @@ -7483,7 +7486,7 @@ def _team_user_spend_sql(*, team_count: int, restrict_to_user: bool) -> str: LEFT JOIN "LiteLLM_UserTable" u ON u.user_id = sl."user" WHERE sl."startTime" >= $1::timestamp AND sl."startTime" < $2::timestamp + INTERVAL '1 day' - AND sl.team_id IN ({team_placeholders}){user_clause} + AND sl.team_id IN ({team_placeholders}){user_clause}{api_key_clause} GROUP BY sl.team_id, sl."user", u.user_email, u.user_alias ORDER BY spend DESC, sl.team_id, sl."user" """ @@ -7555,15 +7558,24 @@ async def get_team_spend_by_user( own_user_only: Final = scope.api_key_filter is not None user_param: Final = (user_api_key_dict.user_id or "",) if own_user_only else () + scoped_api_keys: Final = await spend_log_api_keys_for_teams( + prisma_client, scoped_team_ids, date.fromisoformat(start_date), date.fromisoformat(end_date) + ) + api_key_param: Final = (list(scoped_api_keys),) if scoped_api_keys is not None else () async with db_span("team_user_spend", "LiteLLM_SpendLogs"): rows: Final = _TEAM_USER_SPEND_ROWS.validate_python( await query_raw_within_statement_timeout( prisma_client, - _team_user_spend_sql(team_count=len(scoped_team_ids), restrict_to_user=own_user_only), + _team_user_spend_sql( + team_count=len(scoped_team_ids), + restrict_to_user=own_user_only, + narrow_to_api_keys=scoped_api_keys is not None, + ), start_date, end_date, *scoped_team_ids, *user_param, + *api_key_param, what="Spend per user", ) ) diff --git a/litellm/proxy/spend_tracking/spend_log_api_key_scope.py b/litellm/proxy/spend_tracking/spend_log_api_key_scope.py new file mode 100644 index 00000000000..6507011a977 --- /dev/null +++ b/litellm/proxy/spend_tracking/spend_log_api_key_scope.py @@ -0,0 +1,83 @@ +"""``LiteLLM_SpendLogs`` has no index on ``"user"`` or ``team_id``, so a read filtered by owner scans the +whole date window. The owner's ``LiteLLM_Daily*Spend`` rows (kept after key deletion) plus their live keys +name every api_key that produced those logs, and ``(api_key, "startTime")`` is indexed.""" + +from collections.abc import Mapping, Sequence +from datetime import date, timedelta +from typing import Final + +from pydantic import ConfigDict, TypeAdapter + +from litellm.proxy.utils import PrismaClient, query_raw_within_statement_timeout +from litellm.types.llms.base import LiteLLMBaseModel + +SPEND_LOG_API_KEY_SCOPE_MAX_KEYS: Final = 5000 +_WHAT: Final = "Spend log key lookup" + +_USER_API_KEYS_SQL: Final = """ +SELECT DISTINCT api_key FROM ( + SELECT api_key FROM "LiteLLM_DailyUserSpend" + WHERE user_id = $1 AND date >= $2 AND date <= $3 + UNION ALL + SELECT token AS api_key FROM "LiteLLM_VerificationToken" WHERE user_id = $1 +) owned +LIMIT $4 +""" + +_TEAM_API_KEYS_SQL: Final = """ +SELECT DISTINCT api_key FROM ( + SELECT api_key FROM "LiteLLM_DailyTeamSpend" + WHERE team_id = ANY($1::text[]) AND date >= $2 AND date <= $3 + UNION ALL + SELECT token AS api_key FROM "LiteLLM_VerificationToken" WHERE team_id = ANY($1::text[]) +) owned +LIMIT $4 +""" + + +class _ApiKeyRow(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + api_key: str + + +_API_KEY_ROWS: Final = TypeAdapter(tuple[_ApiKeyRow, ...]) + + +def _narrowing_keys(rows: Sequence[Mapping[str, object]]) -> tuple[str, ...] | None: + """None when the list cannot narrow the read: nothing known about the owner, or too many keys to help.""" + keys: Final = tuple(row.api_key for row in _API_KEY_ROWS.validate_python(rows)) + return None if not keys or len(keys) > SPEND_LOG_API_KEY_SCOPE_MAX_KEYS else keys + + +def _daily_dates_around(start: date, end: date) -> tuple[str, str]: + return (start - timedelta(days=1)).isoformat(), (end + timedelta(days=1)).isoformat() + + +async def spend_log_api_keys_for_user( + prisma_client: PrismaClient, user_id: str, start: date, end: date +) -> tuple[str, ...] | None: + """The api_keys that can carry this user's spend logs in the window, or None when there are too many to help.""" + first, last = _daily_dates_around(start, end) + return _narrowing_keys( + await query_raw_within_statement_timeout( + prisma_client, _USER_API_KEYS_SQL, user_id, first, last, SPEND_LOG_API_KEY_SCOPE_MAX_KEYS + 1, what=_WHAT + ) + ) + + +async def spend_log_api_keys_for_teams( + prisma_client: PrismaClient, team_ids: Sequence[str], start: date, end: date +) -> tuple[str, ...] | None: + first, last = _daily_dates_around(start, end) + return _narrowing_keys( + await query_raw_within_statement_timeout( + prisma_client, + _TEAM_API_KEYS_SQL, + list(team_ids), + first, + last, + SPEND_LOG_API_KEY_SCOPE_MAX_KEYS + 1, + what=_WHAT, + ) + ) diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 7d9a0dcc971..6d87b39afb3 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -70,6 +70,10 @@ from litellm.proxy.spend_tracking.spend_capture_rate import ( # NOTE: Avoid module-level import from common_utils: proxy_server imports this # module while common_utils may pull proxy_server during init, which can leave # those names undefined. Import the helpers locally where they are used. +from litellm.proxy.spend_tracking.spend_log_api_key_scope import ( + spend_log_api_keys_for_teams, + spend_log_api_keys_for_user, +) from litellm.proxy.spend_tracking.spend_tracking_utils import ( get_spend_by_team, get_spend_by_team_and_customer, @@ -257,6 +261,22 @@ async def _query_raw(prisma_client: PrismaClient, sql_query: str, *args: object) return await prisma_client.db.query_raw(sql_query, *args) +async def _scope_api_keys( + prisma_client: PrismaClient, *, where_conditions: Mapping[str, object], start: datetime | None, end: datetime | None +) -> tuple[str, ...] | None: + """The api_keys behind a ``"user"`` / ``team_id`` filter, so the unindexed column is only checked on rows the + ``(api_key, "startTime")`` index already picked.""" + if start is None or end is None or isinstance(where_conditions.get("api_key"), str): + return None + user_id: Final = where_conditions.get("user") + if isinstance(user_id, str): + return await spend_log_api_keys_for_user(prisma_client, user_id, start.date(), end.date()) + team_id: Final = where_conditions.get("team_id") + if isinstance(team_id, str): + return await spend_log_api_keys_for_teams(prisma_client, (team_id,), start.date(), end.date()) + return None + + async def _query_raw_or_none(prisma_client: PrismaClient, sql_query: str, *args: object) -> Sequence[_RowT] | None: """``_query_raw`` for the call sites that guard the result against ``None``.""" return await _query_raw(prisma_client, sql_query, *args) @@ -2950,6 +2970,18 @@ async def ui_view_spend_logs( sql_params.append(f"%{like_escaped_session_id}%") p += 1 + scope_api_keys: Final = ( + await _scope_api_keys( + prisma_client, where_conditions=where_conditions, start=start_date_obj, end=end_date_obj + ) + if not search_owns_window + else None + ) + if scope_api_keys is not None: + sql_conditions.append(f"api_key = ANY(${p}::text[])") + sql_params.append(list(scope_api_keys)) + p += 1 # rebind-ok: advances the file's shared $N placeholder counter + # Status filter if status_filter is not None and not (group_by_session is True and not is_search_lookup): if status_filter == "success": diff --git a/tests/unit/proxy/management_endpoints/test_team_endpoints.py b/tests/unit/proxy/management_endpoints/test_team_endpoints.py index a79ebd307fe..ac46ccc1cfc 100644 --- a/tests/unit/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_team_endpoints.py @@ -17341,6 +17341,19 @@ def _team_spend_by_user_caller(user_id: str, teams: list[str]) -> LiteLLM_UserTa ) +def _team_spend_reads(api_keys: list[str], spend_rows: list[dict] | Exception) -> AsyncMock: + """Raw reads inside the statement-timeout transaction: the key lookup, then the spend log aggregate.""" + + async def read(sql: str, *params: object) -> list[dict]: + if 'FROM "LiteLLM_DailyTeamSpend"' in sql: + return [{"api_key": api_key} for api_key in api_keys] + if isinstance(spend_rows, Exception): + raise spend_rows + return spend_rows + + return AsyncMock(side_effect=read) + + def _team_spend_by_user_db_row(team_id: str, user_id: str, spend: float, requests: int) -> dict: return { "team_id": team_id, @@ -17367,12 +17380,13 @@ async def test_get_team_spend_by_user_admin_groups_spend_logs_by_team_and_user(m mock_db_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[alpha, beta]) transaction = statement_timeout_reads( mock_db_client, - AsyncMock( - return_value=[ + _team_spend_reads( + ["sk-alpha", "sk-beta"], + [ _team_spend_by_user_db_row("team-alpha", "alice", 0.5, 3), _team_spend_by_user_db_row("team-alpha", "bob", 0.25, 2), _team_spend_by_user_db_row("team-beta", "alice", 0.1, 1), - ] + ], ), ) @@ -17386,11 +17400,15 @@ async def test_get_team_spend_by_user_admin_groups_spend_logs_by_team_and_user(m assert any("statement_timeout = 30000" in statement for statement in transaction.statements) mock_db_client.db.tx.assert_not_called() sql, *params = transaction.query_raw.call_args.args - assert params == ["2026-09-01", "2026-09-04", "team-alpha", "team-beta"] + assert params == ["2026-09-01", "2026-09-04", "team-alpha", "team-beta", ["sk-alpha", "sk-beta"]] assert 'FROM "LiteLLM_SpendLogs" sl' in sql assert 'sl."startTime" >= $1::timestamp' in sql assert "sl.\"startTime\" < $2::timestamp + INTERVAL '1 day'" in sql assert "sl.team_id IN ($3, $4)" in sql + assert "sl.api_key = ANY($5::text[])" in sql + lookup_sql, *lookup_params = transaction.query_raw.call_args_list[0].args + assert 'FROM "LiteLLM_DailyTeamSpend"' in lookup_sql and 'FROM "LiteLLM_VerificationToken"' in lookup_sql + assert lookup_params == [["team-alpha", "team-beta"], "2026-08-31", "2026-09-05", 5001] assert 'GROUP BY sl.team_id, sl."user"' in sql assert 'sl."user" = $' not in sql @@ -17412,7 +17430,7 @@ async def test_get_team_spend_by_user_team_admin_sees_every_member(mock_db_clien caller = UserAPIKeyAuth(user_id="alice", user_role=LitellmUserRoles.INTERNAL_USER) alpha = _team_spend_by_user_team("team-alpha", "Team Alpha", Member(user_id="alice", role="admin"), []) mock_db_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[alpha]) - transaction = statement_timeout_reads(mock_db_client, AsyncMock(return_value=[])) + transaction = statement_timeout_reads(mock_db_client, _team_spend_reads(["sk-alpha"], [])) mock_db_client.writer_db.litellm_usertable.find_unique = AsyncMock( return_value=_team_spend_by_user_caller("alice", ["team-alpha"]) ) @@ -17422,8 +17440,9 @@ async def test_get_team_spend_by_user_team_admin_sees_every_member(mock_db_clien ) sql, *params = transaction.query_raw.call_args.args - assert params == ["2026-09-01", "2026-09-04", "team-alpha"] + assert params == ["2026-09-01", "2026-09-04", "team-alpha", ["sk-alpha"]] assert 'sl."user" = $' not in sql + assert "sl.api_key = ANY($4::text[])" in sql @pytest.mark.asyncio @@ -17435,7 +17454,7 @@ async def test_get_team_spend_by_user_plain_member_only_sees_own_row(mock_db_cli mock_db_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[alpha]) mock_db_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) transaction = statement_timeout_reads( - mock_db_client, AsyncMock(return_value=[_team_spend_by_user_db_row("team-alpha", "bob", 0.25, 2)]) + mock_db_client, _team_spend_reads(["sk-bob"], [_team_spend_by_user_db_row("team-alpha", "bob", 0.25, 2)]) ) mock_db_client.writer_db.litellm_usertable.find_unique = AsyncMock( return_value=_team_spend_by_user_caller("bob", ["team-alpha"]) @@ -17446,9 +17465,10 @@ async def test_get_team_spend_by_user_plain_member_only_sees_own_row(mock_db_cli ) sql, *params = transaction.query_raw.call_args.args - assert params == ["2026-09-01", "2026-09-04", "team-alpha", "bob"] + assert params == ["2026-09-01", "2026-09-04", "team-alpha", "bob", ["sk-bob"]] assert "sl.team_id IN ($3)" in sql assert 'AND sl."user" = $4' in sql + assert "sl.api_key = ANY($5::text[])" in sql assert [(r.user_id, r.spend) for r in response.results] == [("bob", 0.25)] @@ -17506,7 +17526,7 @@ async def test_get_team_spend_by_user_statement_timeout_returns_503(mock_db_clie admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) alpha = _team_spend_by_user_team("team-alpha", "Team Alpha", Member(user_id="alice", role="admin"), []) mock_db_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[alpha]) - statement_timeout_reads(mock_db_client, AsyncMock(side_effect=statement_timeout_error())) + statement_timeout_reads(mock_db_client, _team_spend_reads(["sk-alpha"], statement_timeout_error())) with pytest.raises(HTTPException) as exc_info: await get_team_spend_by_user( @@ -17520,6 +17540,30 @@ async def test_get_team_spend_by_user_statement_timeout_returns_503(mock_db_clie assert "longer than 30s" in str(exc_info.value.detail) +@pytest.mark.asyncio +@pytest.mark.parametrize("known_key_count", [0, 5001]) +async def test_get_team_spend_by_user_without_a_usable_key_list_keeps_the_plain_team_filter( + mock_db_client, known_key_count: int +): + from litellm.proxy.management_endpoints.team_endpoints import get_team_spend_by_user + + admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + alpha = _team_spend_by_user_team("team-alpha", "Team Alpha", Member(user_id="alice", role="admin"), []) + mock_db_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[alpha]) + transaction = statement_timeout_reads( + mock_db_client, _team_spend_reads([f"sk-{i}" for i in range(known_key_count)], []) + ) + + await get_team_spend_by_user( + user_api_key_dict=admin, team_ids="team-alpha", start_date="2026-09-01", end_date="2026-09-04" + ) + + sql, *params = transaction.query_raw.call_args.args + assert params == ["2026-09-01", "2026-09-04", "team-alpha"] + assert "sl.team_id IN ($3)" in sql + assert "sl.api_key = ANY(" not in sql + + class _TeamRowWithOrganization(LiteLLM_TeamTable): litellm_organization_table: LiteLLM_OrganizationTable | None = None 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 4a418d87f22..29e2f97d11c 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py @@ -6,7 +6,7 @@ import hashlib import json import re import sqlite3 -from collections.abc import AsyncIterator, Awaitable, Mapping, Sequence +from collections.abc import AsyncIterator, Awaitable, Callable, Mapping, Sequence from datetime import timezone from functools import partial from types import SimpleNamespace @@ -141,6 +141,7 @@ def _reconstruct_ui_where_from_sql(sql_query, params): sess = re.fullmatch(r"session_id LIKE \$(\d+)", cond) status = re.fullmatch(r"status = \$(\d+)", cond) api_key_not_in = re.fullmatch(r"api_key NOT IN \(\$(\d+), \$(\d+)\)", cond) + api_key_any = re.fullmatch(r"api_key = ANY\(\$(\d+)::text\[\]\)", cond) req_or_call = re.fullmatch(r"\(request_id = \$(\d+) OR litellm_call_id = \$\1\)", cond) if req_or_call: where["request_id_or_call_id"] = params[int(req_or_call.group(1)) - 1] @@ -169,6 +170,8 @@ def _reconstruct_ui_where_from_sql(sql_query, params): where["session_id"] = {"contains": str(params[int(sess.group(1)) - 1]).strip("%")} elif status: where["status"] = {"equals": params[int(status.group(1)) - 1]} + elif api_key_any: + where["api_key_in"] = list(params[int(api_key_any.group(1)) - 1]) elif api_key_not_in: where["api_key_not_in"] = [ params[int(api_key_not_in.group(1)) - 1], @@ -218,7 +221,18 @@ def _reconstruct_ui_where_from_sql(sql_query, params): return where -def make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_fn, team_lookup_fn=None, query_observer=None): +def _emulate_owner_api_key_lookup(mock_spend_logs: list[dict], sql_query: str, params: tuple) -> list[dict]: + """The daily spend tables and key table name every api_key an owner's requests ran under, so the + lookup answers with the distinct api_keys of the owner's logs.""" + owner_col = "user" if 'FROM "LiteLLM_DailyUserSpend"' in sql_query else "team_id" + owners = {params[0]} if isinstance(params[0], str) else set(params[0]) + api_keys = {log["api_key"] for log in mock_spend_logs if log.get(owner_col) in owners} + return [{"api_key": api_key} for api_key in sorted(api_keys)] + + +def make_ui_spend_logs_mock_prisma( + mock_spend_logs, filter_fn, team_lookup_fn=None, query_observer=None, owner_api_keys_fn=None +): """ Create a MockPrismaClient for /spend/logs/ui endpoint tests. @@ -228,6 +242,8 @@ def make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_fn, team_lookup_fn=No where_conditions, returns the filtered list of logs. team_lookup_fn: Optional async callable for team RBAC (find_unique). If provided, adds litellm_teamtable to db. + owner_api_keys_fn: Optional Callable[[tuple], list[dict]] answering the + user/team -> api_key lookup instead of deriving it from the logs. """ class MockDB: @@ -251,6 +267,10 @@ def make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_fn, team_lookup_fn=No return [] if 'SELECT DISTINCT "user", team_id' in sql_query: return _emulate_spend_log_owner_lookup(mock_spend_logs, sql_query, params) + if 'FROM "LiteLLM_DailyUserSpend"' in sql_query or 'FROM "LiteLLM_DailyTeamSpend"' in sql_query: + if owner_api_keys_fn is not None: + return owner_api_keys_fn(params) + return _emulate_owner_api_key_lookup(mock_spend_logs, sql_query, params) filtered = filter_fn(_reconstruct_ui_where_from_sql(sql_query, params)) total = len(filtered) if "COUNT(*)" in sql_query: @@ -845,6 +865,196 @@ async def test_ui_view_spend_logs_with_user_id(client, monkeypatch): assert data["data"][0]["user"] == "test_user_1" +def _owned_log(request_id: str, api_key: str, user: str, team_id: str) -> dict: + return { + "id": request_id, + "request_id": request_id, + "api_key": api_key, + "user": user, + "team_id": team_id, + "spend": 0.01, + "startTime": "2026-09-02T10:00:00+00:00", + "model": "gpt-4o", + } + + +def _owner_filtered_logs(logs: list[dict]) -> Callable[[dict], list[dict]]: + """What Postgres returns for the page query: the owner filter and, when present, the api_key list both apply.""" + + def filter_fn(where: dict) -> list[dict]: + return [ + log + for log in logs + if log["user"] == where.get("user", log["user"]) + and log["team_id"] == where.get("team_id", log["team_id"]) + and log["api_key"] in where.get("api_key_in", [log["api_key"]]) + ] + + return filter_fn + + +def _owner_api_keys(daily_rows: dict[str, set[str]], live_keys: dict[str, set[str]]) -> Callable[[tuple], list[dict]]: + """The key lookup's two sources, keyed by owner: api_keys on daily spend rows and api_keys of live keys.""" + + def lookup(params: tuple) -> list[dict]: + owners = {params[0]} if isinstance(params[0], str) else set(params[0]) + api_keys = set().union(*(daily_rows.get(o, set()) | live_keys.get(o, set()) for o in owners)) + return [{"api_key": api_key} for api_key in sorted(api_keys)] + + return lookup + + +def _spend_logs_page_query(queries: list[tuple[str, tuple]]) -> tuple[str, tuple, list[str] | None]: + sql, params = next((sql, params) for sql, params in queries if '"LiteLLM_SpendLogs"' in sql and "ORDER BY" in sql) + key_param = re.search(r"api_key = ANY\(\$(\d+)::text\[\]\)", sql) + return sql, params, None if key_param is None else params[int(key_param.group(1)) - 1] + + +@pytest.mark.asyncio +async def test_ui_view_spend_logs_user_filter_reads_spend_logs_through_the_users_api_keys(client, monkeypatch): + from litellm.proxy.spend_tracking.spend_log_api_key_scope import SPEND_LOG_API_KEY_SCOPE_MAX_KEYS + + logs = [ + _owned_log("req1", "sk-live", "alice", "team1"), + _owned_log("req2", "sk-deleted", "alice", "team1"), + _owned_log("req3", "sk-bob", "bob", "team1"), + ] + queries: list[tuple[str, tuple]] = [] + monkeypatch.setattr( + "litellm.proxy.proxy_server.prisma_client", + make_ui_spend_logs_mock_prisma( + logs, + _owner_filtered_logs(logs), + query_observer=lambda sql, params: queries.append((sql, params)), + owner_api_keys_fn=_owner_api_keys( + daily_rows={"alice": {"sk-deleted"}, "bob": {"sk-bob"}}, live_keys={"alice": {"sk-live"}} + ), + ), + ) + + response = client.get( + "/spend/logs/ui", + params={"user_id": "alice", "start_date": "2026-09-01 00:00:00", "end_date": "2026-09-03 23:59:59"}, + headers={"Authorization": "Bearer sk-test"}, + ) + + assert response.status_code == 200, response.text + assert {row["request_id"] for row in response.json()["data"]} == {"req1", "req2"} + lookup_sql, lookup_params = next((sql, params) for sql, params in queries if '"LiteLLM_DailyUserSpend"' in sql) + assert '"LiteLLM_VerificationToken"' in lookup_sql + assert lookup_params == ("alice", "2026-08-31", "2026-09-04", SPEND_LOG_API_KEY_SCOPE_MAX_KEYS + 1) + page_sql, _, scoped_api_keys = _spend_logs_page_query(queries) + assert scoped_api_keys == ["sk-deleted", "sk-live"] + assert '"user" = $' in page_sql + + +@pytest.mark.asyncio +async def test_ui_view_spend_logs_team_filter_reads_spend_logs_through_the_teams_api_keys(client, monkeypatch): + logs = [ + _owned_log("req1", "sk-alpha-1", "alice", "team-alpha"), + _owned_log("req2", "sk-alpha-2", "bob", "team-alpha"), + _owned_log("req3", "sk-beta", "carol", "team-beta"), + ] + queries: list[tuple[str, tuple]] = [] + monkeypatch.setattr( + "litellm.proxy.proxy_server.prisma_client", + make_ui_spend_logs_mock_prisma( + logs, + _owner_filtered_logs(logs), + query_observer=lambda sql, params: queries.append((sql, params)), + owner_api_keys_fn=_owner_api_keys( + daily_rows={"team-alpha": {"sk-alpha-1"}, "team-beta": {"sk-beta"}}, + live_keys={"team-alpha": {"sk-alpha-2"}}, + ), + ), + ) + monkeypatch.setattr( + "litellm.proxy.spend_tracking.spend_management_endpoints._is_admin_view_safe", + lambda user_api_key_dict: True, + ) + monkeypatch.setitem( + app.dependency_overrides, + ps.user_api_key_auth, + lambda: UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user"), + ) + + response = client.get( + "/spend/logs/ui", + params={"team_id": "team-alpha", "start_date": "2026-09-01 00:00:00", "end_date": "2026-09-03 23:59:59"}, + headers={"Authorization": "Bearer sk-test"}, + ) + + assert response.status_code == 200, response.text + assert {row["request_id"] for row in response.json()["data"]} == {"req1", "req2"} + _, lookup_params = next((sql, params) for sql, params in queries if '"LiteLLM_DailyTeamSpend"' in sql) + assert lookup_params[:3] == (["team-alpha"], "2026-08-31", "2026-09-04") + page_sql, _, scoped_api_keys = _spend_logs_page_query(queries) + assert scoped_api_keys == ["sk-alpha-1", "sk-alpha-2"] + assert "team_id = $" in page_sql + + +@pytest.mark.asyncio +@pytest.mark.parametrize("known_keys", [set(), {f"sk-{i}" for i in range(5001)}]) +async def test_ui_view_spend_logs_without_a_usable_key_list_keeps_the_plain_user_filter( + client, monkeypatch, known_keys: set[str] +): + logs = [_owned_log("req1", "sk-live", "alice", "team1")] + queries: list[tuple[str, tuple]] = [] + monkeypatch.setattr( + "litellm.proxy.proxy_server.prisma_client", + make_ui_spend_logs_mock_prisma( + logs, + _owner_filtered_logs(logs), + query_observer=lambda sql, params: queries.append((sql, params)), + owner_api_keys_fn=_owner_api_keys(daily_rows={"alice": known_keys}, live_keys={}), + ), + ) + + response = client.get( + "/spend/logs/ui", + params={"user_id": "alice", "start_date": "2026-09-01 00:00:00", "end_date": "2026-09-03 23:59:59"}, + headers={"Authorization": "Bearer sk-test"}, + ) + + assert response.status_code == 200, response.text + assert [row["request_id"] for row in response.json()["data"]] == ["req1"] + page_sql, _, scoped_api_keys = _spend_logs_page_query(queries) + assert scoped_api_keys is None + assert '"user" = $' in page_sql + + +@pytest.mark.asyncio +async def test_ui_view_spend_logs_search_with_user_filter_is_not_narrowed_to_the_windows_keys(client, monkeypatch): + logs = [_owned_log("req-old", "sk-retired", "alice", "team1")] + queries: list[tuple[str, tuple]] = [] + monkeypatch.setattr( + "litellm.proxy.proxy_server.prisma_client", + make_ui_spend_logs_mock_prisma( + logs, + _owner_filtered_logs(logs), + query_observer=lambda sql, params: queries.append((sql, params)), + owner_api_keys_fn=_owner_api_keys(daily_rows={}, live_keys={"alice": {"sk-live"}}), + ), + ) + + response = client.get( + "/spend/logs/ui", + params={ + "user_id": "alice", + "search": "req-old", + "start_date": "2026-09-01 00:00:00", + "end_date": "2026-09-03 23:59:59", + }, + headers={"Authorization": "Bearer sk-test"}, + ) + + assert response.status_code == 200, response.text + assert [row["request_id"] for row in response.json()["data"]] == ["req-old"] + page_sql, _, scoped_api_keys = _spend_logs_page_query(queries) + assert scoped_api_keys is None + assert "request_id = $" in page_sql and '"user" = $' in page_sql + + @pytest.mark.asyncio @pytest.mark.parametrize( "session_id_query,expected_request_ids",