mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-26 01:12:21 +00:00
feat(ui): offer reset of custom member budgets when team default changes (#42835)
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
cbe342171c
commit
165dbc2243
11 changed files with 1623 additions and 9 deletions
|
|
@ -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(<ResetMemberBudgetsDialog state={promptingWith(3)} {...defaultHandlers()} />);
|
||||
|
||||
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(<ResetMemberBudgetsDialog state={promptingWith(1)} {...defaultHandlers()} />);
|
||||
|
||||
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(<ResetMemberBudgetsDialog state={promptingWith(2, 1000)} {...defaultHandlers()} />);
|
||||
|
||||
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(<ResetMemberBudgetsDialog state={promptingWith(2)} {...handlers} />);
|
||||
|
||||
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(<ResetMemberBudgetsDialog state={state} {...defaultHandlers()} />);
|
||||
|
||||
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(<ResetMemberBudgetsDialog state={state} {...handlers} />);
|
||||
|
||||
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(<ResetMemberBudgetsDialog state={{ phase: "idle" }} {...defaultHandlers()} />);
|
||||
|
||||
expect(screen.queryByRole("dialog")).not.toBeInTheDocument();
|
||||
});
|
||||
});
|
||||
|
|
@ -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 (
|
||||
<Dialog
|
||||
open={open}
|
||||
onOpenChange={(nextOpen) => {
|
||||
if (!nextOpen && !busy) onDismiss();
|
||||
}}
|
||||
>
|
||||
<DialogContent>
|
||||
<DialogHeader>
|
||||
<DialogTitle>Reset member budgets?</DialogTitle>
|
||||
</DialogHeader>
|
||||
<p className="text-sm text-muted-foreground">
|
||||
{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")}?
|
||||
</p>
|
||||
<DialogFooter>
|
||||
{failed ? (
|
||||
<>
|
||||
<Button variant="outline" onClick={onDismiss}>
|
||||
Cancel
|
||||
</Button>
|
||||
<Button onClick={onRetry}>Retry reset</Button>
|
||||
</>
|
||||
) : (
|
||||
<>
|
||||
<Button variant="ghost" onClick={onDismiss} disabled={busy}>
|
||||
Cancel
|
||||
</Button>
|
||||
<Button variant="outline" onClick={onKeep} disabled={busy}>
|
||||
Keep custom {pluralize(memberCount, "budget", "budgets")}
|
||||
</Button>
|
||||
<Button onClick={onReset} disabled={busy}>
|
||||
{pluralize(memberCount, "Reset", "Reset all")} to ${formatNumberWithCommas(newBudget, 2)}
|
||||
</Button>
|
||||
</>
|
||||
)}
|
||||
</DialogFooter>
|
||||
</DialogContent>
|
||||
</Dialog>
|
||||
);
|
||||
}
|
||||
|
|
@ -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<typeof import("@/lib/http/api")>();
|
||||
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;
|
||||
}) => (
|
||||
<>
|
||||
<button
|
||||
type="button"
|
||||
onClick={() =>
|
||||
onChange({ servers: [], accessGroups: value?.accessGroups ?? [], toolsets: value?.toolsets ?? [] })
|
||||
}
|
||||
>
|
||||
deselect all mcp servers
|
||||
</button>
|
||||
<button type="button" onClick={() => onChange({ servers: value?.servers ?? [], accessGroups: [], toolsets: [] })}>
|
||||
remove all access groups
|
||||
</button>
|
||||
</>
|
||||
),
|
||||
}));
|
||||
|
||||
vi.mock("@/components/team/TeamMemberTab", () => ({
|
||||
default: vi.fn(({ setIsAddMemberModalVisible }) => (
|
||||
<div>
|
||||
<button onClick={() => setIsAddMemberModalVisible(true)}>Add Member</button>
|
||||
</div>
|
||||
)),
|
||||
}));
|
||||
|
||||
vi.mock("@/components/common_components/user_search_modal", () => ({
|
||||
default: vi.fn(({ isVisible, onCancel, onSubmit }) =>
|
||||
isVisible ? (
|
||||
<div>
|
||||
<button onClick={onCancel}>Cancel</button>
|
||||
<button onClick={() => onSubmit({ user_email: "new@test.com", user_id: "new-user", role: "user" })}>
|
||||
Submit
|
||||
</button>
|
||||
</div>
|
||||
) : null,
|
||||
),
|
||||
}));
|
||||
|
||||
vi.mock("@/components/team/EditMembership", () => ({
|
||||
default: vi.fn(({ visible, onCancel, onSubmit }) =>
|
||||
visible ? (
|
||||
<div>
|
||||
<button onClick={onCancel}>Cancel</button>
|
||||
<button onClick={() => onSubmit({ user_email: "edit@test.com", user_id: "edit-user", role: "admin" })}>
|
||||
Submit
|
||||
</button>
|
||||
</div>
|
||||
) : null,
|
||||
),
|
||||
}));
|
||||
|
||||
vi.mock("@/components/common_components/DeleteResourceModal", () => ({
|
||||
default: vi.fn(({ isOpen, onCancel, onOk }) =>
|
||||
isOpen ? (
|
||||
<div>
|
||||
<button onClick={onCancel}>Cancel</button>
|
||||
<button onClick={onOk}>Confirm Delete</button>
|
||||
</div>
|
||||
) : null,
|
||||
),
|
||||
}));
|
||||
|
||||
vi.mock("@/components/team/member_permissions", () => ({
|
||||
default: vi.fn(() => <div>Member Permissions</div>),
|
||||
}));
|
||||
|
||||
vi.mock("@/components/common_components/ModelAliasManager", () => ({
|
||||
default: vi.fn(({ initialModelAliases, onAliasUpdate }) => (
|
||||
<div>
|
||||
<div data-testid="alias-editor-initial">{JSON.stringify(initialModelAliases)}</div>
|
||||
<button type="button" onClick={() => onAliasUpdate({ "gpt-4o": "gpt-4" })}>
|
||||
Set Alias
|
||||
</button>
|
||||
<button type="button" onClick={() => onAliasUpdate({})}>
|
||||
Clear Aliases
|
||||
</button>
|
||||
</div>
|
||||
)),
|
||||
}));
|
||||
|
||||
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 }) => (
|
||||
<button type="button" onClick={() => onChange?.((value ?? []).slice(1))}>
|
||||
remove first unified access group
|
||||
</button>
|
||||
),
|
||||
}));
|
||||
|
||||
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<typeof useAllProxyModels>);
|
||||
mockUseTeam.mockReturnValue({
|
||||
data: undefined,
|
||||
isLoading: false,
|
||||
} as unknown as ReturnType<typeof useTeam>);
|
||||
mockUseOrganization.mockReturnValue({
|
||||
data: undefined,
|
||||
isLoading: false,
|
||||
} as unknown as ReturnType<typeof useOrganization>);
|
||||
mockUseCurrentUser.mockReturnValue({
|
||||
data: { models: [] },
|
||||
isLoading: false,
|
||||
} as unknown as ReturnType<typeof useCurrentUser>);
|
||||
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<typeof useAccessGroups>);
|
||||
mockUseUISettings.mockReturnValue({
|
||||
data: { values: {} },
|
||||
isLoading: false,
|
||||
} as unknown as ReturnType<typeof useUISettings>);
|
||||
mockUseKeys.mockReturnValue({
|
||||
data: { keys: [], total_count: 0, current_page: 1, total_pages: 1 },
|
||||
isPending: false,
|
||||
isFetching: false,
|
||||
refetch: vi.fn(),
|
||||
} as unknown as ReturnType<typeof useKeys>);
|
||||
vi.mocked(useTeamMetadataSchema).mockReturnValue({
|
||||
data: [],
|
||||
isLoading: false,
|
||||
} as unknown as ReturnType<typeof useTeamMetadataSchema>);
|
||||
|
||||
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<typeof userEvent.setup>,
|
||||
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(<TeamInfoView {...props} />);
|
||||
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<typeof userEvent.setup>, 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(<TeamInfoView {...props} />);
|
||||
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(<TeamInfoView {...props} teamId="456" />);
|
||||
|
||||
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(<TeamInfoView {...props} />);
|
||||
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(<TeamInfoView {...props} teamId="456" />);
|
||||
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();
|
||||
});
|
||||
});
|
||||
|
|
@ -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<typeof import("@/lib/http/api")>();
|
||||
return { ...actual, fetchClient: { ...actual.fetchClient, POST: bulkUpdatePOST } };
|
||||
});
|
||||
|
||||
const can = vi.fn();
|
||||
vi.mock("@/app/(dashboard)/hooks/useCan", () => ({
|
||||
default: (...args: unknown[]) => can(...args),
|
||||
|
|
|
|||
|
|
@ -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<string> = new Set([
|
||||
"logging",
|
||||
|
|
@ -884,17 +888,41 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
|
|||
const persistTeamUpdate = async (token: string, updateData: Record<string, unknown>) => {
|
||||
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<TeamInfoProps> = ({
|
|||
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<TeamInfoProps> = ({
|
|||
}
|
||||
}
|
||||
|
||||
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<TeamInfoProps> = ({
|
|||
onOk={handleDeleteConfirm}
|
||||
confirmLoading={isDeleting}
|
||||
/>
|
||||
|
||||
<ResetMemberBudgetsDialog
|
||||
state={memberBudgetReset.state}
|
||||
onReset={memberBudgetReset.reset}
|
||||
onRetry={memberBudgetReset.retry}
|
||||
onKeep={memberBudgetReset.keepCustom}
|
||||
onDismiss={memberBudgetReset.dismiss}
|
||||
/>
|
||||
</div>
|
||||
);
|
||||
};
|
||||
|
|
|
|||
|
|
@ -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: [] }));
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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");
|
||||
});
|
||||
});
|
||||
|
|
@ -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<TeamMembership, "user_id" | "budget_source"> & {
|
||||
litellm_budget_table?: Pick<TeamMembership["litellm_budget_table"], "max_budget"> | 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 = <T>(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);
|
||||
|
|
@ -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<string, unknown>) => {}),
|
||||
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<void>();
|
||||
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<void>();
|
||||
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();
|
||||
});
|
||||
});
|
||||
134
ui/litellm-dashboard/src/components/team/useMemberBudgetReset.ts
Normal file
134
ui/litellm-dashboard/src/components/team/useMemberBudgetReset.ts
Normal file
|
|
@ -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<void>;
|
||||
resetMemberBudgets: (teamId: string, userIds: readonly string[]) => Promise<MemberBudgetBulkResult[]>;
|
||||
refreshTeamData: () => Promise<void>;
|
||||
}
|
||||
|
||||
type ResetRun =
|
||||
| { ok: true; attempted: number; results: readonly MemberBudgetBulkResult[] }
|
||||
| { ok: false; attempted: number; error: unknown };
|
||||
|
||||
export const useMemberBudgetReset = (gateway: MemberBudgetResetGateway) => {
|
||||
const [state, setState] = useState<MemberBudgetResetState>({ phase: "idle" });
|
||||
const activeRun = useRef<object | null>(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<ResetRun> => {
|
||||
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 };
|
||||
};
|
||||
Loading…
Add table
Reference in a new issue