diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 1435033e1be..830e3310d04 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -15,6 +15,7 @@ import re import time from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence from dataclasses import dataclass +from datetime import datetime from functools import partial from itertools import chain from types import MappingProxyType @@ -1657,6 +1658,20 @@ async def get_team_member_default_budget( return budget +def resolve_effective_team_member_budget( + member_budget: LiteLLM_BudgetTable | None, + team_default_budget: LiteLLM_BudgetTable | None, + now: datetime, +) -> float | None: + if member_budget is not None and member_budget.max_budget is not None: + return member_budget.effective_max_budget(now) + if team_default_budget is None or team_default_budget.max_budget is None or team_default_budget.max_budget <= 0: + return None + return team_default_budget.max_budget + ( + member_budget.active_temp_budget_increase(now) if member_budget is not None else 0.0 + ) + + async def resolve_default_end_user_budget( prisma_client: PrismaClient, user_api_key_cache: UserApiKeyCache, @@ -6205,29 +6220,24 @@ async def _check_team_member_budget( # Per-member override wins; otherwise fall back to the team-level # default configured via team.metadata["team_member_budget_id"]. - team_member_budget: float | None = None member_budget_row: Final = loaded_membership.litellm_budget_table if loaded_membership is not None else None now: Final = get_utc_datetime() - if member_budget_row is not None and member_budget_row.max_budget is not None: - team_member_budget = member_budget_row.effective_max_budget(now=now) - else: - default_budget_id: Final = (team_object.metadata or {}).get("team_member_budget_id") - if isinstance(default_budget_id, str): - default_budget: Final = await get_team_member_default_budget( - budget_id=default_budget_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - ) - # Treat 0 on the team default as "no cap". - # Per-member rows still respect 0 as an explicit admin disable. - if ( - default_budget is not None - and default_budget.max_budget is not None - and default_budget.max_budget > 0 - ): - team_member_budget = default_budget.max_budget + ( - member_budget_row.active_temp_budget_increase(now=now) if member_budget_row is not None else 0.0 - ) + default_budget_id: Final = (team_object.metadata or {}).get("team_member_budget_id") + team_default_budget: Final = ( + await get_team_member_default_budget( + budget_id=default_budget_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + if (member_budget_row is None or member_budget_row.max_budget is None) + and isinstance(default_budget_id, str) + else None + ) + team_member_budget: Final = resolve_effective_team_member_budget( + member_budget=member_budget_row, + team_default_budget=team_default_budget, + now=now, + ) if team_member_budget is not None: team_member_spend = (loaded_membership.spend if loaded_membership is not None else 0.0) or 0.0 diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 69a5e2e233a..9c734488960 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -48,6 +48,7 @@ from litellm.proxy._types import ( CommonProxyErrors, DeleteTeamRequest, LiteLLM_AuditLogs, + LiteLLM_BudgetTable, LiteLLM_DeletedTeamTable, Litellm_EntityType, LiteLLM_ManagementEndpoint_MetadataFields, @@ -108,6 +109,7 @@ from litellm.proxy.auth.auth_checks import ( # noqa: F401 # legacy module expo get_team_object, get_user_object, invalidate_team_member_spend_state, + resolve_effective_team_member_budget, ) from litellm.proxy.auth.auth_utils import ( enforce_batch_limits_are_admin_only, @@ -218,6 +220,7 @@ from litellm.types.proxy.management_endpoints.team_endpoints import ( BulkUpdateTeamMemberPermissionsResponse, GetTeamMemberPermissionsResponse, TeamIdSearchMatch, + TeamListCallerMembership, TeamListItem, TeamListResponse, TeamMemberAddResult, @@ -228,6 +231,7 @@ from litellm.types.proxy.management_endpoints.team_endpoints import ( UpdateTeamMemberPermissionsRequest, ) from litellm.types.utils import BudgetConfig +from litellm.utils import get_utc_datetime if TYPE_CHECKING: from prisma import Prisma @@ -276,6 +280,34 @@ class _TeamIdRow(Protocol): def team_id(self) -> str: ... +class _TeamListBudgetRow(Protocol): + @property + def budget_id(self) -> str: ... + + @property + def budget_reset_at(self) -> datetime | None: ... + + @property + def max_budget(self) -> float | None: ... + + @property + def temp_budget_expiry(self) -> datetime | None: ... + + @property + def temp_budget_increase(self) -> float | None: ... + + +class _TeamListMembershipRow(Protocol): + @property + def team_id(self) -> str: ... + + @property + def spend(self) -> float | None: ... + + @property + def litellm_budget_table(self) -> _TeamListBudgetRow | None: ... + + class _CacheableTeamRow(_TeamIdRow, _ModelDumpRow, Protocol): ... @@ -5531,10 +5563,15 @@ def _convert_teams_to_response_models( members_count = len(members_with_roles) keys_count = counts.get(team_dict.get("team_id") or "", 0) team_list.append( - TeamListItem( - **team_dict, - members_count=members_count, - keys_count=keys_count, + TeamListItem.model_validate( + cast( + object, + { + **team_dict, + "members_count": members_count, + "keys_count": keys_count, + }, + ) # cast-ok: TeamListItem validates the Prisma model dump ) ) return team_list @@ -5565,6 +5602,120 @@ async def _get_keys_count_by_team( return {row["team_id"]: row.get("_count", {}).get("team_id", 0) for row in grouped if row.get("team_id")} +def _team_list_budget_model(budget_row: _TeamListBudgetRow | None) -> LiteLLM_BudgetTable | None: + if budget_row is None: + return None + return LiteLLM_BudgetTable( + max_budget=budget_row.max_budget, + temp_budget_expiry=budget_row.temp_budget_expiry, + temp_budget_increase=budget_row.temp_budget_increase, + ) + + +def _team_member_default_budget_id( + team: TeamListItem, + caller_user_id: str, + memberships_by_team: Mapping[str, _TeamListMembershipRow], +) -> str | None: + membership_row: Final = memberships_by_team.get(team.team_id) + if membership_row is None and not any(member.user_id == caller_user_id for member in team.members_with_roles): + return None + member_budget: Final = _team_list_budget_model( + membership_row.litellm_budget_table if membership_row is not None else None + ) + if member_budget is not None and member_budget.max_budget is not None: + return None + metadata: Final = cast("dict[str, object] | None", team.metadata) # cast-ok: team metadata is stored as a JSON object + if metadata is None: + return None + budget_id: Final = metadata.get("team_member_budget_id") + return budget_id if isinstance(budget_id, str) else None + + +def _team_list_caller_membership( + team: TeamListItem, + caller_user_id: str, + memberships_by_team: Mapping[str, _TeamListMembershipRow], + budgets_by_id: Mapping[str, _TeamListBudgetRow], + now: datetime, +) -> TeamListCallerMembership | None: + membership_row: Final = memberships_by_team.get(team.team_id) + if membership_row is None and not any(member.user_id == caller_user_id for member in team.members_with_roles): + return None + member_budget_row: Final = membership_row.litellm_budget_table if membership_row is not None else None + member_budget: Final = _team_list_budget_model(member_budget_row) + default_budget_id: Final = _team_member_default_budget_id(team, caller_user_id, memberships_by_team) + default_budget_row: Final = budgets_by_id.get(default_budget_id) if default_budget_id is not None else None + team_default_budget: Final = _team_list_budget_model(default_budget_row) + max_budget: Final = resolve_effective_team_member_budget( + member_budget=member_budget, + team_default_budget=team_default_budget, + now=now, + ) + budget_reset_at: Final = ( + member_budget_row.budget_reset_at + if member_budget_row is not None and member_budget_row.max_budget is not None + else default_budget_row.budget_reset_at + if max_budget is not None and default_budget_row is not None + else None + ) + return TeamListCallerMembership( + spend=(membership_row.spend if membership_row is not None else 0.0) or 0.0, + max_budget=max_budget, + budget_reset_at=budget_reset_at, + ) + + +async def _enrich_team_list_with_caller_membership( + prisma_client: PrismaClient, + teams: Sequence[TeamListItem], + caller_user_id: str | None, + now: datetime, +) -> list[TeamListItem]: + if caller_user_id is None: + return [team.model_copy(update={"caller_membership": None}) for team in teams] + if not teams: + return [] + + page_team_ids: Final = tuple(team.team_id for team in teams) + membership_rows: Final[Sequence[_TeamListMembershipRow]] = await _team_membership_db(prisma_client).find_many( + # bounded-ok: page_team_ids is capped at 100 teams by list_team_v2 + where={"team_id": {"in": page_team_ids}, "user_id": caller_user_id}, + include={"litellm_budget_table": True}, + ) + memberships_by_team: Final = {row.team_id: row for row in membership_rows} + default_budget_ids: Final = tuple( + dict.fromkeys( + budget_id + for team in teams + if (budget_id := _team_member_default_budget_id(team, caller_user_id, memberships_by_team)) is not None + ) + ) + default_budget_rows: Final[Sequence[_TeamListBudgetRow]] = ( + await _budget_db(prisma_client).find_many( + # bounded-ok: default_budget_ids is derived from this page of at most 100 teams + where={"budget_id": {"in": default_budget_ids}}, + ) + if default_budget_ids + else () + ) + budgets_by_id: Final = {row.budget_id: row for row in default_budget_rows} + return [ + team.model_copy( + update={ + "caller_membership": _team_list_caller_membership( + team=team, + caller_user_id=caller_user_id, + memberships_by_team=memberships_by_team, + budgets_by_id=budgets_by_id, + now=now, + ) + } + ) + for team in teams + ] + + async def _enforce_list_team_v2_access( user_api_key_dict: UserAPIKeyAuth, user_id: str | None, @@ -5678,7 +5829,10 @@ async def list_team_v2( sort_order: str = fastapi.Query(default="asc", description="Sort order ('asc' or 'desc')"), status: str | None = fastapi.Query(default=None, description="Filter by status (e.g. 'deleted')"), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), -): +) -> dict[ + str, + int | list[TeamListItem] | list[TeamListItem | LiteLLM_TeamTable | LiteLLM_DeletedTeamTable], +]: """ Get a paginated list of teams with filtering and sorting options. @@ -5801,7 +5955,7 @@ async def list_team_v2( keys_count_by_team = await _get_keys_count_by_team(prisma_client, teams) # Convert Prisma models to response models with members_count and keys_count - team_list = _convert_teams_to_response_models(teams, use_deleted_table, keys_count_by_team=keys_count_by_team) + team_list: Final = _convert_teams_to_response_models(teams, use_deleted_table, keys_count_by_team=keys_count_by_team) # Resolve resources inherited from access groups (single batch query) if not use_deleted_table: @@ -5823,8 +5977,18 @@ async def list_team_v2( {a for group in team_groups for a in (group.access_agent_ids or [])} ) + team_list_with_caller_membership: Final = ( + await _enrich_team_list_with_caller_membership( + prisma_client=prisma_client, + teams=tuple(team for team in team_list if isinstance(team, TeamListItem)), + caller_user_id=user_api_key_dict.user_id, + now=get_utc_datetime(), + ) + if not use_deleted_table + else team_list + ) return { - "teams": team_list, + "teams": team_list_with_caller_membership, "total": total_count, "page": page, "page_size": page_size, diff --git a/litellm/types/proxy/management_endpoints/team_endpoints.py b/litellm/types/proxy/management_endpoints/team_endpoints.py index 6731b826ebd..ba831aa7ba9 100644 --- a/litellm/types/proxy/management_endpoints/team_endpoints.py +++ b/litellm/types/proxy/management_endpoints/team_endpoints.py @@ -1,3 +1,4 @@ +from datetime import datetime from typing import Any, Final, Literal from pydantic import ConfigDict, Field, field_validator, model_validator @@ -74,11 +75,22 @@ class BulkUpdateTeamMemberPermissionsResponse(LiteLLMBaseModel): permissions_appended: list[str] | None = None +class TeamListCallerMembership(LiteLLMBaseModel): + """The calling user's own spend and effective member budget in a team.""" + + spend: float = 0.0 + max_budget: float | None = None + budget_reset_at: datetime | None = None + + model_config = ConfigDict(frozen=True) + + class TeamListItem(LiteLLM_TeamTable): """A team item in the paginated list response, enriched with computed fields.""" members_count: int = 0 keys_count: int = 0 + caller_membership: TeamListCallerMembership | None = None # Resources inherited from access groups (separate from direct assignments) access_group_models: list[str] | None = None access_group_mcp_server_ids: list[str] | None = None diff --git a/tests/proxy_behavior/management/test_team_list_v2.py b/tests/proxy_behavior/management/test_team_list_v2.py index 81178ad73c0..b5597f02e97 100644 --- a/tests/proxy_behavior/management/test_team_list_v2.py +++ b/tests/proxy_behavior/management/test_team_list_v2.py @@ -1,8 +1,9 @@ -from typing import FrozenSet, Optional +from typing import Final +import httpx import pytest -from .actors import Actor +from .actors import Actor, World pytestmark = pytest.mark.asyncio(loop_scope="session") @@ -63,7 +64,7 @@ _BARE = [ async def test_team_list_v2_bare( actor: Actor, expected_status: int, - expected_visible: Optional[FrozenSet[str]], + expected_visible: frozenset[str] | None, proxy_client, world, ): @@ -98,7 +99,7 @@ _OWN = { "actor,expected_visible", list(_OWN.items()), ids=[a.value for a in _OWN] ) async def test_team_list_v2_own_user_id_query( - actor: Actor, expected_visible: FrozenSet[str], proxy_client, world + actor: Actor, expected_visible: frozenset[str], proxy_client, world ): caller = world.keys[actor] visible = _seeded( @@ -139,3 +140,31 @@ async def test_team_list_v2_invalid_status_is_400(proxy_client, world): headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"}, ) assert resp.status_code == 400, resp.text + + +async def test_team_list_v2_returns_caller_membership_only_for_members( + proxy_client: httpx.AsyncClient, + world: World, +) -> None: + member: Final = world.keys[Actor.TEAM_ADMIN] + member_response: Final = await proxy_client.get( + f"/v2/team/list?team_id={world.team_alpha_id}&user_id={member.user_id}", + headers={"Authorization": f"Bearer {member.cleartext}"}, + ) + + assert member_response.status_code == 200, member_response.text + member_team: Final = member_response.json()["teams"][0] + assert member_team["caller_membership"] == { + "spend": 0.0, + "max_budget": None, + "budget_reset_at": None, + } + + admin: Final = world.keys[Actor.PROXY_ADMIN] + admin_response: Final = await proxy_client.get( + f"/v2/team/list?team_id={world.team_alpha_id}", + headers={"Authorization": f"Bearer {admin.cleartext}"}, + ) + + assert admin_response.status_code == 200, admin_response.text + assert admin_response.json()["teams"][0]["caller_membership"] is None diff --git a/tests/unit/proxy/auth/test_team_member_budget.py b/tests/unit/proxy/auth/test_team_member_budget.py index b38a953d189..6055e4f0356 100644 --- a/tests/unit/proxy/auth/test_team_member_budget.py +++ b/tests/unit/proxy/auth/test_team_member_budget.py @@ -3,8 +3,11 @@ Unit tests for team member budget checks in common_checks. These tests verify the team member budget enforcement without requiring a proxy server. """ -import pytest +from datetime import datetime, timezone +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch + +import pytest from fastapi import Request import litellm @@ -15,7 +18,10 @@ from litellm.proxy._types import ( LiteLLM_UserTable, UserAPIKeyAuth, ) -from litellm.proxy.auth.auth_checks import common_checks, get_team_membership +from litellm.proxy.auth.auth_checks import ( + common_checks, + resolve_effective_team_member_budget, +) @pytest.mark.asyncio @@ -436,3 +442,48 @@ async def test_team_member_budget_check_personal_key_not_team(): # Should pass and get_team_membership should not be called assert result is True mock_get_team_membership.assert_not_called() + + +def test_resolve_effective_team_member_budget_preserves_member_and_default_precedence() -> None: + now: Final = datetime(2030, 1, 1, tzinfo=timezone.utc) + member_budget: Final = LiteLLM_BudgetTable( + max_budget=100, + temp_budget_increase=10, + temp_budget_expiry=datetime(2030, 1, 2, tzinfo=timezone.utc), + ) + expired_member_budget: Final = LiteLLM_BudgetTable( + max_budget=100, + temp_budget_increase=10, + temp_budget_expiry=datetime(2029, 12, 31, tzinfo=timezone.utc), + ) + active_inherited_increase: Final = LiteLLM_BudgetTable( + max_budget=None, + temp_budget_increase=10, + temp_budget_expiry=datetime(2030, 1, 2, tzinfo=timezone.utc), + ) + expired_inherited_increase: Final = LiteLLM_BudgetTable( + max_budget=None, + temp_budget_increase=10, + temp_budget_expiry=datetime(2029, 12, 31, tzinfo=timezone.utc), + ) + + assert resolve_effective_team_member_budget(member_budget, None, now) == 110 + assert resolve_effective_team_member_budget(expired_member_budget, None, now) == 100 + assert resolve_effective_team_member_budget( + active_inherited_increase, + LiteLLM_BudgetTable(max_budget=50), + now, + ) == 60 + assert resolve_effective_team_member_budget( + expired_inherited_increase, + LiteLLM_BudgetTable(max_budget=50), + now, + ) == 50 + assert resolve_effective_team_member_budget( + LiteLLM_BudgetTable(max_budget=0), + LiteLLM_BudgetTable(max_budget=50), + now, + ) == 0 + assert resolve_effective_team_member_budget(None, LiteLLM_BudgetTable(max_budget=50), now) == 50 + assert resolve_effective_team_member_budget(None, LiteLLM_BudgetTable(max_budget=0), now) is None + assert resolve_effective_team_member_budget(None, None, now) is None diff --git a/tests/unit/proxy/management_endpoints/test_team_endpoints.py b/tests/unit/proxy/management_endpoints/test_team_endpoints.py index 5e9c350431c..0251403df83 100644 --- a/tests/unit/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_team_endpoints.py @@ -47,6 +47,7 @@ from litellm.proxy.management_endpoints.team_admin_field_permissions import ( ) from litellm.proxy.management_endpoints.team_endpoints import ( _STRIP_DELETED_TEAM_FROM_USERS_SQL, + _enrich_team_list_with_caller_membership, GetTeamMemberPermissionsResponse, UpdateTeamMemberPermissionsRequest, _build_team_list_where_conditions, @@ -81,6 +82,8 @@ from litellm.router import Router from litellm.types.proxy.management_endpoints.team_endpoints import ( BulkTeamMemberAddRequest, BulkTeamMemberAddResponse, + TeamListCallerMembership, + TeamListItem, TeamMemberAddResult, ) from litellm.types.utils import StandardAuditLogPayload @@ -3821,6 +3824,7 @@ async def test_list_team_v2_security_check_non_admin_user_own_teams(): ] mock_db.litellm_teamtable.find_many = AsyncMock(return_value=mock_teams) mock_db.litellm_teamtable.count = AsyncMock(return_value=2) + mock_db.litellm_teammembership.find_many = AsyncMock(return_value=[]) mock_db.litellm_verificationtoken.group_by = AsyncMock(return_value=[]) with patch( @@ -3878,6 +3882,7 @@ async def test_list_team_v2_security_check_admin_user(): ] mock_db.litellm_teamtable.find_many = AsyncMock(return_value=mock_teams) mock_db.litellm_teamtable.count = AsyncMock(return_value=2) + mock_db.litellm_teammembership.find_many = AsyncMock(return_value=[]) mock_db.litellm_verificationtoken.group_by = AsyncMock(return_value=[]) # Should NOT raise an exception @@ -4021,6 +4026,7 @@ async def test_list_team_v2_includes_litellm_model_table(): side_effect=lambda **kw: [_team_row("team_1", kw.get("include"))] ) mock_db.litellm_teamtable.count = AsyncMock(return_value=1) + mock_db.litellm_teammembership.find_many = AsyncMock(return_value=[]) mock_db.litellm_verificationtoken.group_by = AsyncMock(return_value=[]) result = await list_team_v2( @@ -4114,6 +4120,7 @@ async def test_list_team_v2_org_admin_sees_org_teams(): } mock_db.litellm_teamtable.find_many = AsyncMock(return_value=[mock_team]) mock_db.litellm_teamtable.count = AsyncMock(return_value=1) + mock_db.litellm_teammembership.find_many = AsyncMock(return_value=[]) mock_db.litellm_verificationtoken.group_by = AsyncMock(return_value=[]) mock_db.litellm_usertable.find_unique = AsyncMock(return_value=mock_user) @@ -4214,6 +4221,7 @@ async def test_list_team_v2_org_admin_own_user_id_sees_all_org_teams(): return_value=[mock_team_1, mock_team_2] ) mock_db.litellm_teamtable.count = AsyncMock(return_value=2) + mock_db.litellm_teammembership.find_many = AsyncMock(return_value=[]) mock_db.litellm_verificationtoken.group_by = AsyncMock(return_value=[]) mock_db.litellm_usertable.find_unique = AsyncMock(return_value=mock_user) @@ -4338,6 +4346,7 @@ async def test_list_team_v2_org_admin_own_query_keeps_memberships_in_other_orgs( prisma_client = MagicMock() prisma_client.db.litellm_teamtable.find_many = AsyncMock(side_effect=find_many) prisma_client.db.litellm_teamtable.count = AsyncMock(side_effect=count) + prisma_client.db.litellm_teammembership.find_many = AsyncMock(return_value=[]) prisma_client.db.litellm_verificationtoken.group_by = AsyncMock(return_value=[]) prisma_client.writer_db.litellm_usertable.find_unique = AsyncMock( return_value=LiteLLM_UserTable( @@ -4599,6 +4608,7 @@ async def test_list_team_v2_org_admin_with_user_id_returns_user_teams(): } mock_db.litellm_teamtable.find_many = AsyncMock(return_value=[mock_team]) mock_db.litellm_teamtable.count = AsyncMock(return_value=1) + mock_db.litellm_teammembership.find_many = AsyncMock(return_value=[]) mock_db.litellm_verificationtoken.group_by = AsyncMock(return_value=[]) result = await list_team_v2( @@ -4868,6 +4878,7 @@ async def test_list_team_v2_populates_keys_count(): mock_db.litellm_teamtable.find_many = AsyncMock(return_value=[team_a, team_b]) mock_db.litellm_teamtable.count = AsyncMock(return_value=2) + mock_db.litellm_teammembership.find_many = AsyncMock(return_value=[]) mock_db.litellm_verificationtoken.group_by = AsyncMock( return_value=[ {"team_id": "team_a", "_count": {"team_id": 3}}, @@ -4888,12 +4899,251 @@ async def test_list_team_v2_populates_keys_count(): by_id = {t.team_id: t for t in result["teams"]} assert by_id["team_a"].keys_count == 3 assert by_id["team_b"].keys_count == 0 - # The aggregate is one batched query, filtered by the page's team IDs. group_by_kwargs = mock_db.litellm_verificationtoken.group_by.call_args.kwargs assert group_by_kwargs["by"] == ["team_id"] assert group_by_kwargs["where"] == {"team_id": {"in": ["team_a", "team_b"]}} assert group_by_kwargs["count"] == {"team_id": True} + membership_where = mock_db.litellm_teammembership.find_many.call_args.kwargs["where"] + assert membership_where == { + "team_id": {"in": ("team_a", "team_b")}, + "user_id": "admin_user_123", + } + + +def _team_list_prisma( + membership_rows: Sequence[LiteLLM_TeamMembership], + budget_rows: Sequence[LiteLLM_BudgetTableFull], +) -> tuple[MagicMock, AsyncMock, AsyncMock]: + prisma_client: Final = MagicMock() + membership_query: Final = AsyncMock(return_value=tuple(membership_rows)) + budget_query: Final = AsyncMock(return_value=tuple(budget_rows)) + prisma_client.db.litellm_teammembership.find_many = membership_query + prisma_client.db.litellm_budgettable.find_many = budget_query + return prisma_client, membership_query, budget_query + + +def _team_list_member( + team_id: str, + user_id: str = "caller", + default_budget_id: str | None = None, +) -> TeamListItem: + metadata: Final = {"team_member_budget_id": default_budget_id} if default_budget_id is not None else None + return TeamListItem( + team_id=team_id, + metadata=metadata, + members_with_roles=[Member(user_id=user_id, role="user")], + ) + + +def _budget_table( + budget_id: str, + max_budget: float | None, + budget_reset_at: datetime | None = None, + temp_budget_increase: float | None = None, + temp_budget_expiry: datetime | None = None, +) -> LiteLLM_BudgetTableFull: + return LiteLLM_BudgetTableFull( + budget_id=budget_id, + max_budget=max_budget, + budget_reset_at=budget_reset_at, + created_at=datetime(2026, 1, 1, tzinfo=timezone.utc), + temp_budget_increase=temp_budget_increase, + temp_budget_expiry=temp_budget_expiry, + ) + + +@pytest.mark.asyncio +async def test_team_list_caller_membership_uses_member_spend_budget_and_reset_date() -> None: + reset_at: Final = datetime(2026, 10, 20, tzinfo=timezone.utc) + member_budget: Final = _budget_table("member-budget", 100, reset_at) + membership: Final = LiteLLM_TeamMembership( + user_id="caller", + team_id="team-a", + spend=50, + litellm_budget_table=member_budget, + ) + prisma_client, _, budget_query = _team_list_prisma([membership], []) + + teams: Final = await _enrich_team_list_with_caller_membership( + prisma_client, + [_team_list_member("team-a")], + "caller", + datetime(2026, 10, 1, tzinfo=timezone.utc), + ) + + assert teams[0].caller_membership == TeamListCallerMembership( + spend=50, + max_budget=100, + budget_reset_at=reset_at, + ) + budget_query.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_team_list_caller_membership_includes_roster_members_without_membership_rows() -> None: + prisma_client, membership_query, budget_query = _team_list_prisma((), ()) + + teams: Final = await _enrich_team_list_with_caller_membership( + prisma_client, + [_team_list_member("team-a")], + "caller", + datetime(2026, 10, 1, tzinfo=timezone.utc), + ) + + assert teams[0].caller_membership == TeamListCallerMembership( + spend=0.0, + max_budget=None, + budget_reset_at=None, + ) + membership_query.assert_awaited_once() + budget_query.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_team_list_caller_membership_uses_team_default_budget_and_reset_date() -> None: + reset_at: Final = datetime(2026, 10, 20, tzinfo=timezone.utc) + membership: Final = LiteLLM_TeamMembership(user_id="caller", team_id="team-a", spend=12) + prisma_client, _, budget_query = _team_list_prisma( + [membership], + [_budget_table("team-budget", 50, reset_at)], + ) + + teams: Final = await _enrich_team_list_with_caller_membership( + prisma_client, + [_team_list_member("team-a", default_budget_id="team-budget")], + "caller", + datetime(2026, 10, 1, tzinfo=timezone.utc), + ) + + assert teams[0].caller_membership == TeamListCallerMembership( + spend=12, + max_budget=50, + budget_reset_at=reset_at, + ) + budget_query.assert_awaited_once_with(where={"budget_id": {"in": ("team-budget",)}}) + + +@pytest.mark.asyncio +async def test_team_list_caller_membership_ignores_zero_team_default_budget() -> None: + membership: Final = LiteLLM_TeamMembership(user_id="caller", team_id="team-a") + prisma_client, _, _ = _team_list_prisma( + [membership], + [_budget_table("team-budget", 0, datetime(2026, 10, 20, tzinfo=timezone.utc))], + ) + + teams: Final = await _enrich_team_list_with_caller_membership( + prisma_client, + [_team_list_member("team-a", default_budget_id="team-budget")], + "caller", + datetime(2026, 10, 1, tzinfo=timezone.utc), + ) + + assert teams[0].caller_membership == TeamListCallerMembership(spend=0.0, max_budget=None, budget_reset_at=None) + + +@pytest.mark.parametrize( + ("expiry", "expected_budget"), + ( + (datetime(2030, 1, 2, tzinfo=timezone.utc), 110), + (datetime(2029, 12, 31, tzinfo=timezone.utc), 100), + ), +) +@pytest.mark.asyncio +async def test_team_list_caller_membership_applies_only_active_temporary_increases( + expiry: datetime, + expected_budget: float, +) -> None: + member_budget: Final = _budget_table( + "member-budget", + 100, + datetime(2030, 2, 1, tzinfo=timezone.utc), + temp_budget_increase=10, + temp_budget_expiry=expiry, + ) + membership: Final = LiteLLM_TeamMembership( + user_id="caller", + team_id="team-a", + litellm_budget_table=member_budget, + ) + prisma_client, _, _ = _team_list_prisma([membership], []) + + teams: Final = await _enrich_team_list_with_caller_membership( + prisma_client, + [_team_list_member("team-a")], + "caller", + datetime(2030, 1, 1, tzinfo=timezone.utc), + ) + + assert teams[0].caller_membership is not None + assert teams[0].caller_membership.max_budget == expected_budget + assert teams[0].caller_membership.budget_reset_at == datetime(2030, 2, 1, tzinfo=timezone.utc) + + +@pytest.mark.asyncio +async def test_team_list_caller_membership_is_null_for_nonmember_admins() -> None: + prisma_client, membership_query, budget_query = _team_list_prisma([], []) + team: Final = TeamListItem( + team_id="team-a", + metadata={"team_member_budget_id": "team-budget"}, + members_with_roles=[Member(user_id="another-user", role="user")], + ) + + teams: Final = await _enrich_team_list_with_caller_membership( + prisma_client, + [team], + "caller", + datetime(2026, 10, 1, tzinfo=timezone.utc), + ) + + assert teams[0].caller_membership is None + membership_query.assert_awaited_once() + budget_query.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_team_list_caller_membership_skips_queries_without_caller_user_id() -> None: + prisma_client, membership_query, budget_query = _team_list_prisma((), ()) + + teams: Final = await _enrich_team_list_with_caller_membership( + prisma_client, + [_team_list_member("team-a")], + None, + datetime(2026, 10, 1, tzinfo=timezone.utc), + ) + + assert teams[0].caller_membership is None + membership_query.assert_not_awaited() + budget_query.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_team_list_caller_membership_batches_page_membership_and_budget_queries() -> None: + memberships: Final = ( + LiteLLM_TeamMembership(user_id="caller", team_id="team-a"), + LiteLLM_TeamMembership(user_id="caller", team_id="team-b"), + ) + prisma_client, membership_query, budget_query = _team_list_prisma( + memberships, + [_budget_table("budget-a", 50), _budget_table("budget-b", 75)], + ) + + teams: Final = await _enrich_team_list_with_caller_membership( + prisma_client, + [ + _team_list_member("team-a", default_budget_id="budget-a"), + _team_list_member("team-b", default_budget_id="budget-b"), + ], + "caller", + datetime(2026, 10, 1, tzinfo=timezone.utc), + ) + + assert [team.caller_membership.max_budget for team in teams if team.caller_membership is not None] == [50, 75] + membership_query.assert_awaited_once_with( + where={"team_id": {"in": ("team-a", "team-b")}, "user_id": "caller"}, + include={"litellm_budget_table": True}, + ) + budget_query.assert_awaited_once_with(where={"budget_id": {"in": ("budget-a", "budget-b")}}) @pytest.mark.asyncio diff --git a/ui/litellm-dashboard/src/components/TeamsPage/TeamMemberSpendBudgetCell.test.tsx b/ui/litellm-dashboard/src/components/TeamsPage/TeamMemberSpendBudgetCell.test.tsx new file mode 100644 index 00000000000..01be7407968 --- /dev/null +++ b/ui/litellm-dashboard/src/components/TeamsPage/TeamMemberSpendBudgetCell.test.tsx @@ -0,0 +1,92 @@ +import { render, screen } from "@testing-library/react"; +import { describe, expect, it } from "vitest"; + +import { TeamMemberSpendBudgetCell } from "./TeamMemberSpendBudgetCell"; + +describe("TeamMemberSpendBudgetCell", () => { + it("shows team and member budgets with the member reset date", () => { + render( + , + ); + + expect(screen.getByText("Team")).toBeInTheDocument(); + expect(screen.getByText("$600.00")).toBeInTheDocument(); + expect(screen.getByText("/ $1,000.00")).toBeInTheDocument(); + expect(screen.getByText("Member")).toBeInTheDocument(); + expect(screen.getByText("$50.00")).toBeInTheDocument(); + expect(screen.getByText("/ $100.00")).toBeInTheDocument(); + expect(screen.getByText("· resets Oct 20, 2026")).toBeInTheDocument(); + expect(screen.getAllByRole("meter")[0]).toHaveAttribute("aria-valuetext", "Team $600.00 of $1,000.00"); + expect(screen.getAllByRole("meter")[1]).toHaveAttribute("aria-valuetext", "Member $50.00 of $100.00"); + }); + + it("shows the team budget and Unlimited for a member without a budget", () => { + const { container } = render( + , + ); + + expect(screen.getByText("Team")).toBeInTheDocument(); + expect(screen.getByText("$600.00")).toBeInTheDocument(); + expect(screen.getByText("Member")).toBeInTheDocument(); + expect(screen.getByText("/ Unlimited")).toBeInTheDocument(); + expect(screen.getAllByRole("meter")).toHaveLength(1); + expect(container.querySelectorAll('[data-slot="meter-indicator"]')).toHaveLength(1); + }); + + it("omits the reset text when no member reset date exists", () => { + render( + , + ); + + expect(screen.queryByText(/resets/)).not.toBeInTheDocument(); + }); + + it("renders only the team line when caller membership is absent", () => { + render(); + + expect(screen.getByText("Team")).toBeInTheDocument(); + expect(screen.queryByText("Member")).not.toBeInTheDocument(); + expect(screen.getAllByRole("meter")).toHaveLength(1); + }); + + it("uses the over tone when member spend exceeds the budget", () => { + const { container } = render( + , + ); + const indicators = container.querySelectorAll('[data-slot="meter-indicator"]'); + + expect(indicators[1]).toHaveClass("bg-destructive"); + }); + + it("does not render meters for zero or unlimited budgets", () => { + render( + , + ); + + expect(screen.queryByRole("meter")).not.toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/TeamsPage/TeamMemberSpendBudgetCell.tsx b/ui/litellm-dashboard/src/components/TeamsPage/TeamMemberSpendBudgetCell.tsx new file mode 100644 index 00000000000..337ea383130 --- /dev/null +++ b/ui/litellm-dashboard/src/components/TeamsPage/TeamMemberSpendBudgetCell.tsx @@ -0,0 +1,75 @@ +"use client"; + +import { Meter, MeterIndicator, MeterTrack } from "@/components/shared/Meter"; +import { getSpendBudgetMeterTone } from "@/components/shared/table_cells/spend_budget_cell"; +import { formatBudgetReset } from "@/utils/budgetUtils"; +import { formatNumberWithCommas, getSpendString } from "@/utils/dataUtils"; + +interface CallerMembership { + spend: number; + max_budget: number | null; + budget_reset_at: string | null; +} + +interface TeamMemberSpendBudgetCellProps { + teamSpend: number | null | undefined; + teamMaxBudget: number | null | undefined; + callerMembership?: CallerMembership | null; +} + +interface BudgetLineProps { + label: "Team" | "Member"; + spend: number | null | undefined; + maxBudget: number | null | undefined; + resetAt?: string | null; +} + +function BudgetLine({ label, spend, maxBudget, resetAt }: BudgetLineProps) { + const spendValue = typeof spend === "number" && Number.isFinite(spend) ? spend : 0; + const budget = typeof maxBudget === "number" && Number.isFinite(maxBudget) ? maxBudget : null; + const spendText = spendValue === 0 ? "$0.00" : getSpendString(spendValue, 2); + const budgetText = budget === null ? "Unlimited" : `$${formatNumberWithCommas(budget, 2)}`; + const resetText = label === "Member" ? formatBudgetReset(resetAt) : null; + + return ( +
+
+ {label}{" "} + {spendText}{" "} + / {budgetText} + {resetText && · resets {resetText}} +
+ {budget !== null && budget > 0 && ( + + + + + + )} +
+ ); +} + +export function TeamMemberSpendBudgetCell({ + teamSpend, + teamMaxBudget, + callerMembership, +}: TeamMemberSpendBudgetCellProps) { + return ( +
+ + {callerMembership && ( + + )} +
+ ); +} diff --git a/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.test.tsx b/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.test.tsx index ad89b62d611..e007a8bb21a 100644 --- a/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.test.tsx +++ b/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.test.tsx @@ -112,10 +112,30 @@ it("renders a team row with alias, organization, and spend/budget", async () => expect(screen.getByText("Acme Team")).toBeInTheDocument(); expect(screen.getByText("Test Organization")).toBeInTheDocument(); expect(screen.getByText("$42.50")).toBeInTheDocument(); - expect(screen.getByText("of $100.00")).toBeInTheDocument(); + expect(screen.getByText("/ $100.00")).toBeInTheDocument(); }); }); +it("renders caller membership spend and budget in the team row", async () => { + mockUseTeamsTable.mockReturnValue( + teamsResult([ + { + ...mockTeam, + caller_membership: { + spend: 50, + max_budget: 100, + budget_reset_at: "2026-10-20T12:00:00Z", + }, + }, + ]), + ); + renderTable(); + + expect(await screen.findByText("Member")).toBeInTheDocument(); + expect(screen.getByText("$50.00")).toBeInTheDocument(); + expect(screen.getByText("· resets Oct 20, 2026")).toBeInTheDocument(); +}); + it("renders the Resources cell with member, model, and key counts", () => { renderTable(); @@ -154,7 +174,7 @@ describe("sort contract – only backend-sortable columns are sortable", () => { it("sorts by the backend team_alias field (not the label) when the Team header is clicked", async () => { renderTable(); - fireEvent.click(screen.getByText("Team").closest("button") as HTMLElement); + fireEvent.click(within(screen.getByRole("columnheader", { name: "Team" })).getByRole("button")); await waitFor(() => { expect(mockUseTeamsTable).toHaveBeenLastCalledWith(1, 50, expect.objectContaining({ sortBy: "team_alias" })); @@ -171,7 +191,7 @@ describe("sort contract – only backend-sortable columns are sortable", () => { renderTable(); expect(screen.queryByText("Spend / Budget").closest("button")).toBeNull(); // Team and Created are the only sortable headers. - expect(screen.getByText("Team").closest("button")).not.toBeNull(); + expect(within(screen.getByRole("columnheader", { name: "Team" })).getByRole("button")).toBeInTheDocument(); expect(screen.getByText("Created").closest("button")).not.toBeNull(); }); }); diff --git a/ui/litellm-dashboard/src/components/TeamsPage/teamTableColumns.tsx b/ui/litellm-dashboard/src/components/TeamsPage/teamTableColumns.tsx index 05bfccf6417..462019713b4 100644 --- a/ui/litellm-dashboard/src/components/TeamsPage/teamTableColumns.tsx +++ b/ui/litellm-dashboard/src/components/TeamsPage/teamTableColumns.tsx @@ -4,7 +4,7 @@ import { ColumnDef } from "@tanstack/react-table"; import { Copy, KeyRound, Layers, MoreHorizontal, Pencil, Trash2, Users } from "lucide-react"; import { DataTableSortHeader } from "@/components/shared/DataTable"; -import { DateCell, IdentityCell, SpendBudgetCell } from "@/components/shared/table_cells"; +import { DateCell, IdentityCell } from "@/components/shared/table_cells"; import { buttonVariants } from "@/components/ui/button"; import { DropdownMenu, @@ -20,6 +20,7 @@ import { copyToClipboard, formatNumberWithCommas } from "@/utils/dataUtils"; import { Team } from "../key_team_helpers/key_list"; import { Organization } from "../networking"; +import { TeamMemberSpendBudgetCell } from "./TeamMemberSpendBudgetCell"; interface ResourceTone { icon: typeof Users; @@ -215,11 +216,10 @@ export const getTeamTableColumns = ({ size: 200, enableSorting: false, cell: ({ row }) => ( - ), }, diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx b/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx index 8d39a0fd22b..e6f6f34345a 100644 --- a/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx +++ b/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx @@ -21,6 +21,11 @@ export interface Team { keys_count?: number; members_count?: number; members_with_roles: Member[]; + caller_membership?: { + spend: number; + max_budget: number | null; + budget_reset_at: string | null; + } | null; team_member_permissions?: string[] | null; spend: number; access_group_ids?: string[]; diff --git a/ui/litellm-dashboard/src/components/shared/table_cells/spend_budget_cell.tsx b/ui/litellm-dashboard/src/components/shared/table_cells/spend_budget_cell.tsx index 127c7202121..26f7b818514 100644 --- a/ui/litellm-dashboard/src/components/shared/table_cells/spend_budget_cell.tsx +++ b/ui/litellm-dashboard/src/components/shared/table_cells/spend_budget_cell.tsx @@ -12,7 +12,7 @@ interface SpendBudgetCellProps { budgetDecimals?: number; } -const meterTone = (pct: number): "default" | "warning" | "over" => { +export const getSpendBudgetMeterTone = (pct: number): "default" | "warning" | "over" => { if (pct > 100) return "over"; if (pct >= 80) return "warning"; return "default"; @@ -47,7 +47,7 @@ export function SpendBudgetCell({ aria-valuetext={`${spendText} of $${formatNumberWithCommas(budget, budgetDecimals)}`} > - + )} diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index fa90213a486..9aea67066ae 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -46949,6 +46949,21 @@ export interface components { /** Team Id */ team_id: string; }; + /** + * TeamListCallerMembership + * @description The calling user's own spend and effective member budget in a team. + */ + TeamListCallerMembership: { + /** Budget Reset At */ + budget_reset_at?: string | null; + /** Max Budget */ + max_budget?: number | null; + /** + * Spend + * @default 0 + */ + spend: number; + }; /** * TeamListItem * @description A team item in the paginated list response, enriched with computed fields. @@ -46983,6 +46998,7 @@ export interface components { budget_limits?: components["schemas"]["BudgetLimitEntry"][] | null; /** Budget Reset At */ budget_reset_at?: string | null; + caller_membership?: components["schemas"]["TeamListCallerMembership"] | null; /** Created At */ created_at?: string | null; /** Default Team Member Models */