From b32d1112a6c9af25f0b5a66eae7ba33d8e201c74 Mon Sep 17 00:00:00 2001 From: ryan Date: Sat, 19 Sep 2026 00:31:55 +0000 Subject: [PATCH 1/4] feat(team): show whether a member follows the team default budget and allow resetting to it Adds budget_source (team_default, custom, none) to each membership in /team/info and a POST /team/{team_id}/member/{user_id}/reset_budget route that relinks a member to the team's shared team_member_budget row without touching their spend. The Admin UI team members table shows a Team default or Custom badge next to each member's budget and offers a Use team default action on customized members Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/_types.py | 19 +- .../management_endpoints/team_endpoints.py | 112 ++++++++- .../test_team_endpoints.py | 214 ++++++++++++++++++ .../hooks/teams/useResetTeamMemberBudget.ts | 16 ++ .../src/components/team/TeamInfo.tsx | 7 +- .../components/team/TeamMemberTab.test.tsx | 150 ++++++++++++ .../src/components/team/TeamMemberTab.tsx | 103 ++++++++- ui/litellm-dashboard/src/lib/http/schema.d.ts | 71 ++++++ 8 files changed, 679 insertions(+), 13 deletions(-) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useResetTeamMemberBudget.ts 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; From d3a364d74f8bee7f6133d51e68c3036cde7f8136 Mon Sep 17 00:00:00 2001 From: ryan Date: Sat, 19 Sep 2026 00:49:58 +0000 Subject: [PATCH 2/4] fix(team): report no budget source when the team default row was deleted Derive budget_source from the budget row /team/info actually loaded, so a metadata id whose row was removed via /budget/delete reads as none instead of team_default. Share the /team/info test scaffolding so the added patch calls stay within the TQ008 budget, and allowlist the imperative reset_budget route in the provider endpoint audit Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../management_endpoints/team_endpoints.py | 5 +- .../endpointaudit/coverage_allowlist.txt | 1 + .../test_team_endpoints.py | 112 ++++++++++-------- 3 files changed, 70 insertions(+), 48 deletions(-) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index b8d95167045..a441b3834ed 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -4825,6 +4825,9 @@ async def team_info( prisma_client=prisma_client, team_info_response_object=_team_info, ) + active_default_budget_id: Final = ( + team_member_budget_id if _team_info.team_member_budget_table is not None else None + ) # Resolve resources inherited from access groups resolved_team_info: Final = await _resolve_team_access_group_resources(_team_info) @@ -4856,7 +4859,7 @@ async def team_info( MappingProxyType( { **tm.model_dump(), - "budget_source": _member_budget_source(tm.budget_id, team_member_budget_id), + "budget_source": _member_budget_source(tm.budget_id, active_default_budget_id), } ) ) diff --git a/terraform/provider/tools/endpointaudit/coverage_allowlist.txt b/terraform/provider/tools/endpointaudit/coverage_allowlist.txt index 6bc8947e89f..4ea64b152f1 100644 --- a/terraform/provider/tools/endpointaudit/coverage_allowlist.txt +++ b/terraform/provider/tools/endpointaudit/coverage_allowlist.txt @@ -81,6 +81,7 @@ POST /prompts/test POST /search_tools/test_connection POST /team/bulk_member_add POST /team/{team_id}/member/{user_id}/reset_spend +POST /team/{team_id}/member/{user_id}/reset_budget POST /team/key/bulk_update POST /team/permissions_bulk_update POST /team/{team_id}/disable_logging 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 d55b9f79b5f..484c054fa54 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -14572,40 +14572,52 @@ async def test_reset_team_member_budget_fn_forbidden_for_non_admin(monkeypatch): mock_prisma_client.db.litellm_teammembership.update.assert_not_awaited() +async def _team_info_budget_sources( + team_row: LiteLLM_TeamTable, + memberships: list[LiteLLM_TeamMembership], + default_budget_row: LiteLLM_BudgetTable | None, +) -> dict[str, str]: + from fastapi import Request + + from litellm.proxy.management_endpoints import team_endpoints + + 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=default_budget_row) + mock_prisma.get_data = AsyncMock(return_value=[]) + + with ( + patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + "litellm.proxy.proxy_server.prisma_client", mock_prisma + ), + patch.object( # test-quality-ok: membership lookup is a module-level DB query with no injection point + team_endpoints, "get_all_team_memberships", AsyncMock(return_value=memberships) + ), + ): + response = await team_endpoints.team_info( + http_request=MagicMock(spec=Request), + team_id=team_row.team_id, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) + return {tm.user_id: tm.budget_source for tm in response["team_memberships"]} + + @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) + sources = await _team_info_budget_sources( + 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), + ], + default_budget_row=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"]} == { + assert sources == { "inherits": "team_default", "customized": "custom", "unlinked": "team_default", @@ -14617,30 +14629,36 @@ 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 + sources = await _team_info_budget_sources( + team_row=LiteLLM_TeamTable(team_id="team-1"), + 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), + ], + default_budget_row=None, + ) - from litellm.proxy.management_endpoints import team_endpoints + assert sources == { + "customized": "custom", + "unlinked": "none", + } - 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=[]) +@pytest.mark.asyncio +async def test_team_info_reports_no_budget_source_when_team_default_row_was_deleted(): + """If the budget row named by team_member_budget_id was removed via /budget/delete, nothing is + enforced for unlinked members any more, so /team/info must not keep advertising a team default + that no longer exists.""" + sources = await _team_info_budget_sources( + team_row=_team_with_default_budget("team-1", "deleted-b"), + 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), + ], + default_budget_row=None, + ) - 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"]} == { + assert sources == { "customized": "custom", "unlinked": "none", } From 52a71ff68188b0d6781204144472c613c1c8240b Mon Sep 17 00:00:00 2001 From: ryan Date: Sat, 19 Sep 2026 00:58:39 +0000 Subject: [PATCH 3/4] fix(team): let team admins reach the member reset_budget route and cover it in the behavior suite Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/_types.py | 1 + .../test_team_member_reset_budget.py | 201 ++++++++++++++++++ 2 files changed, 202 insertions(+) create mode 100644 tests/proxy_behavior/management/test_team_member_reset_budget.py diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index a624234cf5f..c21113f8c29 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -865,6 +865,7 @@ class LiteLLMRoutes(enum.Enum): "/management/v1/teams/{team_id}/members/bulk_update", "/team/member_update", "/team/{team_id}/member/{user_id}/reset_spend", + "/team/{team_id}/member/{user_id}/reset_budget", "/team/permissions_list", "/team/permissions_update", "/team/daily/activity", diff --git a/tests/proxy_behavior/management/test_team_member_reset_budget.py b/tests/proxy_behavior/management/test_team_member_reset_budget.py new file mode 100644 index 00000000000..42f327c33ef --- /dev/null +++ b/tests/proxy_behavior/management/test_team_member_reset_budget.py @@ -0,0 +1,201 @@ +import uuid + +import pytest + +from .actors import Actor +from .conftest import create_scratch_team + +pytestmark = pytest.mark.asyncio(loop_scope="session") + +_SEED_SPEND = 5.0 +_TEAM_DEFAULT_MAX_BUDGET = 100.0 +_CUSTOM_MAX_BUDGET = 50.0 + +_MATRIX = [ + ("alpha/proxy_admin", Actor.PROXY_ADMIN, "alpha", 200), + ("alpha/org_admin", Actor.ORG_ADMIN, "alpha", 200), + ("alpha/team_admin", Actor.TEAM_ADMIN, "alpha", 200), + ("alpha/internal_user", Actor.INTERNAL_USER, "alpha", 403), + ("alpha/owner", Actor.OWNER, "alpha", 403), + ("alpha/unrelated_same_org", Actor.UNRELATED_SAME_ORG, "alpha", 403), + ("alpha/cross_org_user", Actor.CROSS_ORG_USER, "alpha", 403), + ("alpha/service_account", Actor.SERVICE_ACCOUNT, "alpha", 403), + ("alpha/org_b_admin", Actor.ORG_B_ADMIN, "alpha", 403), + ("beta/proxy_admin", Actor.PROXY_ADMIN, "beta", 200), + ("beta/org_admin", Actor.ORG_ADMIN, "beta", 403), + ("beta/team_admin", Actor.TEAM_ADMIN, "beta", 403), + ("beta/org_b_admin", Actor.ORG_B_ADMIN, "beta", 200), +] + + +async def _seed_budget(prisma, budget_id: str, max_budget: float) -> str: + await prisma.db.litellm_budgettable.create( + data={ + "budget_id": budget_id, + "max_budget": max_budget, + "created_by": "phase4-scratch", + "updated_by": "phase4-scratch", + } + ) + return budget_id + + +async def _seed_team_with_default_budget(prisma, world, shape: str, team_id: str, scratch) -> str: + default_budget_id = await _seed_budget(prisma, scratch.tag("team-default-budget"), _TEAM_DEFAULT_MAX_BUDGET) + metadata = {"team_member_budget_id": default_budget_id} + if shape == "alpha": + await create_scratch_team( + prisma, + team_id, + organization_id=world.org_a_id, + admin_user_ids=[world.keys[Actor.TEAM_ADMIN].user_id], + metadata=metadata, + ) + elif shape == "beta": + await create_scratch_team(prisma, team_id, organization_id=world.org_b_id, metadata=metadata) + else: # pragma: no cover - guard + pytest.fail(f"unknown shape={shape}") + return default_budget_id + + +async def _seed_custom_member(prisma, team_id: str, member_id: str, scratch) -> str: + custom_budget_id = await _seed_budget(prisma, scratch.tag("custom-budget"), _CUSTOM_MAX_BUDGET) + await prisma.db.litellm_teammembership.create( + data={ + "user_id": member_id, + "team_id": team_id, + "spend": _SEED_SPEND, + "litellm_budget_table": {"connect": {"budget_id": custom_budget_id}}, + } + ) + return custom_budget_id + + +async def _membership(prisma, team_id: str, member_id: str): + row = await prisma.db.litellm_teammembership.find_unique( + where={"user_id_team_id": {"user_id": member_id, "team_id": team_id}} + ) + assert row is not None + return row + + +@pytest.mark.parametrize( + "actor,shape,expected_status", + [(a, sh, s) for (_id, a, sh, s) in _MATRIX], + ids=[s[0] for s in _MATRIX], +) +async def test_team_member_reset_budget_authz_matrix( + actor: Actor, + shape: str, + expected_status: int, + proxy_client, + prisma, + scratch, + world, +): + member_id = scratch.tag("member") + default_budget_id = await _seed_team_with_default_budget(prisma, world, shape, scratch.prefix, scratch) + custom_budget_id = await _seed_custom_member(prisma, scratch.prefix, member_id, scratch) + caller = world.keys[actor] + + resp = await proxy_client.post( + f"/team/{scratch.prefix}/member/{member_id}/reset_budget", + headers={"Authorization": f"Bearer {caller.cleartext}"}, + ) + assert resp.status_code == expected_status, f"{actor.value} {shape}: {resp.status_code} {resp.text}" + + row = await _membership(prisma, scratch.prefix, member_id) + assert row.spend == _SEED_SPEND, "reset_budget must never touch spend" + if expected_status == 200: + assert row.budget_id == default_budget_id + body = resp.json() + assert body["budget_id"] == default_budget_id + assert body["previous_budget_id"] == custom_budget_id + assert body["budget_source"] == "team_default" + else: + assert row.budget_id == custom_budget_id, "denied but budget relinked" + + +async def test_team_member_reset_budget_leaves_shared_default_row_untouched(proxy_client, prisma, scratch, world): + """Relinking must point the member at the shared row, not copy or edit it, so a later + /team/update to team_member_budget reaches this member again.""" + member_id = scratch.tag("member") + default_budget_id = await _seed_team_with_default_budget(prisma, world, "alpha", scratch.prefix, scratch) + await _seed_custom_member(prisma, scratch.prefix, member_id, scratch) + + resp = await proxy_client.post( + f"/team/{scratch.prefix}/member/{member_id}/reset_budget", + headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"}, + ) + assert resp.status_code == 200, resp.text + + default_row = await prisma.db.litellm_budgettable.find_unique(where={"budget_id": default_budget_id}) + assert default_row is not None and default_row.max_budget == _TEAM_DEFAULT_MAX_BUDGET + + info = await proxy_client.get( + f"/team/info?team_id={scratch.prefix}", + headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"}, + ) + assert info.status_code == 200, info.text + memberships = {tm["user_id"]: tm for tm in info.json()["team_memberships"]} + assert memberships[member_id]["budget_source"] == "team_default" + assert memberships[member_id]["litellm_budget_table"]["max_budget"] == _TEAM_DEFAULT_MAX_BUDGET + + +async def test_team_member_reset_budget_without_team_default_detaches_member(proxy_client, prisma, scratch, world): + member_id = scratch.tag("member") + await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id) + await _seed_custom_member(prisma, scratch.prefix, member_id, scratch) + + resp = await proxy_client.post( + f"/team/{scratch.prefix}/member/{member_id}/reset_budget", + headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"}, + ) + assert resp.status_code == 200, resp.text + assert resp.json()["budget_id"] is None + assert resp.json()["budget_source"] == "none" + + row = await _membership(prisma, scratch.prefix, member_id) + assert row.budget_id is None + assert row.spend == _SEED_SPEND + + +async def test_team_member_reset_budget_with_deleted_team_default_detaches_member(proxy_client, prisma, scratch, world): + """metadata.team_member_budget_id can outlive its budget row; a stale id must not be + relinked to (the FK would fail) and must read as no budget, not as the team default.""" + member_id = scratch.tag("member") + await create_scratch_team( + prisma, + scratch.prefix, + organization_id=world.org_a_id, + metadata={"team_member_budget_id": scratch.tag("deleted-budget")}, + ) + await _seed_custom_member(prisma, scratch.prefix, member_id, scratch) + + resp = await proxy_client.post( + f"/team/{scratch.prefix}/member/{member_id}/reset_budget", + headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"}, + ) + assert resp.status_code == 200, resp.text + assert resp.json()["budget_id"] is None + assert resp.json()["budget_source"] == "none" + + row = await _membership(prisma, scratch.prefix, member_id) + assert row.budget_id is None + + +async def test_team_member_reset_budget_missing_team_is_404(proxy_client, world): + resp = await proxy_client.post( + f"/team/behavior-pin-no-such-team/member/{uuid.uuid4().hex}/reset_budget", + headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"}, + ) + assert resp.status_code == 404, resp.text + + +async def test_team_member_reset_budget_missing_membership_is_404(proxy_client, prisma, scratch, world): + await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id) + resp = await proxy_client.post( + f"/team/{scratch.prefix}/member/{uuid.uuid4().hex}/reset_budget", + headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"}, + ) + assert resp.status_code == 404, resp.text From 8bd9d356dcc10632a8efdc1a2229646e97cb301f Mon Sep 17 00:00:00 2001 From: ryan Date: Sat, 19 Sep 2026 01:00:31 +0000 Subject: [PATCH 4/4] test(team): drop docstrings that restate the budget source and reset assertions Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../management/test_team_member_reset_budget.py | 4 ---- .../management_endpoints/test_team_endpoints.py | 16 ---------------- 2 files changed, 20 deletions(-) diff --git a/tests/proxy_behavior/management/test_team_member_reset_budget.py b/tests/proxy_behavior/management/test_team_member_reset_budget.py index 42f327c33ef..1e55b8b6b15 100644 --- a/tests/proxy_behavior/management/test_team_member_reset_budget.py +++ b/tests/proxy_behavior/management/test_team_member_reset_budget.py @@ -117,8 +117,6 @@ async def test_team_member_reset_budget_authz_matrix( async def test_team_member_reset_budget_leaves_shared_default_row_untouched(proxy_client, prisma, scratch, world): - """Relinking must point the member at the shared row, not copy or edit it, so a later - /team/update to team_member_budget reaches this member again.""" member_id = scratch.tag("member") default_budget_id = await _seed_team_with_default_budget(prisma, world, "alpha", scratch.prefix, scratch) await _seed_custom_member(prisma, scratch.prefix, member_id, scratch) @@ -161,8 +159,6 @@ async def test_team_member_reset_budget_without_team_default_detaches_member(pro async def test_team_member_reset_budget_with_deleted_team_default_detaches_member(proxy_client, prisma, scratch, world): - """metadata.team_member_budget_id can outlive its budget row; a stale id must not be - relinked to (the FK would fail) and must read as no budget, not as the team default.""" member_id = scratch.tag("member") await create_scratch_team( prisma, 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 484c054fa54..1f19163933a 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -14443,10 +14443,6 @@ def _team_with_default_budget(team_id: str, budget_id: str) -> LiteLLM_TeamTable @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() @@ -14498,9 +14494,6 @@ async def test_reset_team_member_budget_fn_relinks_custom_member_to_team_default 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) @@ -14604,9 +14597,6 @@ async def _team_info_budget_sources( @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.""" sources = await _team_info_budget_sources( team_row=_team_with_default_budget("team-1", "team-default-b"), memberships=[ @@ -14626,9 +14616,6 @@ async def test_team_info_reports_whether_each_member_follows_the_team_default_bu @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'.""" sources = await _team_info_budget_sources( team_row=LiteLLM_TeamTable(team_id="team-1"), memberships=[ @@ -14646,9 +14633,6 @@ async def test_team_info_reports_no_budget_source_when_team_has_no_default(): @pytest.mark.asyncio async def test_team_info_reports_no_budget_source_when_team_default_row_was_deleted(): - """If the budget row named by team_member_budget_id was removed via /budget/delete, nothing is - enforced for unlinked members any more, so /team/info must not keep advertising a team default - that no longer exists.""" sources = await _team_info_budget_sources( team_row=_team_with_default_budget("team-1", "deleted-b"), memberships=[