diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 76a51627d0c..70a93676d4e 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -4,7 +4,7 @@ import os from collections.abc import Callable, Mapping from datetime import datetime from types import MappingProxyType -from typing import TYPE_CHECKING, Annotated, Any, Final, Literal, NamedTuple +from typing import TYPE_CHECKING, Annotated, Any, Final, Literal, NamedTuple, TypeAlias import httpx from pydantic import ( @@ -4588,11 +4588,26 @@ class TeamInfoResponseObjectTeamTable(LiteLLM_TeamTable): caller_edit_access: TeamEditAccess = Field(default_factory=TeamEditNone) +TeamMemberBudgetSource: TypeAlias = Literal["team_default", "custom", "none"] + + +class TeamInfoMembership(LiteLLM_TeamMembership): + budget_source: TeamMemberBudgetSource + + class TeamInfoResponseObject(TypedDict): team_id: str team_info: TeamInfoResponseObjectTeamTable keys: list - team_memberships: list[LiteLLM_TeamMembership] + team_memberships: ReadOnly[tuple[TeamInfoMembership, ...]] + + +class TeamMemberResetBudgetResponse(BaseModel): + team_id: str + user_id: str + budget_id: str | None + previous_budget_id: str | None + budget_source: TeamMemberBudgetSource class TeamListResponseObject(LiteLLM_TeamTable): diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 28c12173ea7..b8d95167045 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -80,11 +80,14 @@ from litellm.proxy._types import ( TeamEditNone, TeamEditUnrestricted, TeamInfoMember, + TeamInfoMembership, TeamInfoResponseObject, TeamInfoResponseObjectTeamTable, TeamListResponseObject, TeamMemberAddRequest, + TeamMemberBudgetSource, TeamMemberDeleteRequest, + TeamMemberResetBudgetResponse, TeamMemberUpdateRequest, TeamMemberUpdateResponse, TeamModelAddRequest, @@ -3954,6 +3957,99 @@ async def reset_team_member_spend_fn( } +class _TeamMetadataView(BaseModel): + metadata: Mapping[str, object] | None = None + + +def _team_default_budget_id(team: LiteLLM_TeamTable) -> str | None: + view: Final = _TeamMetadataView.model_validate(team, from_attributes=True) + raw: Final = view.metadata.get("team_member_budget_id") if view.metadata is not None else None + return raw if isinstance(raw, str) else None + + +async def _existing_team_default_budget_id(team: LiteLLM_TeamTable, prisma_client: PrismaClient) -> str | None: + budget_id: Final = _team_default_budget_id(team) + if budget_id is None: + return None + row: Final = await _budget_db(prisma_client).find_unique( + where={"budget_id": budget_id}, # mutable-ok: prisma client requires a plain dict where= argument + ) + return budget_id if row is not None else None + + +def _member_budget_source(budget_id: str | None, team_default_budget_id: str | None) -> TeamMemberBudgetSource: + if budget_id is not None and budget_id != team_default_budget_id: + return "custom" + return "team_default" if team_default_budget_id is not None else "none" + + +@router.post( + "/team/{team_id}/member/{user_id}/reset_budget", + tags=["team management"], # mutable-ok: FastAPI's `tags` param is typed as list[str], not Sequence + dependencies=(Depends(user_api_key_auth),), + response_model=TeamMemberResetBudgetResponse, +) +@management_endpoint_wrapper +async def reset_team_member_budget_fn( + team_id: str, + user_id: str, + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], +) -> TeamMemberResetBudgetResponse: + """ + Put a team member back on the team's shared default member budget (`team_member_budget`). + + Drops the member's own budget row link so team-wide changes made through /team/update + reach them again. Leaves the member with no budget when the team has no default. Spend is untouched. + """ + from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache + + if prisma_client is None: + _raise_reset_spend_error(status.HTTP_500_INTERNAL_SERVER_ERROR, "DB not connected. prisma_client is None") + + team_obj: Final = await get_team_object( + team_id=team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=proxy_logging_obj, + check_db_only=True, + ) + await _verify_team_access(team_obj=team_obj, user_api_key_dict=user_api_key_dict) + + membership_where: Final = { # mutable-ok: prisma client requires a plain dict where= argument + "user_id_team_id": {"user_id": user_id, "team_id": team_id} # mutable-ok: same prisma where= argument + } + membership_row: Final = await _team_membership_db(prisma_client).find_unique(where=membership_where) + if membership_row is None: + _raise_reset_spend_error(status.HTTP_404_NOT_FOUND, f"User {user_id} is not a member of team {team_id}.") + + team_default_budget_id: Final = await _existing_team_default_budget_id(team_obj, prisma_client) + budget_link: Final = ( + { + "connect": {"budget_id": team_default_budget_id} + } # mutable-ok: prisma client requires a plain dict data= argument + if team_default_budget_id is not None + else {"disconnect": True} # mutable-ok: same prisma data= argument + ) + await _team_membership_db(prisma_client).update( + where=membership_where, + data={"litellm_budget_table": budget_link}, # mutable-ok: prisma client requires a plain dict data= argument + ) + await invalidate_team_member_spend_state( + user_id=user_id, + team_id=team_id, + user_api_key_cache=user_api_key_cache, + ) + + return TeamMemberResetBudgetResponse( + team_id=team_id, + user_id=user_id, + budget_id=team_default_budget_id, + previous_budget_id=membership_row.budget_id, + budget_source=_member_budget_source(team_default_budget_id, team_default_budget_id), + ) + + def _create_results_from_response( members: list[Member], response: TeamAddMemberResponse, @@ -4722,9 +4818,7 @@ async def team_info( _team_info = TeamInfoResponseObjectTeamTable() ## GET TEAM BUDGET (if exists) ## - team_member_budget_id: Final = ( - _team_info.metadata.get("team_member_budget_id") if _team_info.metadata is not None else None - ) + team_member_budget_id: Final = _team_default_budget_id(_team_info) if team_member_budget_id is not None: _team_info = await _add_team_member_budget_table( team_member_budget_id=team_member_budget_id, @@ -4757,7 +4851,17 @@ async def team_info( team_id=team_id, team_info=hydrated_team_info, keys=keys, - team_memberships=returned_tm, + team_memberships=tuple( + TeamInfoMembership.model_validate( + MappingProxyType( + { + **tm.model_dump(), + "budget_source": _member_budget_source(tm.budget_id, team_member_budget_id), + } + ) + ) + for tm in returned_tm + ), ) return response_object diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index 690b5ae80b6..d55b9f79b5f 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -46,6 +46,7 @@ from litellm.proxy.management_endpoints.team_endpoints import ( _verify_team_access, delete_team, list_available_teams, + reset_team_member_budget_fn, reset_team_member_spend_fn, router, team_member_add_duplication_check, @@ -14432,6 +14433,219 @@ async def test_reset_team_member_spend_fn_proxy_admin_can_reset_own_spend(monkey assert response["spend"] == 0.0 +def _reset_budget_admin() -> UserAPIKeyAuth: + return UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user") + + +def _team_with_default_budget(team_id: str, budget_id: str) -> LiteLLM_TeamTable: + return LiteLLM_TeamTable(team_id=team_id, metadata={"team_member_budget_id": budget_id}) + + +@pytest.mark.asyncio +async def test_reset_team_member_budget_fn_relinks_custom_member_to_team_default(monkeypatch): + """An admin undoing a per-member budget must put the membership back on the team's shared + default row (a connect, not a copy) so later /team/update changes reach the member again, + and must drop the cached membership so the old cap stops being enforced. The shared row and + the member's tracked spend are never written.""" + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + mock_prisma_client = MagicMock() + real_cache = UserApiKeyCache() + await real_cache.async_set_cache(key="team-1_member-1", value="stale-membership") + await real_cache.async_set_cache(key="team_membership:member-1:team-1", value="stale-membership") + + membership_row = LiteLLM_TeamMembership(user_id="member-1", team_id="team-1", spend=10.0, budget_id="custom-b1") + mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=membership_row) + mock_prisma_client.db.litellm_teammembership.update = AsyncMock(return_value=membership_row) + mock_prisma_client.db.litellm_budgettable.find_unique = AsyncMock( + return_value=LiteLLM_BudgetTable(budget_id="team-default-b", max_budget=100.0) + ) + mock_prisma_client.db.litellm_budgettable.update = AsyncMock() + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", real_cache) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()) + + with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + "litellm.proxy.management_endpoints.team_endpoints.get_team_object", + AsyncMock(return_value=_team_with_default_budget("team-1", "team-default-b")), + ): + response = await reset_team_member_budget_fn( + team_id="team-1", user_id="member-1", user_api_key_dict=_reset_budget_admin() + ) + + assert response.budget_id == "team-default-b" + assert response.previous_budget_id == "custom-b1" + assert response.budget_source == "team_default" + mock_prisma_client.db.litellm_teammembership.update.assert_awaited_once_with( + where={"user_id_team_id": {"user_id": "member-1", "team_id": "team-1"}}, + data={"litellm_budget_table": {"connect": {"budget_id": "team-default-b"}}}, + ) + mock_prisma_client.db.litellm_budgettable.update.assert_not_awaited() + assert await real_cache.async_get_cache(key="team-1_member-1") is None + assert await real_cache.async_get_cache(key="team_membership:member-1:team-1") is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "team_obj, default_row", + [ + (LiteLLM_TeamTable(team_id="team-1"), None), + (_team_with_default_budget("team-1", "gone-b"), None), + ], + ids=["no_default_configured", "configured_default_row_missing"], +) +async def test_reset_team_member_budget_fn_detaches_member_when_team_has_no_usable_default( + monkeypatch, team_obj, default_row +): + """With no shared default to link to, reset leaves the member exactly where a freshly added + member would be: no budget row at all, reported as budget_source='none', rather than + connecting to a budget_id that does not exist or leaving the custom cap in place.""" + mock_prisma_client = MagicMock() + membership_row = LiteLLM_TeamMembership(user_id="member-1", team_id="team-1", budget_id="custom-b1") + mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=membership_row) + mock_prisma_client.db.litellm_teammembership.update = AsyncMock(return_value=membership_row) + mock_prisma_client.db.litellm_budgettable.find_unique = AsyncMock(return_value=default_row) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()) + + with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + "litellm.proxy.management_endpoints.team_endpoints.get_team_object", + AsyncMock(return_value=team_obj), + ): + response = await reset_team_member_budget_fn( + team_id="team-1", user_id="member-1", user_api_key_dict=_reset_budget_admin() + ) + + assert response.budget_id is None + assert response.previous_budget_id == "custom-b1" + assert response.budget_source == "none" + mock_prisma_client.db.litellm_teammembership.update.assert_awaited_once_with( + where={"user_id_team_id": {"user_id": "member-1", "team_id": "team-1"}}, + data={"litellm_budget_table": {"disconnect": True}}, + ) + + +@pytest.mark.asyncio +async def test_reset_team_member_budget_fn_membership_not_found(monkeypatch): + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_teammembership.update = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()) + + with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + "litellm.proxy.management_endpoints.team_endpoints.get_team_object", + AsyncMock(return_value=_team_with_default_budget("team-1", "team-default-b")), + ): + with pytest.raises(HTTPException) as exc: + await reset_team_member_budget_fn( + team_id="team-1", user_id="ghost-user", user_api_key_dict=_reset_budget_admin() + ) + assert exc.value.status_code == 404 + mock_prisma_client.db.litellm_teammembership.update.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_reset_team_member_budget_fn_forbidden_for_non_admin(monkeypatch): + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_teammembership.update = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()) + + with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + "litellm.proxy.management_endpoints.team_endpoints.get_team_object", + AsyncMock(return_value=LiteLLM_TeamTable(team_id="team-1", members_with_roles=[])), + ): + with pytest.raises(HTTPException) as exc: + await reset_team_member_budget_fn( + team_id="team-1", + user_id="member-1", + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, api_key="sk-user", user_id="plain-user" + ), + ) + assert exc.value.status_code == 403 + mock_prisma_client.db.litellm_teammembership.update.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_team_info_reports_whether_each_member_follows_the_team_default_budget(): + """/team/info must tell the caller which members still follow the team's shared member budget + and which carry their own row, since budget_id alone only means something to a reader who + also knows the team's team_member_budget_id.""" + from fastapi import Request + + from litellm.proxy.management_endpoints import team_endpoints + + team_row = _team_with_default_budget("team-1", "team-default-b") + memberships = [ + LiteLLM_TeamMembership(user_id="inherits", team_id="team-1", budget_id="team-default-b"), + LiteLLM_TeamMembership(user_id="customized", team_id="team-1", budget_id="own-b"), + LiteLLM_TeamMembership(user_id="unlinked", team_id="team-1", budget_id=None), + ] + + mock_prisma = MagicMock() + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row) + mock_prisma.db.litellm_budgettable.find_unique = AsyncMock( + return_value=LiteLLM_BudgetTable(budget_id="team-default-b", max_budget=100.0) + ) + mock_prisma.get_data = AsyncMock(return_value=[]) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch.object(team_endpoints, "get_all_team_memberships", AsyncMock(return_value=memberships)), + ): + response = await team_endpoints.team_info( + http_request=MagicMock(spec=Request), + team_id="team-1", + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + assert {tm.user_id: tm.budget_source for tm in response["team_memberships"]} == { + "inherits": "team_default", + "customized": "custom", + "unlinked": "team_default", + } + + +@pytest.mark.asyncio +async def test_team_info_reports_no_budget_source_when_team_has_no_default(): + """A team that never set team_member_budget has nothing for members to inherit, so an + unlinked member is 'none' rather than 'team_default', while a member with their own row is + still 'custom'.""" + from fastapi import Request + + from litellm.proxy.management_endpoints import team_endpoints + + memberships = [ + LiteLLM_TeamMembership(user_id="customized", team_id="team-1", budget_id="own-b"), + LiteLLM_TeamMembership(user_id="unlinked", team_id="team-1", budget_id=None), + ] + + mock_prisma = MagicMock() + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=LiteLLM_TeamTable(team_id="team-1")) + mock_prisma.get_data = AsyncMock(return_value=[]) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch.object(team_endpoints, "get_all_team_memberships", AsyncMock(return_value=memberships)), + ): + response = await team_endpoints.team_info( + http_request=MagicMock(spec=Request), + team_id="team-1", + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + assert {tm.user_id: tm.budget_source for tm in response["team_memberships"]} == { + "customized": "custom", + "unlinked": "none", + } + + @pytest.mark.asyncio async def test_team_member_update_invalidates_team_member_spend_state_when_budget_patch_applied(monkeypatch): """Raising a stuck member's max_budget_in_team via the documented /team/member_update diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useResetTeamMemberBudget.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useResetTeamMemberBudget.ts new file mode 100644 index 00000000000..e7cf95440a5 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useResetTeamMemberBudget.ts @@ -0,0 +1,16 @@ +import { useMutation } from "@tanstack/react-query"; +import { fetchClient } from "@/lib/http/api"; + +export interface ResetTeamMemberBudgetParams { + teamId: string; + userId: string; +} + +export const resetTeamMemberBudget = async ({ teamId, userId }: ResetTeamMemberBudgetParams): Promise => { + await fetchClient.POST("/team/{team_id}/member/{user_id}/reset_budget", { + params: { path: { team_id: teamId, user_id: userId } }, + }); +}; + +export const useResetTeamMemberBudget = () => + useMutation({ mutationFn: resetTeamMemberBudget }); diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index 3b2c344c2f3..22cc99b32c8 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -1,4 +1,5 @@ import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import type { components } from "@/lib/http/schema"; import useCan from "@/app/(dashboard)/hooks/useCan"; import { organizationKeys, useOrganizations } from "@/app/(dashboard)/hooks/organizations/useOrganizations"; import { useQueryClient } from "@tanstack/react-query"; @@ -247,10 +248,13 @@ export const retainedMcpToolPermissions = ( export const mcpUnresolvableSaveError = (reason: string): string => `Cannot save MCP tool permissions because ${reason}. Retry once the page has finished loading`; +export type TeamMemberBudgetSource = components["schemas"]["TeamMemberResetBudgetResponse"]["budget_source"]; + export interface TeamMembership { user_id: string; team_id: string; - budget_id: string; + budget_id: string | null; + budget_source: TeamMemberBudgetSource; spend: number; total_spend: number | null; litellm_budget_table: { @@ -1361,6 +1365,7 @@ const TeamInfoView: React.FC = ({ canEditTeam={canEditTeam} handleMemberDelete={handleMemberDelete} onMemberSpendReset={refreshTeamData} + onMemberBudgetReset={refreshTeamData} setSelectedEditMember={setSelectedEditMember} setIsEditMemberModalVisible={setIsEditMemberModalVisible} setIsAddMemberModalVisible={setIsAddMemberModalVisible} diff --git a/ui/litellm-dashboard/src/components/team/TeamMemberTab.test.tsx b/ui/litellm-dashboard/src/components/team/TeamMemberTab.test.tsx index 52cba1e6330..8652ffa7de2 100644 --- a/ui/litellm-dashboard/src/components/team/TeamMemberTab.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamMemberTab.test.tsx @@ -30,6 +30,7 @@ const mockSetSelectedEditMember = vi.fn(); const mockSetIsEditMemberModalVisible = vi.fn(); const mockSetIsAddMemberModalVisible = vi.fn(); const mockOnMemberSpendReset = vi.fn(); +const mockOnMemberBudgetReset = vi.fn(); const budgetResetIso = new Date(2026, 6, 15, 12, 0, 0).toISOString(); @@ -74,6 +75,7 @@ const createMockTeamData = (overrides: Partial = {}): TeamData => ({ user_id: "user1@test.com", team_id: "team-123", budget_id: "budget1", + budget_source: "custom", spend: 100.5, total_spend: 1538.2608, litellm_budget_table: { @@ -126,6 +128,7 @@ describe("TeamMembersComponent", () => { canEditTeam={false} handleMemberDelete={mockHandleMemberDelete} onMemberSpendReset={mockOnMemberSpendReset} + onMemberBudgetReset={mockOnMemberBudgetReset} setSelectedEditMember={mockSetSelectedEditMember} setIsEditMemberModalVisible={mockSetIsEditMemberModalVisible} setIsAddMemberModalVisible={mockSetIsAddMemberModalVisible} @@ -142,6 +145,7 @@ describe("TeamMembersComponent", () => { canEditTeam={false} handleMemberDelete={mockHandleMemberDelete} onMemberSpendReset={mockOnMemberSpendReset} + onMemberBudgetReset={mockOnMemberBudgetReset} setSelectedEditMember={mockSetSelectedEditMember} setIsEditMemberModalVisible={mockSetIsEditMemberModalVisible} setIsAddMemberModalVisible={mockSetIsAddMemberModalVisible} @@ -161,6 +165,7 @@ describe("TeamMembersComponent", () => { canEditTeam={false} handleMemberDelete={mockHandleMemberDelete} onMemberSpendReset={mockOnMemberSpendReset} + onMemberBudgetReset={mockOnMemberBudgetReset} setSelectedEditMember={mockSetSelectedEditMember} setIsEditMemberModalVisible={mockSetIsEditMemberModalVisible} setIsAddMemberModalVisible={mockSetIsAddMemberModalVisible} @@ -180,6 +185,7 @@ describe("TeamMembersComponent", () => { canEditTeam: false, handleMemberDelete: mockHandleMemberDelete, onMemberSpendReset: mockOnMemberSpendReset, + onMemberBudgetReset: mockOnMemberBudgetReset, setSelectedEditMember: mockSetSelectedEditMember, setIsEditMemberModalVisible: mockSetIsEditMemberModalVisible, setIsAddMemberModalVisible: mockSetIsAddMemberModalVisible, @@ -204,6 +210,7 @@ describe("TeamMembersComponent", () => { canEditTeam={true} handleMemberDelete={mockHandleMemberDelete} onMemberSpendReset={mockOnMemberSpendReset} + onMemberBudgetReset={mockOnMemberBudgetReset} setSelectedEditMember={mockSetSelectedEditMember} setIsEditMemberModalVisible={mockSetIsEditMemberModalVisible} setIsAddMemberModalVisible={mockSetIsAddMemberModalVisible} @@ -231,6 +238,7 @@ describe("TeamMembersComponent", () => { canEditTeam={false} handleMemberDelete={mockHandleMemberDelete} onMemberSpendReset={mockOnMemberSpendReset} + onMemberBudgetReset={mockOnMemberBudgetReset} setSelectedEditMember={mockSetSelectedEditMember} setIsEditMemberModalVisible={mockSetIsEditMemberModalVisible} setIsAddMemberModalVisible={mockSetIsAddMemberModalVisible} @@ -258,6 +266,7 @@ describe("TeamMembersComponent", () => { canEditTeam={false} handleMemberDelete={mockHandleMemberDelete} onMemberSpendReset={mockOnMemberSpendReset} + onMemberBudgetReset={mockOnMemberBudgetReset} setSelectedEditMember={mockSetSelectedEditMember} setIsEditMemberModalVisible={mockSetIsEditMemberModalVisible} setIsAddMemberModalVisible={mockSetIsAddMemberModalVisible} @@ -274,6 +283,7 @@ describe("TeamMembersComponent", () => { canEditTeam={false} handleMemberDelete={mockHandleMemberDelete} onMemberSpendReset={mockOnMemberSpendReset} + onMemberBudgetReset={mockOnMemberBudgetReset} setSelectedEditMember={mockSetSelectedEditMember} setIsEditMemberModalVisible={mockSetIsEditMemberModalVisible} setIsAddMemberModalVisible={mockSetIsAddMemberModalVisible} @@ -293,6 +303,7 @@ describe("TeamMembersComponent", () => { canEditTeam={false} handleMemberDelete={mockHandleMemberDelete} onMemberSpendReset={mockOnMemberSpendReset} + onMemberBudgetReset={mockOnMemberBudgetReset} setSelectedEditMember={mockSetSelectedEditMember} setIsEditMemberModalVisible={mockSetIsEditMemberModalVisible} setIsAddMemberModalVisible={mockSetIsAddMemberModalVisible} @@ -309,6 +320,7 @@ describe("TeamMembersComponent", () => { canEditTeam={false} handleMemberDelete={mockHandleMemberDelete} onMemberSpendReset={mockOnMemberSpendReset} + onMemberBudgetReset={mockOnMemberBudgetReset} setSelectedEditMember={mockSetSelectedEditMember} setIsEditMemberModalVisible={mockSetIsEditMemberModalVisible} setIsAddMemberModalVisible={mockSetIsAddMemberModalVisible} @@ -326,6 +338,7 @@ describe("TeamMembersComponent", () => { canEditTeam={false} handleMemberDelete={mockHandleMemberDelete} onMemberSpendReset={mockOnMemberSpendReset} + onMemberBudgetReset={mockOnMemberBudgetReset} setSelectedEditMember={mockSetSelectedEditMember} setIsEditMemberModalVisible={mockSetIsEditMemberModalVisible} setIsAddMemberModalVisible={mockSetIsAddMemberModalVisible} @@ -346,6 +359,7 @@ describe("TeamMembersComponent", () => { canEditTeam={true} handleMemberDelete={mockHandleMemberDelete} onMemberSpendReset={mockOnMemberSpendReset} + onMemberBudgetReset={mockOnMemberBudgetReset} setSelectedEditMember={mockSetSelectedEditMember} setIsEditMemberModalVisible={mockSetIsEditMemberModalVisible} setIsAddMemberModalVisible={mockSetIsAddMemberModalVisible} @@ -381,6 +395,7 @@ describe("TeamMembersComponent", () => { canEditTeam={true} handleMemberDelete={mockHandleMemberDelete} onMemberSpendReset={mockOnMemberSpendReset} + onMemberBudgetReset={mockOnMemberBudgetReset} setSelectedEditMember={mockSetSelectedEditMember} setIsEditMemberModalVisible={mockSetIsEditMemberModalVisible} setIsAddMemberModalVisible={mockSetIsAddMemberModalVisible} @@ -435,6 +450,7 @@ describe("TeamMembersComponent", () => { canEditTeam={true} handleMemberDelete={mockHandleMemberDelete} onMemberSpendReset={mockOnMemberSpendReset} + onMemberBudgetReset={mockOnMemberBudgetReset} setSelectedEditMember={mockSetSelectedEditMember} setIsEditMemberModalVisible={mockSetIsEditMemberModalVisible} setIsAddMemberModalVisible={mockSetIsAddMemberModalVisible} @@ -466,6 +482,7 @@ describe("TeamMembersComponent", () => { canEditTeam={true} handleMemberDelete={mockHandleMemberDelete} onMemberSpendReset={mockOnMemberSpendReset} + onMemberBudgetReset={mockOnMemberBudgetReset} setSelectedEditMember={mockSetSelectedEditMember} setIsEditMemberModalVisible={mockSetIsEditMemberModalVisible} setIsAddMemberModalVisible={mockSetIsAddMemberModalVisible} @@ -486,6 +503,7 @@ describe("TeamMembersComponent", () => { canEditTeam={true} handleMemberDelete={mockHandleMemberDelete} onMemberSpendReset={mockOnMemberSpendReset} + onMemberBudgetReset={mockOnMemberBudgetReset} setSelectedEditMember={mockSetSelectedEditMember} setIsEditMemberModalVisible={mockSetIsEditMemberModalVisible} setIsAddMemberModalVisible={mockSetIsAddMemberModalVisible} @@ -503,6 +521,7 @@ describe("TeamMembersComponent", () => { canEditTeam={false} handleMemberDelete={mockHandleMemberDelete} onMemberSpendReset={mockOnMemberSpendReset} + onMemberBudgetReset={mockOnMemberBudgetReset} setSelectedEditMember={mockSetSelectedEditMember} setIsEditMemberModalVisible={mockSetIsEditMemberModalVisible} setIsAddMemberModalVisible={mockSetIsAddMemberModalVisible} @@ -521,6 +540,7 @@ describe("TeamMembersComponent", () => { canEditTeam={true} handleMemberDelete={mockHandleMemberDelete} onMemberSpendReset={mockOnMemberSpendReset} + onMemberBudgetReset={mockOnMemberBudgetReset} setSelectedEditMember={mockSetSelectedEditMember} setIsEditMemberModalVisible={mockSetIsEditMemberModalVisible} setIsAddMemberModalVisible={mockSetIsAddMemberModalVisible} @@ -603,4 +623,134 @@ describe("TeamMembersComponent", () => { expect(screen.getByTestId("reset-member-spend")).toBeVisible(); }); }); + + describe("budget source", () => { + const teamDataWithDefault = () => { + const base = createMockTeamData(); + return createMockTeamData({ + team_info: { + ...base.team_info, + team_member_budget_table: { max_budget: 25, budget_duration: null, tpm_limit: null, rpm_limit: null }, + }, + team_memberships: [ + base.team_memberships[0], + { + user_id: "user2@test.com", + team_id: "team-123", + budget_id: "team-default-budget", + budget_source: "team_default", + spend: 0, + total_spend: null, + litellm_budget_table: { + budget_id: "team-default-budget", + soft_budget: null, + max_budget: 25, + max_parallel_requests: null, + tpm_limit: null, + rpm_limit: null, + model_max_budget: null, + budget_duration: null, + budget_reset_at: null, + }, + }, + ], + }); + }; + + const renderTab = (teamData: TeamData, canEditTeam = true) => + renderWithProviders( + , + ); + + it("labels each member's budget as Custom or Team default and shows the team amount for inherited members", () => { + renderTab(teamDataWithDefault()); + + const customRow = screen.getByRole("row", { name: /user1@test\.com/ }); + const inheritedRow = screen.getByRole("row", { name: /user2@test\.com/ }); + expect(within(customRow).getByTestId("member-budget-source")).toHaveTextContent("Custom"); + expect(customRow).toHaveTextContent("$1,000.00"); + expect(within(inheritedRow).getByTestId("member-budget-source")).toHaveTextContent("Team default"); + expect(inheritedRow).toHaveTextContent("$25.00"); + }); + + it("shows no source label for a member with neither a custom nor a team budget", () => { + renderTab(createMockTeamData({ team_memberships: [] })); + + expect(screen.queryByTestId("member-budget-source")).not.toBeInTheDocument(); + expect(screen.queryByTestId("reset-member-budget")).not.toBeInTheDocument(); + }); + + it("only offers Use team default on customized members, and only to editors", () => { + const { unmount } = renderTab(teamDataWithDefault()); + + expect( + within(screen.getByRole("row", { name: /user1@test\.com/ })).getByTestId("reset-member-budget"), + ).toBeVisible(); + expect( + within(screen.getByRole("row", { name: /user2@test\.com/ })).queryByTestId("reset-member-budget"), + ).not.toBeInTheDocument(); + + unmount(); + renderTab(teamDataWithDefault(), false); + expect(screen.queryByTestId("reset-member-budget")).not.toBeInTheDocument(); + }); + + it("puts the member back on the team default after confirming, then refreshes the team", async () => { + const user = userEvent.setup(); + POST.mockResolvedValue({ data: {} }); + renderTab(teamDataWithDefault()); + + await user.click(screen.getByTestId("reset-member-budget")); + + const dialog = await screen.findByRole("dialog", { name: "Reset Team Member Budget" }); + expect(dialog).toHaveTextContent("user1@test.com"); + expect(dialog).toHaveTextContent("team default of $25.00"); + expect(dialog).toHaveTextContent("Custom budget: $1,000.00"); + expect(POST).not.toHaveBeenCalled(); + + await user.click(within(dialog).getByRole("button", { name: "Use team default" })); + + await waitFor(() => expect(mockOnMemberBudgetReset).toHaveBeenCalledTimes(1)); + expect(POST).toHaveBeenCalledExactlyOnceWith("/team/{team_id}/member/{user_id}/reset_budget", { + params: { path: { team_id: "team-123", user_id: "user1@test.com" } }, + }); + expect(mockOnMemberSpendReset).not.toHaveBeenCalled(); + expect(screen.queryByRole("dialog")).not.toBeInTheDocument(); + }); + + it("keeps the dialog open and does not refresh the team when the reset fails", async () => { + const user = userEvent.setup(); + POST.mockRejectedValue(new Error("Team admin cannot reset budgets")); + renderTab(teamDataWithDefault()); + + await user.click(screen.getByTestId("reset-member-budget")); + const dialog = await screen.findByRole("dialog", { name: "Reset Team Member Budget" }); + await user.click(within(dialog).getByRole("button", { name: "Use team default" })); + + await waitFor(() => expect(POST).toHaveBeenCalledTimes(1)); + expect(mockOnMemberBudgetReset).not.toHaveBeenCalled(); + expect(screen.getByRole("dialog", { name: "Reset Team Member Budget" })).toBeInTheDocument(); + }); + + it("does not call the API when the dialog is cancelled", async () => { + const user = userEvent.setup(); + renderTab(teamDataWithDefault()); + + await user.click(screen.getByTestId("reset-member-budget")); + const dialog = await screen.findByRole("dialog", { name: "Reset Team Member Budget" }); + await user.click(within(dialog).getByRole("button", { name: "Cancel" })); + + await waitFor(() => expect(screen.queryByRole("dialog")).not.toBeInTheDocument()); + expect(POST).not.toHaveBeenCalled(); + }); + }); }); diff --git a/ui/litellm-dashboard/src/components/team/TeamMemberTab.tsx b/ui/litellm-dashboard/src/components/team/TeamMemberTab.tsx index a869c1ad624..660416504fe 100644 --- a/ui/litellm-dashboard/src/components/team/TeamMemberTab.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamMemberTab.tsx @@ -1,6 +1,8 @@ +import { useResetTeamMemberBudget } from "@/app/(dashboard)/hooks/teams/useResetTeamMemberBudget"; import { useResetTeamMemberSpend } from "@/app/(dashboard)/hooks/teams/useResetTeamMemberSpend"; import { useUISettings } from "@/app/(dashboard)/hooks/uiSettings/useUISettings"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import { Badge } from "@/components/ui/badge"; import { Button } from "@/components/ui/button"; import { Dialog, DialogContent, DialogFooter, DialogHeader, DialogTitle } from "@/components/ui/dialog"; import { SimpleTooltip } from "@/components/ui/tooltip"; @@ -13,7 +15,15 @@ import { formatNumberWithCommas } from "@/utils/dataUtils"; import { isProxyAdminRole, isUserTeamAdminForSingleTeam } from "@/utils/roles"; import { CircleHelp } from "lucide-react"; import { useState, type ComponentProps } from "react"; -import { TeamData, TeamMembership } from "./TeamInfo"; +import { TeamData, TeamMemberBudgetSource, TeamMembership } from "./TeamInfo"; + +const BUDGET_SOURCE_LABELS: Record, string> = { + team_default: "Team default", + custom: "Custom", +}; + +const formatBudget = (value: number | null): string => + value === null ? "Unlimited" : `$${formatNumberWithCommas(value, 2)}`; export const seedMemberBudgetFields = ( record: Member, @@ -37,6 +47,7 @@ interface TeamMemberTabProps { setIsEditMemberModalVisible: (visible: boolean) => void; setIsAddMemberModalVisible: (visible: boolean) => void; onMemberSpendReset: () => void; + onMemberBudgetReset: () => void; } export default function TeamMemberTab({ @@ -47,9 +58,13 @@ export default function TeamMemberTab({ setIsEditMemberModalVisible, setIsAddMemberModalVisible, onMemberSpendReset, + onMemberBudgetReset, }: TeamMemberTabProps) { const [memberToResetSpend, setMemberToResetSpend] = useState(null); + const [memberToResetBudget, setMemberToResetBudget] = useState(null); const { mutate: resetMemberSpend, isPending: isResettingSpend } = useResetTeamMemberSpend(); + const { mutate: resetMemberBudget, isPending: isResettingBudget } = useResetTeamMemberBudget(); + const teamDefaultBudget = teamData.team_info.team_member_budget_table?.max_budget ?? null; const formatNumber = (value: number | null): string => { if (value === null || value === undefined) return "0"; @@ -82,10 +97,19 @@ export default function TeamMemberTab({ return membership?.total_spend ?? 0; }; + const getUserBudgetSource = (userId: string | null): TeamMemberBudgetSource => { + if (!userId) return "none"; + const membership = teamData.team_memberships.find((tm) => tm.user_id === userId); + return membership?.budget_source ?? "none"; + }; + const getUserBudget = (userId: string | null): number | null => { if (!userId) return null; const membership = teamData.team_memberships.find((tm) => tm.user_id === userId); - return membership?.litellm_budget_table?.max_budget ?? null; + return ( + membership?.litellm_budget_table?.max_budget ?? + (membership?.budget_source === "team_default" ? teamDefaultBudget : null) + ); }; // Helper function to get rate limits for a user @@ -182,12 +206,40 @@ export default function TeamMemberTab({ render: (record: Member) => , }, { - title: "Team Member Budget (USD)", + title: ( + + Team Member Budget (USD) + + + + + ), key: "budget", sortValue: (record: Member) => getUserBudget(record.user_id), - render: (record: Member) => ( - - ), + render: (record: Member) => { + const source = getUserBudgetSource(record.user_id); + return ( + + + {source !== "none" && ( + + {BUDGET_SOURCE_LABELS[source]} + + )} + {source === "custom" && canEditTeam && ( + + )} + + ); + }, }, { title: "Budget Reset", @@ -224,6 +276,21 @@ export default function TeamMemberTab({ ); }; + const handleResetBudget = () => { + if (!memberToResetBudget?.user_id) return; + resetMemberBudget( + { teamId: teamData.team_id, userId: memberToResetBudget.user_id }, + { + onSuccess: () => { + toast.success("Team member budget reset to the team default"); + setMemberToResetBudget(null); + onMemberBudgetReset(); + }, + onError: (error) => toast.fromError(parseErrorMessage(error)), + }, + ); + }; + return ( <> + !open && setMemberToResetBudget(null)}> + + + Reset Team Member Budget + +

+ Remove the custom budget for{" "} + {memberToResetBudget?.user_email || memberToResetBudget?.user_id} and put them back on the + team default of {formatBudget(teamDefaultBudget)}? +

+

+ Custom budget: {formatBudget(getUserBudget(memberToResetBudget?.user_id ?? null))}. Their + spend is kept. Future changes to the team's member budget will apply to them again. +

+ + + + + + ); } diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 136b8f26784..aa6776d12bd 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -16272,6 +16272,29 @@ export interface paths { patch?: never; trace?: never; }; + "/team/{team_id}/member/{user_id}/reset_budget": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + /** + * Reset Team Member Budget Fn + * @description Put a team member back on the team's shared default member budget (`team_member_budget`). + * + * Drops the member's own budget row link so team-wide changes made through /team/update + * reach them again. Leaves the member with no budget when the team has no default. Spend is untouched. + */ + post: operations["reset_team_member_budget_fn_team__team_id__member__user_id__reset_budget_post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/team/{team_id}/member/{user_id}/reset_spend": { parameters: { query?: never; @@ -38547,6 +38570,22 @@ export interface components { /** User Id */ user_id?: string | null; }; + /** TeamMemberResetBudgetResponse */ + TeamMemberResetBudgetResponse: { + /** Budget Id */ + budget_id: string | null; + /** + * Budget Source + * @enum {string} + */ + budget_source: "team_default" | "custom" | "none"; + /** Previous Budget Id */ + previous_budget_id: string | null; + /** Team Id */ + team_id: string; + /** User Id */ + user_id: string; + }; /** TeamMemberUpdateRequest */ TeamMemberUpdateRequest: { /** @@ -61730,6 +61769,38 @@ export interface operations { }; }; }; + reset_team_member_budget_fn_team__team_id__member__user_id__reset_budget_post: { + parameters: { + query?: never; + header?: never; + path: { + team_id: string; + user_id: string; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["TeamMemberResetBudgetResponse"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; reset_team_member_spend_fn_team__team_id__member__user_id__reset_spend_post: { parameters: { query?: never;