From bc4c89bfc3c13a499d71363fd2e26fcbeb814d31 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Wed, 4 Nov 2026 23:51:07 +0000 Subject: [PATCH] fix(proxy): cancel /team/spend/by_user in Postgres after 30s (#49922) Backport of #49922 to lts/1.108.x. Cherry-picked from 391c8013c9aa (main). Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/db/exception_handler.py | 13 +++ .../management_endpoints/team_endpoints.py | 39 ++++++-- tests/unit/proxy/db/test_exception_handler.py | 15 +++ .../test_team_endpoints.py | 94 ++++++++++++++++--- .../EntityUsage/TeamUserSpendCard.test.tsx | 39 ++++++++ .../EntityUsage/TeamUserSpendCard.tsx | 26 +++-- 6 files changed, 193 insertions(+), 33 deletions(-) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/TeamUserSpendCard.test.tsx diff --git a/litellm/proxy/db/exception_handler.py b/litellm/proxy/db/exception_handler.py index 5dfcaa4b2f8..669df7454e3 100644 --- a/litellm/proxy/db/exception_handler.py +++ b/litellm/proxy/db/exception_handler.py @@ -34,6 +34,8 @@ _CONNECTION_CAPACITY_PHRASES: Final = ( _PRISMA_POOL_TIMEOUT_CODE: Final = "P2024" _LOCK_TIMEOUT_SQLSTATE: Final = "55P03" _LOCK_TIMEOUT_PHRASE: Final = "canceling statement due to lock timeout" +_STATEMENT_TIMEOUT_SQLSTATE: Final = "57014" +_STATEMENT_TIMEOUT_PHRASE: Final = "canceling statement due to statement timeout" def _exception_chain(e: BaseException) -> Iterator[BaseException]: @@ -284,6 +286,17 @@ class PrismaDBExceptionHandler: return True return _LOCK_TIMEOUT_PHRASE in str(e).lower() + @staticmethod + def is_statement_timeout_error(e: Exception) -> bool: + """True iff Postgres cancelled a statement for exceeding ``statement_timeout`` (SQLSTATE 57014).""" + import prisma + + if not isinstance(e, _exception_types(prisma.errors.PrismaError)): + return False + if PrismaDBExceptionHandler.postgres_sqlstate(e) == _STATEMENT_TIMEOUT_SQLSTATE: + return True + return _STATEMENT_TIMEOUT_PHRASE in str(e).lower() + @staticmethod def postgres_sqlstate(e: Exception) -> str | None: """The SQLSTATE Postgres attached to a failed statement, as prisma surfaces it, or None.""" diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 0f40662ad83..99d168315c6 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -126,6 +126,7 @@ from litellm.proxy.common_utils.callback_utils import encrypt_callback_vars from litellm.proxy.common_utils.json_merge_patch import apply_json_merge_patch from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.db.db_span import db_span +from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks from litellm.proxy.hooks.model_max_budget_limiter import ( build_model_max_budget_usage, @@ -6969,6 +6970,19 @@ class _TeamUserSpendDbRow(TypedDict): failed_requests: ReadOnly[int] +_TEAM_USER_SPEND_STATEMENT_TIMEOUT_SECONDS: Final = 30 + + +async def _team_user_spend_rows( + prisma_client: PrismaClient, sql: str, params: tuple[str, ...] +) -> Sequence[_TeamUserSpendDbRow]: + async with prisma_client.read_db.tx( + timeout=timedelta(seconds=_TEAM_USER_SPEND_STATEMENT_TIMEOUT_SECONDS + 5) + ) as tx: + await tx.execute_raw(f"SET LOCAL statement_timeout = {_TEAM_USER_SPEND_STATEMENT_TIMEOUT_SECONDS * 1000}") + return await tx.query_raw(sql, *params) + + @router.get( "/team/spend/by_user", response_model=TeamUserSpendResponse, @@ -7018,14 +7032,23 @@ 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 () - async with db_span("team_user_spend", "LiteLLM_SpendLogs"): - rows: Final[Sequence[_TeamUserSpendDbRow]] = await prisma_client.db.query_raw( - _team_user_spend_sql(team_count=len(scoped_team_ids), restrict_to_user=own_user_only), - start_date, - end_date, - *scoped_team_ids, - *user_param, - ) + try: + async with db_span("team_user_spend", "LiteLLM_SpendLogs"): + rows: Final = await _team_user_spend_rows( + prisma_client, + _team_user_spend_sql(team_count=len(scoped_team_ids), restrict_to_user=own_user_only), + (start_date, end_date, *scoped_team_ids, *user_param), + ) + except Exception as e: + if PrismaDBExceptionHandler.is_statement_timeout_error(e): + raise _daily_activity_error( + status_code=503, + message=( + f"Spend per user took longer than {_TEAM_USER_SPEND_STATEMENT_TIMEOUT_SECONDS}s. " + "Pick a shorter date range or fewer teams" + ), + ) from e + raise results: Final = tuple( TeamUserSpendRow( team_id=row["team_id"], diff --git a/tests/unit/proxy/db/test_exception_handler.py b/tests/unit/proxy/db/test_exception_handler.py index 75346ca08a6..a4c16c0949b 100644 --- a/tests/unit/proxy/db/test_exception_handler.py +++ b/tests/unit/proxy/db/test_exception_handler.py @@ -951,3 +951,18 @@ def test_is_lock_timeout_error_recognises_a_statement_cancelled_under_lock_timeo ) def test_is_lock_timeout_error_excludes_other_failures(error: Exception) -> None: assert PrismaDBExceptionHandler.is_lock_timeout_error(error) is False + + +@pytest.mark.parametrize( + ("error", "expected"), + [ + (_lock_timeout_error(message="canceling statement due to statement timeout", code="57014"), True), + (_lock_timeout_error(message="canceling statement due to statement timeout", code=None), True), + (_lock_timeout_error(), False), + (_capacity_error("sorry, too many clients already"), False), + (RuntimeError("canceling statement due to statement timeout"), False), + ], + ids=["meta_57014", "message_only", "lock_timeout_55P03", "capacity_53300", "not_prisma"], +) +def test_is_statement_timeout_error(error: Exception, expected: bool) -> None: + assert PrismaDBExceptionHandler.is_statement_timeout_error(error) is expected diff --git a/tests/unit/proxy/management_endpoints/test_team_endpoints.py b/tests/unit/proxy/management_endpoints/test_team_endpoints.py index a379b424aeb..5bbc7290380 100644 --- a/tests/unit/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_team_endpoints.py @@ -13,6 +13,7 @@ import litellm import pytest from fastapi import HTTPException, Request from fastapi.testclient import TestClient +from prisma.errors import DataError from pydantic import ValidationError from litellm._uuid import uuid @@ -16388,6 +16389,28 @@ def _team_spend_by_user_db_row(team_id: str, user_id: str, spend: float, request } +class _RecordingTransaction: + def __init__(self, query_raw: AsyncMock) -> None: + self.statements: tuple[str, ...] = () + self.query_raw: AsyncMock = query_raw + + async def execute_raw(self, sql: str) -> int: + self.statements = (*self.statements, sql) + return 0 + + async def __aenter__(self) -> "_RecordingTransaction": + return self + + async def __aexit__(self, *exc_info: object) -> None: + pass + + +def _team_spend_by_user_query(mock_db_client: MagicMock, query_raw: AsyncMock) -> _RecordingTransaction: + transaction: Final = _RecordingTransaction(query_raw) + mock_db_client.read_db.tx = MagicMock(return_value=transaction) + return transaction + + @pytest.mark.asyncio async def test_get_team_spend_by_user_admin_groups_spend_logs_by_team_and_user(mock_db_client): from litellm.proxy.management_endpoints.team_endpoints import get_team_spend_by_user @@ -16396,12 +16419,15 @@ async def test_get_team_spend_by_user_admin_groups_spend_logs_by_team_and_user(m alpha = _team_spend_by_user_team("team-alpha", "Team Alpha", Member(user_id="alice", role="admin"), []) beta = _team_spend_by_user_team("team-beta", "Team Beta", Member(user_id="alice", role="user"), []) mock_db_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[alpha, beta]) - mock_db_client.db.query_raw = AsyncMock( - return_value=[ - _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), - ] + transaction = _team_spend_by_user_query( + mock_db_client, + AsyncMock( + return_value=[ + _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), + ] + ), ) response = await get_team_spend_by_user( @@ -16411,7 +16437,9 @@ async def test_get_team_spend_by_user_admin_groups_spend_logs_by_team_and_user(m end_date="2026-09-04", ) - sql, *params = mock_db_client.db.query_raw.call_args.args + 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 'FROM "LiteLLM_SpendLogs" sl' in sql assert 'sl."startTime" >= $1::timestamp' in sql @@ -16438,7 +16466,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]) - mock_db_client.db.query_raw = AsyncMock(return_value=[]) + transaction = _team_spend_by_user_query(mock_db_client, AsyncMock(return_value=[])) mock_db_client.writer_db.litellm_usertable.find_unique = AsyncMock( return_value=_team_spend_by_user_caller("alice", ["team-alpha"]) ) @@ -16447,7 +16475,7 @@ async def test_get_team_spend_by_user_team_admin_sees_every_member(mock_db_clien user_api_key_dict=caller, team_ids="team-alpha", start_date="2026-09-01", end_date="2026-09-04" ) - sql, *params = mock_db_client.db.query_raw.call_args.args + sql, *params = transaction.query_raw.call_args.args assert params == ["2026-09-01", "2026-09-04", "team-alpha"] assert 'sl."user" = $' not in sql @@ -16460,7 +16488,9 @@ async def test_get_team_spend_by_user_plain_member_only_sees_own_row(mock_db_cli alpha = _team_spend_by_user_team("team-alpha", "Team Alpha", Member(user_id="bob", role="user"), ["/key/info"]) mock_db_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[alpha]) mock_db_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) - mock_db_client.db.query_raw = AsyncMock(return_value=[_team_spend_by_user_db_row("team-alpha", "bob", 0.25, 2)]) + transaction = _team_spend_by_user_query( + mock_db_client, AsyncMock(return_value=[_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"]) ) @@ -16469,7 +16499,7 @@ async def test_get_team_spend_by_user_plain_member_only_sees_own_row(mock_db_cli user_api_key_dict=caller, team_ids="team-alpha", start_date="2026-09-01", end_date="2026-09-04" ) - sql, *params = mock_db_client.db.query_raw.call_args.args + sql, *params = transaction.query_raw.call_args.args assert params == ["2026-09-01", "2026-09-04", "team-alpha", "bob"] assert "sl.team_id IN ($3)" in sql assert 'AND sl."user" = $4' in sql @@ -16481,7 +16511,7 @@ async def test_get_team_spend_by_user_member_of_other_team_gets_404(mock_db_clie from litellm.proxy.management_endpoints.team_endpoints import get_team_spend_by_user caller = UserAPIKeyAuth(user_id="bob", user_role=LitellmUserRoles.INTERNAL_USER) - mock_db_client.db.query_raw = AsyncMock(return_value=[]) + transaction = _team_spend_by_user_query(mock_db_client, AsyncMock(return_value=[])) mock_db_client.writer_db.litellm_usertable.find_unique = AsyncMock( return_value=_team_spend_by_user_caller("bob", ["team-alpha"]) ) @@ -16492,7 +16522,7 @@ async def test_get_team_spend_by_user_member_of_other_team_gets_404(mock_db_clie ) assert exc_info.value.status_code == 404 - mock_db_client.db.query_raw.assert_not_called() + transaction.query_raw.assert_not_called() @pytest.mark.asyncio @@ -16510,7 +16540,7 @@ async def test_get_team_spend_by_user_member_of_other_team_gets_404(mock_db_clie async def test_get_team_spend_by_user_rejects_bad_input(mock_db_client, team_ids, start_date, end_date, expected_error): from litellm.proxy.management_endpoints.team_endpoints import get_team_spend_by_user - mock_db_client.db.query_raw = AsyncMock(return_value=[]) + transaction = _team_spend_by_user_query(mock_db_client, AsyncMock(return_value=[])) admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) with pytest.raises(HTTPException) as exc_info: @@ -16520,7 +16550,41 @@ async def test_get_team_spend_by_user_rejects_bad_input(mock_db_client, team_ids assert exc_info.value.status_code == 400 assert expected_error in str(exc_info.value.detail) - mock_db_client.db.query_raw.assert_not_called() + transaction.query_raw.assert_not_called() + + +def _statement_timeout_error() -> DataError: + message = "canceling statement due to statement timeout" + return DataError( + data={ + "user_facing_error": { + "is_panic": False, + "message": f"Error querying the database: {message}", + "meta": {"code": "57014", "message": message}, + } + } + ) + + +@pytest.mark.asyncio +async def test_get_team_spend_by_user_statement_timeout_returns_503(mock_db_client): + 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]) + _team_spend_by_user_query(mock_db_client, AsyncMock(side_effect=_statement_timeout_error())) + + with pytest.raises(HTTPException) as exc_info: + 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", + ) + + assert exc_info.value.status_code == 503 + assert "longer than 30s" in str(exc_info.value.detail) class _TeamRowWithOrganization(LiteLLM_TeamTable): diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/TeamUserSpendCard.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/TeamUserSpendCard.test.tsx new file mode 100644 index 00000000000..c8d14ab4461 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/TeamUserSpendCard.test.tsx @@ -0,0 +1,39 @@ +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { render, screen } from "@testing-library/react"; +import React from "react"; +import { describe, expect, it, vi } from "vitest"; + +import { teamSpendByUserCall } from "@/components/networking"; + +import TeamUserSpendCard from "./TeamUserSpendCard"; + +vi.mock("@/components/networking", async (importOriginal) => ({ + ...(await importOriginal()), + teamSpendByUserCall: vi.fn(), +})); + +const renderCard = () => { + const client = new QueryClient({ defaultOptions: { queries: { retry: false, gcTime: 0 } } }); + return render( + + + , + ); +}; + +describe("TeamUserSpendCard", () => { + it("shows the server's error instead of an empty table when the request fails", async () => { + const timeoutMessage = "Spend per user took longer than 30s. Pick a shorter date range or fewer teams"; + vi.mocked(teamSpendByUserCall).mockRejectedValueOnce(new Error(timeoutMessage)); + + renderCard(); + + expect(await screen.findByText(timeoutMessage)).toBeInTheDocument(); + expect(screen.queryByText("No user spend in this range")).not.toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/TeamUserSpendCard.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/TeamUserSpendCard.tsx index cd488cbd4ea..eb48088f919 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/TeamUserSpendCard.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/TeamUserSpendCard.tsx @@ -73,7 +73,7 @@ const columns: ColumnDef[] = [ const TeamUserSpendCard: React.FC = ({ accessToken, startTime, endTime, teamIds }) => { const hasTeams = teamIds.length > 0; - const { data, isLoading } = useQuery({ + const { data, isLoading, error } = useQuery({ queryKey: ["teamSpendByUser", startTime?.toISOString(), endTime?.toISOString(), teamIds], queryFn: () => accessToken && startTime && endTime ? teamSpendByUserCall(accessToken, startTime, endTime, teamIds) : null, @@ -100,15 +100,21 @@ const TeamUserSpendCard: React.FC = ({ accessToken, star

Attributed per request from spend logs, so it includes JWT/SSO traffic that does not use a virtual key

- + {error !== null ? ( +

+ {error instanceof Error ? error.message : "Spend per user failed to load"} +

+ ) : ( + + )} ); };