fix(proxy): cancel /team/spend/by_user in Postgres after 30s (#49922)

* fix(proxy): let the Admin UI cancel /team/spend/by_user at 60s with an opt-in statement timeout

GET /team/spend/by_user takes an optional timeout_seconds (1 to 600). When set, the raw
LiteLLM_SpendLogs aggregation runs under SET LOCAL statement_timeout and Postgres cancels it,
and the route returns 503 instead of letting the query run for hours. Callers that omit the
param are unchanged. The Admin UI sends timeout_seconds=60 and the Spend Per User Within Team
card shows the server's error instead of an empty table

* fix(proxy): cancel /team/spend/by_user in Postgres after 30s

Prisma's HTTP client already gives up on the statement at 30s, so the
query now runs under SET LOCAL statement_timeout = 30000 for every caller
and the opt-in timeout_seconds param is gone. The Admin UI card shows the
request error instead of the empty-state text.

* fix(proxy): keep /team/spend/by_user on the read replica inside its timeout transaction

---------

Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Kerry Lu 2026-11-03 14:03:28 -08:00 • committed by GitHub
parent 6db7001c05
commit 391c8013c9
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 180 additions and 24 deletions

View file

@ -34,6 +34,8 @@ _CONNECTION_CAPACITY_PHRASES: Final = (
_PRISMA_POOL_TIMEOUT_CODE: Final = "P2024" _PRISMA_POOL_TIMEOUT_CODE: Final = "P2024"
_LOCK_TIMEOUT_SQLSTATE: Final = "55P03" _LOCK_TIMEOUT_SQLSTATE: Final = "55P03"
_LOCK_TIMEOUT_PHRASE: Final = "canceling statement due to lock timeout" _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]: def exception_chain(e: BaseException) -> Iterator[BaseException]:
@ -302,6 +304,17 @@ class PrismaDBExceptionHandler:
return True return True
return _LOCK_TIMEOUT_PHRASE in str(e).lower() 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 @staticmethod
def postgres_sqlstate(e: Exception) -> str | None: def postgres_sqlstate(e: Exception) -> str | None:
"""The SQLSTATE Postgres attached to a failed statement, as prisma surfaces it, or None.""" """The SQLSTATE Postgres attached to a failed statement, as prisma surfaces it, or None."""

View file

@ -140,6 +140,7 @@ from litellm.proxy.common_utils.user_api_key_cache import (
team_membership_reservation_cache_key, team_membership_reservation_cache_key,
) )
from litellm.proxy.db.db_span import db_span 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.key_management_event_hooks import KeyManagementEventHooks
from litellm.proxy.hooks.model_max_budget_limiter import ( from litellm.proxy.hooks.model_max_budget_limiter import (
RoutedModels, RoutedModels,
@ -7473,6 +7474,19 @@ class _TeamUserSpendDbRow(TypedDict):
failed_requests: ReadOnly[int] 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( @router.get(
"/team/spend/by_user", "/team/spend/by_user",
response_model=TeamUserSpendResponse, response_model=TeamUserSpendResponse,
@ -7522,14 +7536,23 @@ async def get_team_spend_by_user(
own_user_only: Final = scope.api_key_filter is not None 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 () user_param: Final = (user_api_key_dict.user_id or "",) if own_user_only else ()
async with db_span("team_user_spend", "LiteLLM_SpendLogs"): try:
rows: Final[Sequence[_TeamUserSpendDbRow]] = await prisma_client.db.query_raw( async with db_span("team_user_spend", "LiteLLM_SpendLogs"):
_team_user_spend_sql(team_count=len(scoped_team_ids), restrict_to_user=own_user_only), rows: Final = await _team_user_spend_rows(
start_date, prisma_client,
end_date, _team_user_spend_sql(team_count=len(scoped_team_ids), restrict_to_user=own_user_only),
*scoped_team_ids, (start_date, end_date, *scoped_team_ids, *user_param),
*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( results: Final = tuple(
TeamUserSpendRow( TeamUserSpendRow(
team_id=row["team_id"], team_id=row["team_id"],

View file

@ -953,6 +953,21 @@ def test_is_lock_timeout_error_excludes_other_failures(error: Exception) -> None
assert PrismaDBExceptionHandler.is_lock_timeout_error(error) is False 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
@pytest.mark.parametrize( @pytest.mark.parametrize(
("body", "retryable"), ("body", "retryable"),
[ [

View file

@ -13,6 +13,7 @@ import litellm
import pytest import pytest
from fastapi import HTTPException, Request from fastapi import HTTPException, Request
from fastapi.testclient import TestClient from fastapi.testclient import TestClient
from prisma.errors import DataError
from pydantic import ValidationError from pydantic import ValidationError
from litellm._uuid import uuid from litellm._uuid import uuid
@ -17278,6 +17279,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 @pytest.mark.asyncio
async def test_get_team_spend_by_user_admin_groups_spend_logs_by_team_and_user(mock_db_client): 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 from litellm.proxy.management_endpoints.team_endpoints import get_team_spend_by_user
@ -17286,12 +17309,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"), []) 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"), []) 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.litellm_teamtable.find_many = AsyncMock(return_value=[alpha, beta])
mock_db_client.db.query_raw = AsyncMock( transaction = _team_spend_by_user_query(
return_value=[ mock_db_client,
_team_spend_by_user_db_row("team-alpha", "alice", 0.5, 3), AsyncMock(
_team_spend_by_user_db_row("team-alpha", "bob", 0.25, 2), return_value=[
_team_spend_by_user_db_row("team-beta", "alice", 0.1, 1), _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( response = await get_team_spend_by_user(
@ -17301,7 +17327,9 @@ async def test_get_team_spend_by_user_admin_groups_spend_logs_by_team_and_user(m
end_date="2026-09-04", 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 params == ["2026-09-01", "2026-09-04", "team-alpha", "team-beta"]
assert 'FROM "LiteLLM_SpendLogs" sl' in sql assert 'FROM "LiteLLM_SpendLogs" sl' in sql
assert 'sl."startTime" >= $1::timestamp' in sql assert 'sl."startTime" >= $1::timestamp' in sql
@ -17328,7 +17356,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) 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"), []) 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.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( mock_db_client.writer_db.litellm_usertable.find_unique = AsyncMock(
return_value=_team_spend_by_user_caller("alice", ["team-alpha"]) return_value=_team_spend_by_user_caller("alice", ["team-alpha"])
) )
@ -17337,7 +17365,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" 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 params == ["2026-09-01", "2026-09-04", "team-alpha"]
assert 'sl."user" = $' not in sql assert 'sl."user" = $' not in sql
@ -17350,7 +17378,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"]) 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_teamtable.find_many = AsyncMock(return_value=[alpha])
mock_db_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) 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( mock_db_client.writer_db.litellm_usertable.find_unique = AsyncMock(
return_value=_team_spend_by_user_caller("bob", ["team-alpha"]) return_value=_team_spend_by_user_caller("bob", ["team-alpha"])
) )
@ -17359,7 +17389,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" 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 params == ["2026-09-01", "2026-09-04", "team-alpha", "bob"]
assert "sl.team_id IN ($3)" in sql assert "sl.team_id IN ($3)" in sql
assert 'AND sl."user" = $4' in sql assert 'AND sl."user" = $4' in sql
@ -17371,7 +17401,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 from litellm.proxy.management_endpoints.team_endpoints import get_team_spend_by_user
caller = UserAPIKeyAuth(user_id="bob", user_role=LitellmUserRoles.INTERNAL_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( mock_db_client.writer_db.litellm_usertable.find_unique = AsyncMock(
return_value=_team_spend_by_user_caller("bob", ["team-alpha"]) return_value=_team_spend_by_user_caller("bob", ["team-alpha"])
) )
@ -17382,7 +17412,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 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 @pytest.mark.asyncio
@ -17400,7 +17430,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): 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 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) admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN)
with pytest.raises(HTTPException) as exc_info: with pytest.raises(HTTPException) as exc_info:
@ -17410,7 +17440,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 exc_info.value.status_code == 400
assert expected_error in str(exc_info.value.detail) 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): class _TeamRowWithOrganization(LiteLLM_TeamTable):

View file

@ -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<typeof import("@/components/networking")>()),
teamSpendByUserCall: vi.fn(),
}));
const renderCard = () => {
const client = new QueryClient({ defaultOptions: { queries: { retry: false, gcTime: 0 } } });
return render(
<QueryClientProvider client={client}>
<TeamUserSpendCard
accessToken="sk-test"
startTime={new Date("2026-09-01T00:00:00Z")}
endTime={new Date("2026-09-30T00:00:00Z")}
teamIds={["team-alpha"]}
/>
</QueryClientProvider>,
);
};
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();
});
});

View file

@ -73,7 +73,7 @@ const columns: ColumnDef<TeamUserSpendRow>[] = [
const TeamUserSpendCard: React.FC<TeamUserSpendCardProps> = ({ accessToken, startTime, endTime, teamIds }) => { const TeamUserSpendCard: React.FC<TeamUserSpendCardProps> = ({ accessToken, startTime, endTime, teamIds }) => {
const hasTeams = teamIds.length > 0; const hasTeams = teamIds.length > 0;
const { data, isLoading } = useQuery({ const { data, isLoading, error } = useQuery({
queryKey: ["teamSpendByUser", startTime?.toISOString(), endTime?.toISOString(), teamIds], queryKey: ["teamSpendByUser", startTime?.toISOString(), endTime?.toISOString(), teamIds],
queryFn: () => queryFn: () =>
accessToken && startTime && endTime ? teamSpendByUserCall(accessToken, startTime, endTime, teamIds) : null, accessToken && startTime && endTime ? teamSpendByUserCall(accessToken, startTime, endTime, teamIds) : null,
@ -105,6 +105,8 @@ const TeamUserSpendCard: React.FC<TeamUserSpendCardProps> = ({ accessToken, star
data={rows} data={rows}
getRowId={teamUserSpendRowId} getRowId={teamUserSpendRowId}
isLoading={isLoading} isLoading={isLoading}
isError={error !== null}
errorMessage={error instanceof Error ? error.message : "Spend per user failed to load"}
maxBodyHeight={320} maxBodyHeight={320}
noDataMessage={teamIds.length === 0 ? "Select a team to see spend per user" : "No user spend in this range"} noDataMessage={teamIds.length === 0 ? "Select a team to see spend per user" : "No user spend in this range"}
size="compact" size="compact"