diff --git a/ui/litellm-dashboard/src/components/team/ResetMemberBudgetsDialog.test.tsx b/ui/litellm-dashboard/src/components/team/ResetMemberBudgetsDialog.test.tsx new file mode 100644 index 00000000000..7132ec13fbb --- /dev/null +++ b/ui/litellm-dashboard/src/components/team/ResetMemberBudgetsDialog.test.tsx @@ -0,0 +1,103 @@ +import { render, screen } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { describe, expect, it, vi } from "vitest"; +import ResetMemberBudgetsDialog from "./ResetMemberBudgetsDialog"; +import type { MemberBudgetResetState } from "./useMemberBudgetReset"; + +const pendingFor = (memberCount: number, newBudget = 10) => ({ + teamId: "team-123", + updateData: {}, + userIds: Array.from({ length: memberCount }, (_, i) => `u-${i}`), + newBudget, +}); + +const promptingWith = (memberCount: number, newBudget = 10): MemberBudgetResetState => ({ + phase: "prompting", + pending: pendingFor(memberCount, newBudget), +}); + +const defaultHandlers = () => ({ + onReset: vi.fn(), + onRetry: vi.fn(), + onKeep: vi.fn(), + onDismiss: vi.fn(), +}); + +describe("ResetMemberBudgetsDialog", () => { + it("names the members and the new default so the prompt is self-explanatory", () => { + render(); + + expect( + screen.getByText( + "3 members have a custom budget, so the new team default of $10.00 will not apply to them. " + + "Reset them to the default, or keep the custom budgets?", + ), + ).toBeInTheDocument(); + }); + + it("switches to singular for one member", () => { + render(); + + expect( + screen.getByText( + "1 member has a custom budget, so the new team default of $10.00 will not apply to that member. " + + "Reset it to the default, or keep the custom budget?", + ), + ).toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Reset to $10.00" })).toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Keep custom budget" })).toBeInTheDocument(); + }); + + it("formats large budgets like the rest of the team page", () => { + render(); + + expect(screen.getByRole("button", { name: "Reset all to $1,000.00" })).toBeInTheDocument(); + expect(screen.getByText(/new team default of \$1,000\.00/)).toBeInTheDocument(); + }); + + it("routes each choice to the matching action", async () => { + const handlers = defaultHandlers(); + const user = userEvent.setup(); + render(); + + await user.click(screen.getByRole("button", { name: "Reset all to $10.00" })); + expect(handlers.onReset).toHaveBeenCalledTimes(1); + + await user.click(screen.getByRole("button", { name: "Keep custom budgets" })); + expect(handlers.onKeep).toHaveBeenCalledTimes(1); + + await user.click(screen.getByRole("button", { name: "Cancel" })); + expect(handlers.onDismiss).toHaveBeenCalledTimes(1); + }); + + it("locks every button while the reset is in flight", () => { + const state: MemberBudgetResetState = { phase: "resetting", pending: pendingFor(2), attempted: 0 }; + render(); + + expect(screen.getByRole("button", { name: "Reset all to $10.00" })).toBeDisabled(); + expect(screen.getByRole("button", { name: "Keep custom budgets" })).toBeDisabled(); + expect(screen.getByRole("button", { name: "Cancel" })).toBeDisabled(); + }); + + it("offers retry or cancel after a failed reset", async () => { + const handlers = defaultHandlers(); + const user = userEvent.setup(); + const state: MemberBudgetResetState = { phase: "resetFailed", pending: pendingFor(2), attempted: 0 }; + render(); + + await user.click(screen.getByRole("button", { name: "Retry reset" })); + expect(handlers.onRetry).toHaveBeenCalledTimes(1); + + await user.click(screen.getByRole("button", { name: "Cancel" })); + expect(handlers.onDismiss).toHaveBeenCalledTimes(1); + + expect(screen.queryByRole("button", { name: "Keep custom budgets" })).not.toBeInTheDocument(); + expect(screen.queryByRole("button", { name: "Reset all to $10.00" })).not.toBeInTheDocument(); + }); + + it("renders nothing while idle", () => { + render(); + + expect(screen.queryByRole("dialog")).not.toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/team/ResetMemberBudgetsDialog.tsx b/ui/litellm-dashboard/src/components/team/ResetMemberBudgetsDialog.tsx new file mode 100644 index 00000000000..95714e3df33 --- /dev/null +++ b/ui/litellm-dashboard/src/components/team/ResetMemberBudgetsDialog.tsx @@ -0,0 +1,71 @@ +import React from "react"; +import { Button } from "@/components/ui/button"; +import { Dialog, DialogContent, DialogFooter, DialogHeader, DialogTitle } from "@/components/ui/dialog"; +import { formatNumberWithCommas } from "@/utils/dataUtils"; +import { pluralize } from "./memberBudgetReset"; +import type { MemberBudgetResetState } from "./useMemberBudgetReset"; + +interface ResetMemberBudgetsDialogProps { + state: MemberBudgetResetState; + onReset: () => void; + onRetry: () => void; + onKeep: () => void; + onDismiss: () => void; +} + +export default function ResetMemberBudgetsDialog({ + state, + onReset, + onRetry, + onKeep, + onDismiss, +}: ResetMemberBudgetsDialogProps) { + const open = state.phase !== "idle"; + const busy = state.phase === "resetting"; + const failed = state.phase === "resetFailed"; + const memberCount = state.phase === "idle" ? 0 : state.pending.userIds.length; + const newBudget = state.phase === "idle" ? 0 : state.pending.newBudget; + + return ( + { + if (!nextOpen && !busy) onDismiss(); + }} + > + + + Reset member budgets? + +

+ {memberCount} {pluralize(memberCount, "member has", "members have")} a custom budget, so the new team default + of ${formatNumberWithCommas(newBudget, 2)} will not apply to {pluralize(memberCount, "that member", "them")}. + Reset {pluralize(memberCount, "it", "them")} to the default, or keep the custom{" "} + {pluralize(memberCount, "budget", "budgets")}? +

+ + {failed ? ( + <> + + + + ) : ( + <> + + + + + )} + + + + ); +} diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.integration.test.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.integration.test.tsx new file mode 100644 index 00000000000..2d84d78d56c --- /dev/null +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.integration.test.tsx @@ -0,0 +1,722 @@ +import { useTeamMetadataSchema } from "@/app/(dashboard)/hooks/teams/useTeamMetadataSchema"; +import * as networking from "@/components/networking"; +import { fireEvent, screen, waitFor } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { renderWithProviders } from "../../../tests/test-utils"; +import { toast } from "@/lib/toast"; +import type { Team } from "../key_team_helpers/key_list"; +import { MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES } from "./memberBudgetReset"; +import TeamInfoView, { type TeamData } from "./TeamInfo"; + +const authState = vi.hoisted(() => ({ userRole: "Admin" })); + +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: () => ({ + token: "123", + accessToken: "123", + userId: "user-1", + userEmail: "user@example.com", + userRole: authState.userRole, + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }), +})); + +vi.mock("next/navigation", () => ({ useRouter: () => ({ push: vi.fn() }) })); + +vi.mock("@/components/networking", () => ({ + serverRootPath: "", + teamInfoCall: vi.fn(), + teamMemberDeleteCall: vi.fn(), + teamMemberAddCall: vi.fn(), + teamMemberUpdateCall: vi.fn(), + teamUpdateCall: vi.fn(), + getGuardrailsList: vi.fn(), + getPoliciesList: vi.fn(), + getPolicyInfoWithGuardrails: vi.fn(), + fetchMCPAccessGroups: vi.fn(), + getTeamPermissionsCall: vi.fn(), + organizationInfoCall: vi.fn(), + getRouterSettingsCall: vi.fn().mockResolvedValue({ fields: [] }), + getPassThroughEndpointsCall: vi.fn().mockResolvedValue({ endpoints: [] }), + fetchMCPServers: vi.fn().mockResolvedValue([]), + fetchMCPToolsets: vi.fn().mockResolvedValue([]), + listMCPTools: vi.fn().mockResolvedValue({ tools: [] }), + vectorStoreListCall: vi.fn().mockResolvedValue({ data: [] }), + getAgentsList: vi.fn().mockResolvedValue({ agents: [] }), + getClaudeCodePluginsList: vi.fn().mockResolvedValue({ plugins: [], count: 0 }), +})); + +const { bulkUpdatePOST } = vi.hoisted(() => ({ bulkUpdatePOST: vi.fn() })); +vi.mock("@/lib/http/api", async (importOriginal) => { + const actual = await importOriginal(); + return { ...actual, fetchClient: { ...actual.fetchClient, POST: bulkUpdatePOST } }; +}); + +const can = vi.fn(); +vi.mock("@/app/(dashboard)/hooks/useCan", () => ({ + default: (...args: unknown[]) => can(...args), +})); + +vi.mock("@/components/utils/dataUtils", () => ({ + copyToClipboard: vi.fn().mockResolvedValue(true), + formatNumberWithCommas: vi.fn((value: number) => value.toLocaleString()), +})); + +vi.mock("@/app/(dashboard)/hooks/teams/useTeamMetadataSchema", () => ({ + useTeamMetadataSchema: vi.fn(() => ({ data: [], isLoading: false })), +})); + +vi.mock("@/app/(dashboard)/hooks/uiSettings/useUISettings", () => ({ + useUISettings: vi.fn(), +})); + +vi.mock("@/app/(dashboard)/hooks/models/useModels", () => ({ + useAllProxyModels: vi.fn(), +})); + +vi.mock("@/app/(dashboard)/hooks/teams/useTeams", () => ({ + useTeam: vi.fn(), +})); + +vi.mock("@/app/(dashboard)/hooks/organizations/useOrganizations", () => ({ + organizationKeys: { all: ["organizations"] }, + useOrganization: vi.fn(), + useOrganizations: vi.fn().mockReturnValue({ data: [], isLoading: false }), +})); + +vi.mock("@/app/(dashboard)/hooks/users/useCurrentUser", () => ({ + useCurrentUser: vi.fn(), +})); + +vi.mock("@/app/(dashboard)/hooks/mcpServers/useMCPServers", () => ({ + useMCPServers: vi.fn(), +})); + +vi.mock("@/app/(dashboard)/hooks/mcpServers/useMCPToolsets", () => ({ + useMCPToolsets: vi.fn(), +})); + +vi.mock("@/components/mcp_server_management/MCPServerSelector", () => ({ + default: ({ + value, + onChange, + }: { + value?: { servers: string[]; accessGroups: string[]; toolsets?: string[] }; + onChange: (next: { servers: string[]; accessGroups: string[]; toolsets: string[] }) => void; + }) => ( + <> + + + + ), +})); + +vi.mock("@/components/team/TeamMemberTab", () => ({ + default: vi.fn(({ setIsAddMemberModalVisible }) => ( +
+ +
+ )), +})); + +vi.mock("@/components/common_components/user_search_modal", () => ({ + default: vi.fn(({ isVisible, onCancel, onSubmit }) => + isVisible ? ( +
+ + +
+ ) : null, + ), +})); + +vi.mock("@/components/team/EditMembership", () => ({ + default: vi.fn(({ visible, onCancel, onSubmit }) => + visible ? ( +
+ + +
+ ) : null, + ), +})); + +vi.mock("@/components/common_components/DeleteResourceModal", () => ({ + default: vi.fn(({ isOpen, onCancel, onOk }) => + isOpen ? ( +
+ + +
+ ) : null, + ), +})); + +vi.mock("@/components/team/member_permissions", () => ({ + default: vi.fn(() =>
Member Permissions
), +})); + +vi.mock("@/components/common_components/ModelAliasManager", () => ({ + default: vi.fn(({ initialModelAliases, onAliasUpdate }) => ( +
+
{JSON.stringify(initialModelAliases)}
+ + +
+ )), +})); + +vi.mock("@/app/(dashboard)/hooks/accessGroups/useAccessGroups", () => ({ + useAccessGroups: vi.fn().mockReturnValue({ + data: [ + { access_group_id: "ag-1", access_group_name: "Group 1", access_mcp_server_ids: [] }, + { access_group_id: "ag-2", access_group_name: "Group 2", access_mcp_server_ids: [] }, + ], + isLoading: false, + isError: false, + }), +})); + +vi.mock("@/components/common_components/AccessGroupSelector", () => ({ + default: ({ value, onChange }: { value?: string[]; onChange?: (next: string[]) => void }) => ( + + ), +})); + +vi.mock("@/app/(dashboard)/hooks/keys/useKeys", () => { + const useKeysResult = { + data: { keys: [], total_count: 0, current_page: 1, total_pages: 1 }, + isPending: false, + isFetching: false, + refetch: vi.fn(), + }; + return { useKeys: vi.fn().mockReturnValue(useKeysResult) }; +}); + +vi.mock("../key_team_helpers/filter_helpers", () => ({ + fetchTeamFilterOptions: vi.fn().mockResolvedValue({ + keyAliases: [], + organizationIds: [], + userIds: [], + }), + fetchAllKeyAliases: vi.fn().mockResolvedValue([]), + fetchAllOrganizations: vi.fn().mockResolvedValue([]), +})); + +import { useAllProxyModels } from "@/app/(dashboard)/hooks/models/useModels"; +import { useKeys } from "@/app/(dashboard)/hooks/keys/useKeys"; +import { useOrganization } from "@/app/(dashboard)/hooks/organizations/useOrganizations"; +import { useTeam } from "@/app/(dashboard)/hooks/teams/useTeams"; +import { useCurrentUser } from "@/app/(dashboard)/hooks/users/useCurrentUser"; +import { useMCPServers } from "@/app/(dashboard)/hooks/mcpServers/useMCPServers"; +import { useMCPToolsets } from "@/app/(dashboard)/hooks/mcpServers/useMCPToolsets"; +import { useAccessGroups } from "@/app/(dashboard)/hooks/accessGroups/useAccessGroups"; +import { useUISettings } from "@/app/(dashboard)/hooks/uiSettings/useUISettings"; + +const mockUseAllProxyModels = vi.mocked(useAllProxyModels); +const mockUseKeys = vi.mocked(useKeys); +const mockUseTeam = vi.mocked(useTeam); +const mockUseOrganization = vi.mocked(useOrganization); +const mockUseCurrentUser = vi.mocked(useCurrentUser); +const mockUseMCPServers = vi.mocked(useMCPServers); +const mockUseMCPToolsets = vi.mocked(useMCPToolsets); +const mockUseAccessGroups = vi.mocked(useAccessGroups); +const mockUseUISettings = vi.mocked(useUISettings); + +const createMockTeamData = (overrides = {}) => ({ + team_id: "123", + team_info: { + team_alias: "Test Team", + team_id: "123", + organization_id: null, + admins: ["admin@test.com"], + members: ["user1@test.com"], + members_with_roles: [ + { + user_id: "user1@test.com", + user_email: "user1@test.com", + role: "member", + spend: 0, + budget_id: "budget1", + }, + ], + metadata: {}, + tpm_limit: null, + rpm_limit: null, + max_budget: null, + budget_duration: null, + models: [], + blocked: false, + spend: 0, + max_parallel_requests: null, + budget_reset_at: null, + model_id: null, + litellm_model_table: null, + created_at: "2024-01-01T00:00:00Z", + team_member_budget_table: null, + guardrails: [], + policies: [], + object_permission: null, + ...overrides, + }, + keys: [], + team_memberships: [], +}); + +const seedDefaultMocks = () => { + mockUseAllProxyModels.mockReturnValue({ + data: { data: [] }, + isLoading: false, + } as unknown as ReturnType); + mockUseTeam.mockReturnValue({ + data: undefined, + isLoading: false, + } as unknown as ReturnType); + mockUseOrganization.mockReturnValue({ + data: undefined, + isLoading: false, + } as unknown as ReturnType); + mockUseCurrentUser.mockReturnValue({ + data: { models: [] }, + isLoading: false, + } as unknown as ReturnType); + mockUseMCPServers.mockReturnValue({ data: [], isLoading: false, isError: false } as unknown as ReturnType< + typeof useMCPServers + >); + mockUseMCPToolsets.mockReturnValue({ data: [], isLoading: false, isError: false } as unknown as ReturnType< + typeof useMCPToolsets + >); + mockUseAccessGroups.mockReturnValue({ + data: [ + { access_group_id: "ag-1", access_group_name: "Group 1", access_mcp_server_ids: [] }, + { access_group_id: "ag-2", access_group_name: "Group 2", access_mcp_server_ids: [] }, + ], + isLoading: false, + isError: false, + } as unknown as ReturnType); + mockUseUISettings.mockReturnValue({ + data: { values: {} }, + isLoading: false, + } as unknown as ReturnType); + mockUseKeys.mockReturnValue({ + data: { keys: [], total_count: 0, current_page: 1, total_pages: 1 }, + isPending: false, + isFetching: false, + refetch: vi.fn(), + } as unknown as ReturnType); + vi.mocked(useTeamMetadataSchema).mockReturnValue({ + data: [], + isLoading: false, + } as unknown as ReturnType); + + can.mockReturnValue(true); + vi.mocked(networking.getGuardrailsList).mockResolvedValue({ guardrails: [] }); + vi.mocked(networking.getPoliciesList).mockResolvedValue({ policies: [] }); + vi.mocked(networking.fetchMCPAccessGroups).mockResolvedValue([]); + vi.mocked(networking.getTeamPermissionsCall).mockResolvedValue({ + all_available_permissions: [], + team_member_permissions: [], + }); +}; +describe("TeamInfoView - member budget reset prompt", () => { + const props = { + teamId: "123", + onUpdate: vi.fn(), + onClose: vi.fn(), + accessToken: "test-token", + is_team_admin: true, + is_proxy_admin: true, + userModels: ["gpt-4", "gpt-3.5-turbo"], + editTeam: false, + premiumUser: false, + }; + + const customBudgetMembership = ( + userId: string, + maxBudget: number | null = 50, + ): TeamData["team_memberships"][number] => ({ + user_id: userId, + team_id: "123", + budget_id: `budget-${userId}`, + budget_source: "custom", + spend: 0, + total_spend: 0, + litellm_budget_table: { + budget_id: `budget-${userId}`, + soft_budget: null, + max_budget: maxBudget, + max_parallel_requests: null, + tpm_limit: null, + rpm_limit: null, + model_max_budget: null, + budget_duration: null, + budget_reset_at: null, + }, + }); + + const savedTeam: Team = { + team_id: "123", + team_alias: "Test Team", + models: [], + max_budget: null, + budget_duration: null, + tpm_limit: null, + rpm_limit: null, + organization_id: "org-1", + created_at: "2024-01-01T00:00:00Z", + keys: [], + members_with_roles: [], + spend: 0, + }; + + const openEditorWithCustomMembers = async ( + user: ReturnType, + userIds: string[] = ["user-custom"], + maxBudget: number | null = 50, + ) => { + const data = { + ...createMockTeamData({ + team_member_budget_table: { max_budget: 10, budget_duration: null, tpm_limit: null, rpm_limit: null }, + }), + team_memberships: userIds.map((id) => customBudgetMembership(id, maxBudget)), + } as TeamData; + vi.mocked(networking.teamInfoCall).mockResolvedValue(data); + vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: savedTeam, team_id: "123" }); + + renderWithProviders(); + await waitFor(() => expect(screen.queryAllByText("Test Team").length).toBeGreaterThan(0)); + await user.click(screen.getByRole("tab", { name: "Settings" })); + await user.click(await screen.findByRole("button", { name: /edit settings/i })); + await user.click(screen.getByText("Team Member Settings")); + return await screen.findByLabelText("Default Budget (USD)"); + }; + + const submitNewDefault = async (user: ReturnType, input: HTMLElement, value: string) => { + fireEvent.change(input, { target: { value } }); + await user.click(screen.getByRole("button", { name: /save changes/i })); + }; + + beforeEach(() => { + seedDefaultMocks(); + bulkUpdatePOST.mockResolvedValue({ data: { data: [] } }); + }); + + afterEach(() => { + vi.clearAllMocks(); + }); + + it("holds the save behind the prompt when the default changes while a member has a custom budget", async () => { + const user = userEvent.setup({ delay: null }); + const input = await openEditorWithCustomMembers(user); + + await submitNewDefault(user, input, "20"); + + expect(await screen.findByText("Reset member budgets?")).toBeInTheDocument(); + expect(screen.getByText(/1 member has a custom budget/)).toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Reset to $20.00" })).toBeInTheDocument(); + expect(networking.teamUpdateCall).not.toHaveBeenCalled(); + expect(bulkUpdatePOST).not.toHaveBeenCalled(); + }); + + it("saves the new default and leaves custom budgets alone on keep", async () => { + const user = userEvent.setup({ delay: null }); + const input = await openEditorWithCustomMembers(user); + + await submitNewDefault(user, input, "20"); + await user.click(await screen.findByRole("button", { name: "Keep custom budget" })); + + await waitFor(() => expect(networking.teamUpdateCall).toHaveBeenCalled()); + expect(vi.mocked(networking.teamUpdateCall).mock.calls[0][1].team_member_budget).toBe(20); + expect(bulkUpdatePOST).not.toHaveBeenCalled(); + await waitFor(() => expect(screen.queryByText("Reset member budgets?")).not.toBeInTheDocument()); + }); + + it("saves the new default then clears each custom budget through the bulk endpoint on reset all", async () => { + const user = userEvent.setup({ delay: null }); + const input = await openEditorWithCustomMembers(user, ["user-a", "user-b"]); + bulkUpdatePOST.mockResolvedValue({ + data: { + data: [ + { success: true, user_id: "user-a" }, + { success: true, user_id: "user-b" }, + ], + }, + }); + + await submitNewDefault(user, input, "20"); + await user.click(await screen.findByRole("button", { name: "Reset all to $20.00" })); + + await waitFor(() => expect(bulkUpdatePOST).toHaveBeenCalledTimes(1)); + expect(networking.teamUpdateCall).toHaveBeenCalled(); + expect(vi.mocked(networking.teamUpdateCall).mock.calls[0][1].team_member_budget).toBe(20); + expect(bulkUpdatePOST).toHaveBeenCalledWith("/management/v1/teams/{team_id}/members/bulk_update", { + params: { path: { team_id: "123" } }, + body: { + members: [ + { user_id: "user-a", max_budget_in_team: null }, + { user_id: "user-b", max_budget_in_team: null }, + ], + }, + }); + await waitFor(() => expect(toast.success).toHaveBeenCalledWith("Reset 2 member budgets to the team default")); + }); + + it("keeps the prompt mounted while the member-budget reset is in flight", async () => { + const user = userEvent.setup({ delay: null }); + const input = await openEditorWithCustomMembers(user); + const bulk = Promise.withResolvers<{ data: { data: { success: boolean; user_id: string }[] } }>(); + vi.mocked(networking.teamInfoCall).mockImplementationOnce(() => new Promise(() => {})); + bulkUpdatePOST.mockImplementationOnce(() => bulk.promise); + + await submitNewDefault(user, input, "20"); + await user.click(await screen.findByRole("button", { name: "Reset to $20.00" })); + + await waitFor(() => expect(bulkUpdatePOST).toHaveBeenCalled()); + expect(screen.getByText("Reset member budgets?")).toBeInTheDocument(); + expect(screen.queryByText("Loading...")).not.toBeInTheDocument(); + + bulk.resolve({ data: { data: [{ success: true, user_id: "user-custom" }] } }); + await waitFor(() => expect(screen.queryByText("Reset member budgets?")).not.toBeInTheDocument()); + }); + + it("offers a retry after a failed reset and never re-saves the team", async () => { + const user = userEvent.setup({ delay: null }); + const input = await openEditorWithCustomMembers(user); + bulkUpdatePOST + .mockRejectedValueOnce(new Error("bulk update down")) + .mockResolvedValueOnce({ data: { data: [{ success: true, user_id: "user-custom" }] } }); + + await submitNewDefault(user, input, "20"); + await user.click(await screen.findByRole("button", { name: "Reset to $20.00" })); + + await waitFor(() => + expect(toast.fromError).toHaveBeenCalledWith("Team updated, but member budgets could not be reset"), + ); + expect(screen.getByText("Reset member budgets?")).toBeInTheDocument(); + expect(screen.queryByText("Loading...")).not.toBeInTheDocument(); + await waitFor(() => expect(vi.mocked(networking.teamInfoCall).mock.calls.length).toBeGreaterThan(1)); + + await user.click(screen.getByRole("button", { name: "Retry reset" })); + + await waitFor(() => expect(bulkUpdatePOST).toHaveBeenCalledTimes(2)); + expect(vi.mocked(networking.teamUpdateCall).mock.calls).toHaveLength(1); + await waitFor(() => expect(toast.success).toHaveBeenCalledWith("Reset 1 member budget to the team default")); + await waitFor(() => expect(screen.queryByText("Reset member budgets?")).not.toBeInTheDocument()); + }); + + it("returns to the prompt without its own toast when the team save fails", async () => { + const user = userEvent.setup({ delay: null }); + const input = await openEditorWithCustomMembers(user); + vi.mocked(networking.teamUpdateCall).mockRejectedValueOnce(new Error("save failed")); + + await submitNewDefault(user, input, "20"); + await user.click(await screen.findByRole("button", { name: "Reset to $20.00" })); + + expect(await screen.findByRole("button", { name: "Reset to $20.00" })).toBeEnabled(); + expect(bulkUpdatePOST).not.toHaveBeenCalled(); + expect(toast.fromError).not.toHaveBeenCalled(); + expect(toast.error).not.toHaveBeenCalled(); + expect(screen.queryByText("Loading...")).not.toBeInTheDocument(); + }); + + it("reports how many budgets were already reset when a later batch fails", async () => { + const user = userEvent.setup({ delay: null }); + const userIds = Array.from({ length: MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES + 1 }, (_, i) => `user-${i}`); + const input = await openEditorWithCustomMembers(user, userIds); + bulkUpdatePOST + .mockResolvedValueOnce({ + data: { + data: userIds.slice(0, MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES).map((user_id) => ({ success: true, user_id })), + }, + }) + .mockRejectedValueOnce(new Error("second batch down")) + .mockResolvedValueOnce({ + data: { data: [{ success: true, user_id: `user-${MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES}` }] }, + }); + + await submitNewDefault(user, input, "20"); + await user.click(await screen.findByRole("button", { name: "Reset all to $20.00" })); + + await waitFor(() => + expect(toast.error).toHaveBeenCalledWith( + `Reset ${MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES} of ${MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES + 1} member budgets; the rest could not be reset`, + ), + ); + expect(bulkUpdatePOST).toHaveBeenCalledTimes(2); + + await user.click(screen.getByRole("button", { name: "Retry reset" })); + + await waitFor(() => expect(bulkUpdatePOST).toHaveBeenCalledTimes(3)); + expect(vi.mocked(bulkUpdatePOST).mock.calls[2][1]).toEqual({ + params: { path: { team_id: "123" } }, + body: { + members: [{ user_id: `user-${MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES}`, max_budget_in_team: null }], + }, + }); + expect(vi.mocked(networking.teamUpdateCall).mock.calls).toHaveLength(1); + await waitFor(() => + expect(toast.success).toHaveBeenCalledWith( + `Reset ${MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES + 1} member budgets to the team default`, + ), + ); + }); + + it("surfaces a failure toast when some members cannot be reset", async () => { + const user = userEvent.setup({ delay: null }); + const input = await openEditorWithCustomMembers(user, ["user-a", "user-b"]); + bulkUpdatePOST.mockResolvedValue({ + data: { + data: [ + { success: true, user_id: "user-a" }, + { success: false, user_id: "user-b", error: "no such member" }, + ], + }, + }); + + await submitNewDefault(user, input, "20"); + await user.click(await screen.findByRole("button", { name: "Reset all to $20.00" })); + + await waitFor(() => + expect(toast.error).toHaveBeenCalledWith("Team updated, but 1 member budget could not be reset"), + ); + }); + + it("aborts the save entirely on cancel", async () => { + const user = userEvent.setup({ delay: null }); + const input = await openEditorWithCustomMembers(user); + + await submitNewDefault(user, input, "20"); + await user.click(await screen.findByRole("button", { name: "Cancel" })); + + await waitFor(() => expect(screen.queryByText("Reset member budgets?")).not.toBeInTheDocument()); + expect(networking.teamUpdateCall).not.toHaveBeenCalled(); + expect(bulkUpdatePOST).not.toHaveBeenCalled(); + }); + + it("saves directly when no member carries a custom budget", async () => { + const user = userEvent.setup({ delay: null }); + const input = await openEditorWithCustomMembers(user, []); + + await submitNewDefault(user, input, "20"); + + await waitFor(() => expect(networking.teamUpdateCall).toHaveBeenCalled()); + expect(vi.mocked(networking.teamUpdateCall).mock.calls[0][1].team_member_budget).toBe(20); + expect(screen.queryByText("Reset member budgets?")).not.toBeInTheDocument(); + expect(bulkUpdatePOST).not.toHaveBeenCalled(); + }); + + it("saves directly when the default is resubmitted unchanged", async () => { + const user = userEvent.setup({ delay: null }); + await openEditorWithCustomMembers(user); + + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => expect(networking.teamUpdateCall).toHaveBeenCalled()); + expect(vi.mocked(networking.teamUpdateCall).mock.calls[0][1].team_member_budget).toBe(10); + expect(screen.queryByText("Reset member budgets?")).not.toBeInTheDocument(); + expect(bulkUpdatePOST).not.toHaveBeenCalled(); + }); + + it("dismisses a pending prompt when the selected team changes", async () => { + const user = userEvent.setup({ delay: null }); + const data = { + ...createMockTeamData({ + team_member_budget_table: { max_budget: 10, budget_duration: null, tpm_limit: null, rpm_limit: null }, + }), + team_memberships: [customBudgetMembership("user-custom")], + } as TeamData; + vi.mocked(networking.teamInfoCall).mockResolvedValue(data); + vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: savedTeam, team_id: "123" }); + + const { rerender } = renderWithProviders(); + await waitFor(() => expect(screen.queryAllByText("Test Team").length).toBeGreaterThan(0)); + await user.click(screen.getByRole("tab", { name: "Settings" })); + await user.click(await screen.findByRole("button", { name: /edit settings/i })); + await user.click(screen.getByText("Team Member Settings")); + const input = await screen.findByLabelText("Default Budget (USD)"); + + await submitNewDefault(user, input, "20"); + expect(await screen.findByText("Reset member budgets?")).toBeInTheDocument(); + + rerender(); + + await waitFor(() => expect(screen.queryByText("Reset member budgets?")).not.toBeInTheDocument()); + expect(bulkUpdatePOST).not.toHaveBeenCalled(); + }); + + it("does not reopen the dialog or refetch the old team when the team changes mid-reset", async () => { + const user = userEvent.setup({ delay: null }); + const data = { + ...createMockTeamData({ + team_member_budget_table: { max_budget: 10, budget_duration: null, tpm_limit: null, rpm_limit: null }, + }), + team_memberships: [customBudgetMembership("user-custom")], + } as TeamData; + vi.mocked(networking.teamInfoCall).mockResolvedValue(data); + vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: savedTeam, team_id: "123" }); + const bulkDone = Promise.withResolvers<{ data: { data: { success: boolean; user_id: string }[] } }>(); + bulkUpdatePOST.mockReturnValue(bulkDone.promise); + + const { rerender } = renderWithProviders(); + await waitFor(() => expect(screen.queryAllByText("Test Team").length).toBeGreaterThan(0)); + await user.click(screen.getByRole("tab", { name: "Settings" })); + await user.click(await screen.findByRole("button", { name: /edit settings/i })); + await user.click(screen.getByText("Team Member Settings")); + const input = await screen.findByLabelText("Default Budget (USD)"); + + await submitNewDefault(user, input, "20"); + await user.click(await screen.findByRole("button", { name: "Reset to $20.00" })); + await waitFor(() => expect(bulkUpdatePOST).toHaveBeenCalled()); + + const infoCallsBeforeSwitch = vi.mocked(networking.teamInfoCall).mock.calls.length; + rerender(); + bulkDone.resolve({ data: { data: [{ success: true, user_id: "user-custom" }] } }); + + await waitFor(() => expect(screen.queryByText("Reset member budgets?")).not.toBeInTheDocument()); + await waitFor(() => + expect( + vi + .mocked(networking.teamInfoCall) + .mock.calls.slice(infoCallsBeforeSwitch) + .some((call) => call[1] === "456"), + ).toBe(true), + ); + const infoCallsAfterSwitch = vi.mocked(networking.teamInfoCall).mock.calls.slice(infoCallsBeforeSwitch); + expect(infoCallsAfterSwitch.every((call) => call[1] === "456")).toBe(true); + expect(screen.queryByRole("button", { name: "Retry reset" })).not.toBeInTheDocument(); + }); + + it("saves directly when a custom member's cap is null and already inherits the default", async () => { + const user = userEvent.setup({ delay: null }); + const input = await openEditorWithCustomMembers(user, ["user-limits-only"], null); + + await submitNewDefault(user, input, "20"); + + await waitFor(() => expect(networking.teamUpdateCall).toHaveBeenCalled()); + expect(vi.mocked(networking.teamUpdateCall).mock.calls[0][1].team_member_budget).toBe(20); + expect(screen.queryByText("Reset member budgets?")).not.toBeInTheDocument(); + expect(bulkUpdatePOST).not.toHaveBeenCalled(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx index 18b1505e342..37f63433d2a 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx @@ -55,6 +55,12 @@ vi.mock("@/components/networking", () => ({ getClaudeCodePluginsList: vi.fn().mockResolvedValue({ plugins: [], count: 0 }), })); +const { bulkUpdatePOST } = vi.hoisted(() => ({ bulkUpdatePOST: vi.fn() })); +vi.mock("@/lib/http/api", async (importOriginal) => { + const actual = await importOriginal(); + return { ...actual, fetchClient: { ...actual.fetchClient, POST: bulkUpdatePOST } }; +}); + const can = vi.fn(); vi.mock("@/app/(dashboard)/hooks/useCan", () => ({ default: (...args: unknown[]) => can(...args), diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index d9e308e9d6f..78ed507216a 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -119,6 +119,10 @@ import { } from "./tabVisibilityUtils"; import TeamMembersComponent from "./TeamMemberTab"; import { TeamVirtualKeysTable } from "./TeamVirtualKeysTable"; +import ResetMemberBudgetsDialog from "./ResetMemberBudgetsDialog"; +import { customBudgetMemberUserIds, shouldPromptMemberBudgetReset } from "./memberBudgetReset"; +import { useMemberBudgetReset } from "./useMemberBudgetReset"; +import { fetchClient } from "@/lib/http/api"; const UI_MANAGED_METADATA_KEYS: ReadonlySet = new Set([ "logging", @@ -884,17 +888,41 @@ const TeamInfoView: React.FC = ({ const persistTeamUpdate = async (token: string, updateData: Record) => { await teamUpdateCall(token, updateData); queryClient.invalidateQueries({ queryKey: organizationKeys.all }); - - toast.success("Team settings updated successfully"); setIsEditing(false); - fetchTeamInfo(); }; + const memberBudgetReset = useMemberBudgetReset({ + saveTeam: async (updateData) => { + if (!accessToken) return; + setIsTeamSaving(true); + try { + await persistTeamUpdate(accessToken, updateData); + } finally { + setIsTeamSaving(false); + } + }, + resetMemberBudgets: async (bulkTeamId, userIds) => { + const { data } = await fetchClient.POST("/management/v1/teams/{team_id}/members/bulk_update", { + params: { path: { team_id: bulkTeamId } }, + body: { members: userIds.map((user_id) => ({ user_id, max_budget_in_team: null })) }, + }); + return data?.data ?? []; + }, + refreshTeamData, + }); + + const { dismiss: dismissMemberBudgetReset } = memberBudgetReset; + useEffect(() => { + dismissMemberBudgetReset(); + }, [teamId, dismissMemberBudgetReset]); + const saveTeamAdminSettings = async (changes: TeamAdminSettingsChanges) => { if (!accessToken) return; setIsTeamSaving(true); try { await persistTeamUpdate(accessToken, { team_id: teamId, ...changes }); + toast.success("Team settings updated successfully"); + await fetchTeamInfo(); } catch (error) { console.error("Error updating team:", error); } finally { @@ -1005,8 +1033,10 @@ const TeamInfoView: React.FC = ({ updateData.max_budget = mapEmptyStringToNull(updateData.max_budget); updateData.team_member_budget_duration = values.team_member_budget_duration; - if (values.team_member_budget !== undefined) { - updateData.team_member_budget = Number(values.team_member_budget); + const newTeamMemberBudget = + values.team_member_budget !== undefined ? Number(values.team_member_budget) : undefined; + if (newTeamMemberBudget !== undefined) { + updateData.team_member_budget = newTeamMemberBudget; } if (values.team_member_key_duration !== undefined) { @@ -1152,7 +1182,28 @@ const TeamInfoView: React.FC = ({ } } + const customBudgetUserIds = customBudgetMemberUserIds(teamData?.team_memberships ?? []); + if ( + newTeamMemberBudget !== undefined && + shouldPromptMemberBudgetReset( + newTeamMemberBudget, + info.team_member_budget_table?.max_budget, + customBudgetUserIds, + ) + ) { + const pendingReset = { + teamId, + updateData, + userIds: customBudgetUserIds, + newBudget: newTeamMemberBudget, + }; + memberBudgetReset.prompt(pendingReset); + return; + } + await persistTeamUpdate(accessToken, updateData); + toast.success("Team settings updated successfully"); + await fetchTeamInfo(); } catch (error) { console.error("Error updating team:", error); } finally { @@ -2426,6 +2477,14 @@ const TeamInfoView: React.FC = ({ onOk={handleDeleteConfirm} confirmLoading={isDeleting} /> + + ); }; diff --git a/ui/litellm-dashboard/src/components/team/TeamMemberTab.test.tsx b/ui/litellm-dashboard/src/components/team/TeamMemberTab.test.tsx index 8652ffa7de2..bd7398bc6c8 100644 --- a/ui/litellm-dashboard/src/components/team/TeamMemberTab.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamMemberTab.test.tsx @@ -682,6 +682,44 @@ describe("TeamMembersComponent", () => { expect(inheritedRow).toHaveTextContent("$25.00"); }); + it("caps a Custom member at the team default when the private row has no budget limit", () => { + const base = createMockTeamData(); + renderTab( + createMockTeamData({ + team_info: { + ...base.team_info, + team_member_budget_table: { max_budget: 20, budget_duration: null, tpm_limit: null, rpm_limit: null }, + }, + team_memberships: [ + { + user_id: "user2@test.com", + team_id: "team-123", + budget_id: "budget2", + budget_source: "custom", + spend: 0, + total_spend: 0, + litellm_budget_table: { + budget_id: "budget3", + soft_budget: null, + max_budget: null, + max_parallel_requests: null, + tpm_limit: null, + rpm_limit: 100, + model_max_budget: null, + budget_duration: null, + budget_reset_at: null, + }, + }, + ], + }), + ); + + const row = screen.getByRole("row", { name: /user2@test\.com/ }); + expect(within(row).getByTestId("member-budget-source")).toHaveTextContent("Custom"); + expect(row).toHaveTextContent("$20.00"); + expect(row).not.toHaveTextContent("Unlimited"); + }); + it("shows no source label for a member with neither a custom nor a team budget", () => { renderTab(createMockTeamData({ team_memberships: [] })); diff --git a/ui/litellm-dashboard/src/components/team/TeamMemberTab.tsx b/ui/litellm-dashboard/src/components/team/TeamMemberTab.tsx index 660416504fe..e8576ce6dfe 100644 --- a/ui/litellm-dashboard/src/components/team/TeamMemberTab.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamMemberTab.tsx @@ -106,10 +106,7 @@ export default function TeamMemberTab({ 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 ?? - (membership?.budget_source === "team_default" ? teamDefaultBudget : null) - ); + return membership?.litellm_budget_table?.max_budget ?? teamDefaultBudget; }; // Helper function to get rate limits for a user diff --git a/ui/litellm-dashboard/src/components/team/memberBudgetReset.test.ts b/ui/litellm-dashboard/src/components/team/memberBudgetReset.test.ts new file mode 100644 index 00000000000..e673840203f --- /dev/null +++ b/ui/litellm-dashboard/src/components/team/memberBudgetReset.test.ts @@ -0,0 +1,101 @@ +import { describe, expect, it } from "vitest"; +import { + chunk, + customBudgetMemberUserIds, + MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES, + pluralize, + shouldPromptMemberBudgetReset, +} from "./memberBudgetReset"; + +describe("customBudgetMemberUserIds", () => { + const customRow = (user_id: string, max_budget: number | null = 50) => ({ + user_id, + budget_source: "custom" as const, + litellm_budget_table: { max_budget }, + }); + + it("returns only members whose budget_source is custom with a private cap", () => { + const memberships = [ + customRow("u-custom"), + { user_id: "u-default", budget_source: "team_default" as const }, + { user_id: "u-none", budget_source: "none" as const }, + ]; + + expect(customBudgetMemberUserIds(memberships)).toEqual(["u-custom"]); + }); + + it("skips a custom row whose cap is null since it already inherits the default", () => { + const memberships = [customRow("u-rate-limits-only", null), customRow("u-custom")]; + + expect(customBudgetMemberUserIds(memberships)).toEqual(["u-custom"]); + }); + + it("skips a custom row that has no budget table at all", () => { + const memberships = [ + { user_id: "u-no-table", budget_source: "custom" as const, litellm_budget_table: null }, + customRow("u-custom"), + ]; + + expect(customBudgetMemberUserIds(memberships)).toEqual(["u-custom"]); + }); +}); + +describe("shouldPromptMemberBudgetReset", () => { + it("prompts when the default changes and custom-budget members exist", () => { + expect(shouldPromptMemberBudgetReset(10, 5, ["u-1"])).toBe(true); + }); + + it("prompts when a default is set for the first time", () => { + expect(shouldPromptMemberBudgetReset(10, null, ["u-1"])).toBe(true); + expect(shouldPromptMemberBudgetReset(10, undefined, ["u-1"])).toBe(true); + }); + + it("does not prompt when the submitted budget is unchanged", () => { + expect(shouldPromptMemberBudgetReset(10, 10, ["u-1"])).toBe(false); + }); + + it("does not prompt when the budget is cleared or zero", () => { + expect(shouldPromptMemberBudgetReset(undefined, 10, ["u-1"])).toBe(false); + expect(shouldPromptMemberBudgetReset(0, 10, ["u-1"])).toBe(false); + }); + + it("does not prompt when no member has a custom budget", () => { + expect(shouldPromptMemberBudgetReset(10, 5, [])).toBe(false); + }); +}); + +describe("chunk", () => { + it("splits selections larger than the bulk endpoint limit", () => { + const ids = Array.from({ length: MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES + 1 }, (_, i) => `u-${i}`); + + const chunks = chunk(ids, MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES); + + expect(chunks).toHaveLength(2); + expect(chunks[0]).toHaveLength(MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES); + expect(chunks[1]).toEqual([`u-${MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES}`]); + }); + + it("keeps a selection under the limit in a single chunk", () => { + expect(chunk(["u-1", "u-2"], MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES)).toEqual([["u-1", "u-2"]]); + }); + + it("returns no chunks for an empty selection", () => { + expect(chunk([], MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES)).toEqual([]); + }); + + it("returns no chunks for a non-positive size", () => { + expect(chunk(["u-1"], 0)).toEqual([]); + expect(chunk(["u-1"], -3)).toEqual([]); + }); +}); + +describe("pluralize", () => { + it("uses the singular form for exactly one", () => { + expect(pluralize(1, "budget", "budgets")).toBe("budget"); + }); + + it("uses the plural form for zero and for more than one", () => { + expect(pluralize(0, "budget", "budgets")).toBe("budgets"); + expect(pluralize(3, "budget", "budgets")).toBe("budgets"); + }); +}); diff --git a/ui/litellm-dashboard/src/components/team/memberBudgetReset.ts b/ui/litellm-dashboard/src/components/team/memberBudgetReset.ts new file mode 100644 index 00000000000..2b24a5ebbb0 --- /dev/null +++ b/ui/litellm-dashboard/src/components/team/memberBudgetReset.ts @@ -0,0 +1,38 @@ +import type { components } from "@/lib/http/schema"; +import type { TeamMembership } from "./TeamInfo"; + +export const MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES = 500; + +export type TeamUpdatePayload = components["schemas"]["UpdateTeamRequest"]; + +export interface MemberBudgetResetPending { + readonly teamId: string; + readonly updateData: TeamUpdatePayload; + readonly userIds: readonly string[]; + readonly newBudget: number; +} + +type MembershipBudgetRow = Pick & { + litellm_budget_table?: Pick | null; +}; + +export const customBudgetMemberUserIds = (memberships: readonly MembershipBudgetRow[]): string[] => + memberships + .filter((m) => m.budget_source === "custom" && m.litellm_budget_table?.max_budget != null) + .map((m) => m.user_id); + +export const shouldPromptMemberBudgetReset = ( + nextBudget: number | undefined, + previousBudget: number | null | undefined, + customBudgetUserIds: string[], +): boolean => { + const budgetChanged = typeof nextBudget === "number" && nextBudget > 0 && nextBudget !== previousBudget; + return budgetChanged && customBudgetUserIds.length > 0; +}; + +export const chunk = (items: readonly T[], size: number): T[][] => { + if (size <= 0) return []; + return Array.from({ length: Math.ceil(items.length / size) }, (_, i) => items.slice(i * size, i * size + size)); +}; + +export const pluralize = (count: number, singular: string, plural: string): string => (count === 1 ? singular : plural); diff --git a/ui/litellm-dashboard/src/components/team/useMemberBudgetReset.test.tsx b/ui/litellm-dashboard/src/components/team/useMemberBudgetReset.test.tsx new file mode 100644 index 00000000000..71318b635a3 --- /dev/null +++ b/ui/litellm-dashboard/src/components/team/useMemberBudgetReset.test.tsx @@ -0,0 +1,345 @@ +import { act, renderHook } from "@testing-library/react"; +import { afterEach, describe, expect, it, vi } from "vitest"; +import { toast } from "@/lib/toast"; +import { MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES } from "./memberBudgetReset"; +import { useMemberBudgetReset, type MemberBudgetResetGateway } from "./useMemberBudgetReset"; + +const buildGateway = () => ({ + saveTeam: vi.fn(async (_updateData: Record) => {}), + resetMemberBudgets: vi.fn(async (_teamId: string, userIds: readonly string[]) => + userIds.map((user_id) => ({ success: true, user_id })), + ), + refreshTeamData: vi.fn(async () => {}), +}); + +const pendingFor = (userIds: string[]) => ({ + teamId: "team-123", + updateData: { team_id: "team-123", team_member_budget: 20 }, + userIds, + newBudget: 20, +}); + +const renderReset = (gateway: MemberBudgetResetGateway) => renderHook(() => useMemberBudgetReset(gateway)); + +describe("useMemberBudgetReset", () => { + afterEach(() => { + vi.clearAllMocks(); + }); + + it("starts idle and prompts with the pending update", async () => { + const gateway = buildGateway(); + const { result } = renderReset(gateway); + + expect(result.current.state.phase).toBe("idle"); + + await act(async () => { + result.current.prompt(pendingFor(["u-1"])); + }); + + expect(result.current.state).toEqual({ + phase: "prompting", + pending: pendingFor(["u-1"]), + }); + expect(gateway.saveTeam).not.toHaveBeenCalled(); + }); + + it("saves the team once, resets every member, and refreshes on success", async () => { + const gateway = buildGateway(); + const { result } = renderReset(gateway); + + await act(async () => { + result.current.prompt(pendingFor(["u-1", "u-2"])); + }); + await act(async () => { + await result.current.reset(); + }); + + expect(gateway.saveTeam).toHaveBeenCalledTimes(1); + expect(gateway.saveTeam).toHaveBeenCalledWith({ team_id: "team-123", team_member_budget: 20 }); + expect(gateway.resetMemberBudgets).toHaveBeenCalledWith("team-123", ["u-1", "u-2"]); + expect(toast.success).toHaveBeenCalledWith("Reset 2 member budgets to the team default"); + expect(gateway.refreshTeamData).toHaveBeenCalledTimes(1); + expect(result.current.state.phase).toBe("idle"); + }); + + it("resets members in batches no larger than the bulk endpoint limit", async () => { + const userIds = Array.from({ length: MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES + 1 }, (_, i) => `u-${i}`); + const gateway = buildGateway(); + const { result } = renderReset(gateway); + + await act(async () => { + result.current.prompt(pendingFor(userIds)); + }); + await act(async () => { + await result.current.reset(); + }); + + expect(gateway.resetMemberBudgets).toHaveBeenCalledTimes(2); + expect(gateway.resetMemberBudgets.mock.calls[0][0]).toBe("team-123"); + expect(gateway.resetMemberBudgets.mock.calls[0][1]).toHaveLength(MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES); + expect(gateway.resetMemberBudgets.mock.calls[1][1]).toEqual([`u-${MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES}`]); + expect(result.current.state.phase).toBe("idle"); + }); + + it("resets against the team carried by the pending update", async () => { + const gateway = buildGateway(); + const { result } = renderReset(gateway); + + const pending = { + teamId: "team-999", + updateData: { team_id: "team-999", team_member_budget: 20 }, + userIds: ["u-1"], + newBudget: 20, + }; + await act(async () => { + result.current.prompt(pending); + }); + await act(async () => { + await result.current.reset(); + }); + + expect(gateway.resetMemberBudgets).toHaveBeenCalledWith("team-999", ["u-1"]); + expect(gateway.resetMemberBudgets).not.toHaveBeenCalledWith("team-123", expect.anything()); + }); + + it("returns to prompting without its own toast when the team save fails", async () => { + const gateway = buildGateway(); + gateway.saveTeam.mockRejectedValueOnce(new Error("team update failed")); + const { result } = renderReset(gateway); + + await act(async () => { + result.current.prompt(pendingFor(["u-1"])); + }); + await act(async () => { + await result.current.reset(); + }); + + expect(result.current.state.phase).toBe("prompting"); + expect(gateway.resetMemberBudgets).not.toHaveBeenCalled(); + expect(gateway.refreshTeamData).not.toHaveBeenCalled(); + expect(toast.error).not.toHaveBeenCalled(); + expect(toast.fromError).not.toHaveBeenCalled(); + expect(toast.success).not.toHaveBeenCalled(); + }); + + it("retries only the unsent members when a later batch fails, without re-saving the team", async () => { + const userIds = Array.from({ length: MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES + 1 }, (_, i) => `u-${i}`); + const gateway = buildGateway(); + gateway.resetMemberBudgets + .mockResolvedValueOnce( + userIds.slice(0, MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES).map((user_id) => ({ success: true, user_id })), + ) + .mockRejectedValueOnce(new Error("second batch failed")); + const { result } = renderReset(gateway); + + await act(async () => { + result.current.prompt(pendingFor(userIds)); + }); + await act(async () => { + await result.current.reset(); + }); + + expect(result.current.state.phase).toBe("resetFailed"); + expect(toast.error).toHaveBeenCalledWith( + `Reset ${MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES} of ${MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES + 1} member budgets; the rest could not be reset`, + ); + + gateway.resetMemberBudgets.mockResolvedValueOnce([ + { success: true, user_id: `u-${MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES}` }, + ]); + await act(async () => { + await result.current.retry(); + }); + + expect(gateway.saveTeam).toHaveBeenCalledTimes(1); + expect(gateway.resetMemberBudgets).toHaveBeenCalledTimes(3); + expect(gateway.resetMemberBudgets.mock.calls[2][0]).toBe("team-123"); + expect(gateway.resetMemberBudgets.mock.calls[2][1]).toEqual([`u-${MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES}`]); + expect(toast.success).toHaveBeenCalledWith( + `Reset ${MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES + 1} member budgets to the team default`, + ); + expect(result.current.state.phase).toBe("idle"); + }); + + it("reports partial progress when the first batch fails", async () => { + const gateway = buildGateway(); + gateway.resetMemberBudgets.mockRejectedValueOnce(new Error("batch failed")); + const { result } = renderReset(gateway); + + await act(async () => { + result.current.prompt(pendingFor(["u-1", "u-2"])); + }); + await act(async () => { + await result.current.reset(); + }); + + expect(result.current.state.phase).toBe("resetFailed"); + expect(toast.fromError).toHaveBeenCalledWith("Team updated, but member budgets could not be reset"); + expect(gateway.refreshTeamData).toHaveBeenCalledTimes(1); + }); + + it("reports members the backend could not reset and closes", async () => { + const gateway = buildGateway(); + gateway.resetMemberBudgets.mockResolvedValueOnce([ + { success: true, user_id: "u-1" }, + { success: false, user_id: "u-2" }, + ]); + const { result } = renderReset(gateway); + + await act(async () => { + result.current.prompt(pendingFor(["u-1", "u-2"])); + }); + await act(async () => { + await result.current.reset(); + }); + + expect(toast.error).toHaveBeenCalledWith("Team updated, but 1 member budget could not be reset"); + expect(toast.success).not.toHaveBeenCalled(); + expect(result.current.state.phase).toBe("idle"); + expect(gateway.refreshTeamData).toHaveBeenCalledTimes(1); + }); + + it("saves the team once and closes on keep custom", async () => { + const gateway = buildGateway(); + const { result } = renderReset(gateway); + + await act(async () => { + result.current.prompt(pendingFor(["u-1"])); + }); + await act(async () => { + await result.current.keepCustom(); + }); + + expect(gateway.saveTeam).toHaveBeenCalledTimes(1); + expect(toast.success).toHaveBeenCalledWith("Team settings updated successfully"); + expect(gateway.resetMemberBudgets).not.toHaveBeenCalled(); + expect(gateway.refreshTeamData).toHaveBeenCalledTimes(1); + expect(result.current.state.phase).toBe("idle"); + }); + + it("abandons the pending update on dismiss", async () => { + const gateway = buildGateway(); + const { result } = renderReset(gateway); + + await act(async () => { + result.current.prompt(pendingFor(["u-1"])); + }); + await act(async () => { + result.current.dismiss(); + }); + + expect(result.current.state.phase).toBe("idle"); + expect(gateway.saveTeam).not.toHaveBeenCalled(); + expect(gateway.resetMemberBudgets).not.toHaveBeenCalled(); + }); + + it("skips the refresh and stays idle when dismissed while a reset is in flight", async () => { + const gateway = buildGateway(); + const bulkDone = Promise.withResolvers<{ success: boolean; user_id: string }[]>(); + gateway.resetMemberBudgets.mockReturnValueOnce(bulkDone.promise); + const { result } = renderReset(gateway); + + await act(async () => { + result.current.prompt(pendingFor(["u-1"])); + }); + let resetPromise = Promise.resolve(); + await act(async () => { + resetPromise = result.current.reset(); + }); + expect(result.current.state.phase).toBe("resetting"); + + act(() => { + result.current.dismiss(); + }); + await act(async () => { + bulkDone.resolve([{ success: true, user_id: "u-1" }]); + await resetPromise; + }); + + expect(result.current.state.phase).toBe("idle"); + expect(gateway.refreshTeamData).not.toHaveBeenCalled(); + expect(toast.success).not.toHaveBeenCalled(); + }); + + it("does not return to prompting when the team save fails after a dismiss", async () => { + const gateway = buildGateway(); + const saveDone = Promise.withResolvers(); + gateway.saveTeam.mockReturnValueOnce(saveDone.promise); + const { result } = renderReset(gateway); + + await act(async () => { + result.current.prompt(pendingFor(["u-1"])); + }); + let resetPromise = Promise.resolve(); + await act(async () => { + resetPromise = result.current.reset(); + }); + + act(() => { + result.current.dismiss(); + }); + await act(async () => { + saveDone.reject(new Error("team update failed")); + await resetPromise; + }); + + expect(result.current.state.phase).toBe("idle"); + expect(gateway.resetMemberBudgets).not.toHaveBeenCalled(); + expect(gateway.refreshTeamData).not.toHaveBeenCalled(); + }); + + it("stays idle without a refresh when the reset fails after a dismiss", async () => { + const gateway = buildGateway(); + const bulkDone = Promise.withResolvers<{ success: boolean; user_id: string }[]>(); + gateway.resetMemberBudgets.mockReturnValueOnce(bulkDone.promise); + const { result } = renderReset(gateway); + + await act(async () => { + result.current.prompt(pendingFor(["u-1"])); + }); + let resetPromise = Promise.resolve(); + await act(async () => { + resetPromise = result.current.reset(); + }); + + act(() => { + result.current.dismiss(); + }); + await act(async () => { + bulkDone.reject(new Error("bulk failed")); + await resetPromise; + }); + + expect(result.current.state.phase).toBe("idle"); + expect(gateway.refreshTeamData).not.toHaveBeenCalled(); + expect(toast.error).not.toHaveBeenCalled(); + expect(toast.fromError).not.toHaveBeenCalled(); + }); + + it("skips the refresh when dismissed while keep-custom is saving", async () => { + const gateway = buildGateway(); + const saveDone = Promise.withResolvers(); + gateway.saveTeam.mockReturnValueOnce(saveDone.promise); + const { result } = renderReset(gateway); + + await act(async () => { + result.current.prompt(pendingFor(["u-1"])); + }); + let keepPromise = Promise.resolve(); + await act(async () => { + keepPromise = result.current.keepCustom(); + }); + + act(() => { + result.current.dismiss(); + }); + await act(async () => { + saveDone.resolve(); + await keepPromise; + }); + + expect(result.current.state.phase).toBe("idle"); + expect(gateway.refreshTeamData).not.toHaveBeenCalled(); + expect(toast.success).not.toHaveBeenCalled(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/team/useMemberBudgetReset.ts b/ui/litellm-dashboard/src/components/team/useMemberBudgetReset.ts new file mode 100644 index 00000000000..f0ff832b34a --- /dev/null +++ b/ui/litellm-dashboard/src/components/team/useMemberBudgetReset.ts @@ -0,0 +1,134 @@ +import { useCallback, useRef, useState } from "react"; +import type { components } from "@/lib/http/schema"; +import { toast } from "@/lib/toast"; +import { + chunk, + MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES, + pluralize, + type MemberBudgetResetPending, + type TeamUpdatePayload, +} from "./memberBudgetReset"; + +export type MemberBudgetBulkResult = components["schemas"]["TeamMemberBudgetUpdateResult"]; + +export type MemberBudgetResetState = + | { phase: "idle" } + | { phase: "prompting"; pending: MemberBudgetResetPending } + | { phase: "resetting"; pending: MemberBudgetResetPending; attempted: number } + | { phase: "resetFailed"; pending: MemberBudgetResetPending; attempted: number }; + +export interface MemberBudgetResetGateway { + saveTeam: (updateData: TeamUpdatePayload) => Promise; + resetMemberBudgets: (teamId: string, userIds: readonly string[]) => Promise; + refreshTeamData: () => Promise; +} + +type ResetRun = + | { ok: true; attempted: number; results: readonly MemberBudgetBulkResult[] } + | { ok: false; attempted: number; error: unknown }; + +export const useMemberBudgetReset = (gateway: MemberBudgetResetGateway) => { + const [state, setState] = useState({ phase: "idle" }); + const activeRun = useRef(null); + const isCurrent = (run: object) => activeRun.current === run; + + const runReset = async (pending: MemberBudgetResetPending, attempted: number, run: object) => { + const { resetMemberBudgets, refreshTeamData } = gateway; + + const runChunks = ( + chunks: readonly (readonly string[])[], + attemptedSoFar: number, + results: readonly MemberBudgetBulkResult[], + ): Promise => { + const [ids, ...rest] = chunks; + if (ids === undefined) return Promise.resolve({ ok: true, attempted: attemptedSoFar, results }); + return resetMemberBudgets(pending.teamId, ids).then( + (batch) => runChunks(rest, attemptedSoFar + ids.length, [...results, ...batch]), + (error: unknown) => ({ ok: false as const, attempted: attemptedSoFar, error }), + ); + }; + + const outcome = await runChunks( + chunk(pending.userIds.slice(attempted), MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES), + attempted, + [], + ); + + if (!isCurrent(run)) return; + if (!outcome.ok) { + console.error("Error resetting member budgets:", outcome.error); + const total = pending.userIds.length; + if (outcome.attempted > 0) { + toast.error( + `Reset ${outcome.attempted} of ${total} member ${pluralize(total, "budget", "budgets")}; the rest could not be reset`, + ); + } else { + toast.fromError("Team updated, but member budgets could not be reset"); + } + setState({ phase: "resetFailed", pending, attempted: outcome.attempted }); + await refreshTeamData(); + return; + } + const failed = outcome.results.filter((r) => !r.success); + if (failed.length > 0) { + toast.error( + `Team updated, but ${failed.length} member ${pluralize(failed.length, "budget", "budgets")} could not be reset`, + ); + } else { + toast.success( + `Reset ${outcome.attempted} member ${pluralize(outcome.attempted, "budget", "budgets")} to the team default`, + ); + } + setState({ phase: "idle" }); + await refreshTeamData(); + }; + + const prompt = (pending: MemberBudgetResetPending) => setState({ phase: "prompting", pending }); + + const reset = async () => { + if (state.phase !== "prompting") return; + const { pending } = state; + const run = {}; + activeRun.current = run; + setState({ phase: "resetting", pending, attempted: 0 }); + try { + await gateway.saveTeam(pending.updateData); + } catch (error) { + console.error("Error updating team:", error); + if (isCurrent(run)) setState({ phase: "prompting", pending }); + return; + } + await runReset(pending, 0, run); + }; + + const retry = async () => { + if (state.phase !== "resetFailed") return; + const { pending, attempted } = state; + const run = {}; + activeRun.current = run; + setState({ phase: "resetting", pending, attempted }); + await runReset(pending, attempted, run); + }; + + const keepCustom = async () => { + if (state.phase !== "prompting") return; + const { pending } = state; + const run = {}; + activeRun.current = run; + setState({ phase: "idle" }); + try { + await gateway.saveTeam(pending.updateData); + if (isCurrent(run)) toast.success("Team settings updated successfully"); + } catch (error) { + console.error("Error updating team:", error); + } + if (isCurrent(run)) await gateway.refreshTeamData(); + }; + + const dismiss = useCallback(() => { + activeRun.current = null; + setState({ phase: "idle" }); + }, []); + + return { state, prompt, reset, retry, keepCustom, dismiss }; +};