Merge 4e90a80275 into d8c0e2c715

ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.tsx

@@ -41,6 +41,7 @@ interface ViewUserDashboardProps {
   userRole: string | null;
   userID: string | null;
   teams: any[] | null;
+  premiumUser?: boolean;
   orgAdminOrgIds?: Array<{ organization_id: string; organization_alias: string }> | null;
 }
 
@@ -56,6 +57,7 @@ const ViewUserDashboard: React.FC<ViewUserDashboardProps> = ({
   userRole,
   userID,
   teams,
+  premiumUser = false,
   orgAdminOrgIds,
 }) => {
   const isProxyAdmin = userRole ? isProxyAdminRole(userRole) : false;
@@ -329,7 +331,12 @@ const ViewUserDashboard: React.FC<ViewUserDashboardProps> = ({
           {!userListQuery.isLoading && userID && accessToken && (
             <>
               {isProxyAdmin && (
-                <CreateUserButton userID={userID} accessToken={accessToken} possibleUIRoles={possibleUIRoles} />
+                <CreateUserButton
+                  userID={userID}
+                  accessToken={accessToken}
+                  possibleUIRoles={possibleUIRoles}
+                  premiumUser={premiumUser}
+                />
               )}
 
               {isProxyAdmin && (

ui/litellm-dashboard/src/app/(dashboard)/users/page.tsx

@@ -5,7 +5,7 @@ import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
 import { useTeams } from "@/app/(dashboard)/hooks/teams/useTeams";
 
 export default function UsersPage() {
-  const { accessToken, token, userRole, userId } = useAuthorized();
+  const { accessToken, token, userRole, userId, premiumUser } = useAuthorized();
   const { data: teams } = useTeams();
   return (
     <ViewUserDashboard
@@ -14,6 +14,7 @@ export default function UsersPage() {
       token={token}
       teams={teams ?? null}
       accessToken={accessToken}
+      premiumUser={premiumUser === true}
     />
   );
 }

ui/litellm-dashboard/src/components/CreateUserButton.test.tsx

@@ -1,4 +1,4 @@
-import { render, screen, waitFor, within } from "@testing-library/react";
+import { fireEvent, render, screen, waitFor, within } from "@testing-library/react";
 import userEvent, { PointerEventsCheckLevel } from "@testing-library/user-event";
 import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
 import { describe, it, expect, vi, beforeEach } from "vitest";
@@ -6,6 +6,8 @@ import { CreateUserButton } from "./CreateUserButton";
 import * as networking from "./networking";
 import { toast } from "@/lib/toast";
 import { expectControlBesideLabel } from "../../tests/fieldOrientation";
+import { chooseSelectOption } from "../../tests/test-utils";
+import { MODEL_MAX_BUDGET_PREMIUM_HINT } from "./key_team_helpers/ModelMaxBudgetEditor";
 
 vi.mock("./networking", () => ({
   userCreateCall: vi.fn(),
@@ -751,4 +753,150 @@ describe("CreateUserButton", () => {
       expect(submittedPayload().models).toEqual(["no-default-models"]);
     });
   });
+
+  describe("per-model budgets", () => {
+    const openStandaloneModal = async (user: ReturnType<typeof userEvent.setup>) => {
+      expect(await screen.findByRole("button", { name: /\+ invite user/i })).toBeInTheDocument();
+      await user.click(screen.getByRole("button", { name: /\+ invite user/i }));
+      return screen.getByRole("dialog", { name: /invite user/i });
+    };
+
+    it("locks the per-model budget editor when the proxy has no enterprise license", async () => {
+      const user = userEvent.setup({ pointerEventsCheck: PointerEventsCheckLevel.Never });
+      renderWithProviders(<CreateUserButton {...defaultProps} premiumUser={false} />);
+      const dialog = await openStandaloneModal(user);
+
+      expect(within(dialog).getByRole("button", { name: /Add Model Budget/i })).toBeDisabled();
+      expect(within(dialog).getByText(MODEL_MAX_BUDGET_PREMIUM_HINT)).toBeInTheDocument();
+    });
+
+    it("sends model_max_budget with budget_limit and time_period when a budget is entered", async () => {
+      const user = userEvent.setup({ pointerEventsCheck: PointerEventsCheckLevel.Never });
+      vi.mocked(networking.modelAvailableCall).mockResolvedValue({ data: [{ id: "gpt-4" }] } as any);
+      mockUserCreateCall.mockResolvedValue({ data: { user_id: "budget-user" } });
+      mockInvitationCreateCall.mockResolvedValue({
+        id: "inv-budget",
+        user_id: "budget-user",
+        has_user_setup_sso: false,
+      } as any);
+
+      renderWithProviders(
+        <CreateUserButton
+          {...defaultProps}
+          premiumUser
+          possibleUIRoles={{ proxy_user: { ui_label: "User", description: "" } }}
+        />,
+      );
+      const dialog = await openStandaloneModal(user);
+
+      await user.type(within(dialog).getByLabelText(/user email/i), "budget@example.com");
+      await user.click(within(dialog).getByRole("combobox", { name: /global proxy role/i }));
+      await user.click(screen.getByText("User"));
+      await user.click(within(dialog).getByRole("button", { name: /Add Model Budget/i }));
+      await chooseSelectOption(user, within(dialog).getByPlaceholderText("Select model"), "gpt-4");
+      fireEvent.change(within(dialog).getByPlaceholderText("Max spend ($)"), { target: { value: "3" } });
+      await user.click(within(dialog).getByRole("button", { name: /invite user/i }));
+
+      await waitFor(() => {
+        expect(mockUserCreateCall).toHaveBeenCalled();
+      });
+      expect(mockUserCreateCall.mock.calls[0][2].model_max_budget).toStrictEqual({
+        "gpt-4": { budget_limit: 3, time_period: "30d" },
+      });
+    });
+
+    it("omits model_max_budget from the payload when no budget is entered", async () => {
+      const user = userEvent.setup({ pointerEventsCheck: PointerEventsCheckLevel.Never });
+      mockUserCreateCall.mockResolvedValue({ data: { user_id: "no-budget-user" } });
+      mockInvitationCreateCall.mockResolvedValue({
+        id: "inv-nobudget",
+        user_id: "no-budget-user",
+        has_user_setup_sso: false,
+      } as any);
+
+      renderWithProviders(
+        <CreateUserButton
+          {...defaultProps}
+          premiumUser
+          possibleUIRoles={{ proxy_user: { ui_label: "User", description: "" } }}
+        />,
+      );
+      const dialog = await openStandaloneModal(user);
+
+      await user.type(within(dialog).getByLabelText(/user email/i), "nobudget@example.com");
+      await user.click(within(dialog).getByRole("combobox", { name: /global proxy role/i }));
+      await user.click(screen.getByText("User"));
+      await user.click(within(dialog).getByRole("button", { name: /invite user/i }));
+
+      await waitFor(() => {
+        expect(mockUserCreateCall).toHaveBeenCalled();
+      });
+      expect(mockUserCreateCall.mock.calls[0][2]).not.toHaveProperty("model_max_budget");
+    });
+
+    it("clears the per-model budget rows after a user is created in standalone mode", async () => {
+      const user = userEvent.setup({ pointerEventsCheck: PointerEventsCheckLevel.Never });
+      vi.mocked(networking.modelAvailableCall).mockResolvedValue({ data: [{ id: "gpt-4" }] } as any);
+      mockUserCreateCall.mockResolvedValue({ data: { user_id: "budget-user" } });
+      mockInvitationCreateCall.mockResolvedValue({
+        id: "inv-budget",
+        user_id: "budget-user",
+        has_user_setup_sso: false,
+      } as any);
+
+      renderWithProviders(
+        <CreateUserButton
+          {...defaultProps}
+          premiumUser
+          possibleUIRoles={{ proxy_user: { ui_label: "User", description: "" } }}
+        />,
+      );
+      const dialog = await openStandaloneModal(user);
+
+      await user.type(within(dialog).getByLabelText(/user email/i), "budget@example.com");
+      await user.click(within(dialog).getByRole("combobox", { name: /global proxy role/i }));
+      await user.click(screen.getByText("User"));
+      await user.click(within(dialog).getByRole("button", { name: /Add Model Budget/i }));
+      await chooseSelectOption(user, within(dialog).getByPlaceholderText("Select model"), "gpt-4");
+      fireEvent.change(within(dialog).getByPlaceholderText("Max spend ($)"), { target: { value: "3" } });
+      await user.click(within(dialog).getByRole("button", { name: /invite user/i }));
+
+      await waitFor(() => {
+        expect(mockUserCreateCall).toHaveBeenCalled();
+      });
+      expect(within(dialog).queryByPlaceholderText("Max spend ($)")).not.toBeInTheDocument();
+      expect(within(dialog).getByRole("button", { name: /Add Model Budget/i, hidden: true })).toBeInTheDocument();
+    });
+
+    it("clears the per-model budget rows after a user is created in embedded mode", async () => {
+      const user = userEvent.setup({ pointerEventsCheck: PointerEventsCheckLevel.Never });
+      const onUserCreated = vi.fn();
+      vi.mocked(networking.modelAvailableCall).mockResolvedValue({ data: [{ id: "gpt-4" }] } as any);
+      mockUserCreateCall.mockResolvedValue({ data: { user_id: "embedded-budget-user" } });
+
+      renderWithProviders(
+        <CreateUserButton
+          {...defaultProps}
+          premiumUser
+          onUserCreated={onUserCreated}
+          possibleUIRoles={{ proxy_user: { ui_label: "User", description: "" } }}
+          isEmbedded
+        />,
+      );
+
+      await user.type(screen.getByLabelText(/user email/i), "embedded-budget@example.com");
+      await user.click(screen.getByRole("combobox", { name: /user role/i }));
+      await user.click(screen.getByText("User"));
+      await user.click(screen.getByRole("button", { name: /Add Model Budget/i }));
+      await chooseSelectOption(user, screen.getByPlaceholderText("Select model"), "gpt-4");
+      fireEvent.change(screen.getByPlaceholderText("Max spend ($)"), { target: { value: "3" } });
+      await user.click(screen.getByRole("button", { name: /create user/i }));
+
+      await waitFor(() => {
+        expect(mockUserCreateCall).toHaveBeenCalled();
+      });
+      expect(screen.queryByPlaceholderText("Max spend ($)")).not.toBeInTheDocument();
+      expect(screen.getByRole("button", { name: /Add Model Budget/i })).toBeInTheDocument();
+    });
+  });
 });

ui/litellm-dashboard/src/components/CreateUserButton.tsx

@@ -17,6 +17,7 @@ import React, { useEffect, useState } from "react";
 import { useForm } from "react-hook-form";
 import TeamDropdown from "./common_components/team_dropdown";
 import { getModelDisplayName } from "./key_team_helpers/fetch_available_models_team_key";
+import { ModelMaxBudget, ModelMaxBudgetField } from "./key_team_helpers/ModelMaxBudgetEditor";
 import { toast } from "@/lib/toast";
 import {
   getProxyBaseUrl,
@@ -46,6 +47,7 @@ interface CreateuserProps {
   possibleUIRoles: null | Record<string, Record<string, string>>;
   onUserCreated?: (userId: string) => void;
   isEmbedded?: boolean;
+  premiumUser?: boolean;
 }
 
 // Define an interface for the UI settings
@@ -64,6 +66,7 @@ interface CreateUserFormValues {
   metadata?: string;
   send_invite_email: boolean;
   models?: string[];
+  model_max_budget?: ModelMaxBudget;
 }
 
 const EMBEDDED_DEFAULTS: CreateUserFormValues = {
@@ -153,6 +156,7 @@ export const CreateUserButton: React.FC<CreateuserProps> = ({
   possibleUIRoles,
   onUserCreated,
   isEmbedded = false,
+  premiumUser = false,
 }) => {
   const queryClient = useQueryClient();
   const [uiSettings, setUISettings] = useState<UISettings | null>(null);
@@ -161,6 +165,8 @@ export const CreateUserButton: React.FC<CreateuserProps> = ({
   const [isModalVisible, setIsModalVisible] = useState(false);
   const [apiuser, setApiuser] = useState<boolean>(false);
   const [userModels, setUserModels] = useState<string[]>([]);
+  const [modelMaxBudget, setModelMaxBudget] = useState<ModelMaxBudget>({});
+  const [modelMaxBudgetKey, setModelMaxBudgetKey] = useState(0);
   const [isPersonalKeyOpen, setIsPersonalKeyOpen] = useState(false);
   const [isInvitationLinkModalVisible, setIsInvitationLinkModalVisible] = useState(false);
   const [invitationLinkData, setInvitationLinkData] = useState<InvitationLink | null>(null);
@@ -193,10 +199,16 @@ export const CreateUserButton: React.FC<CreateuserProps> = ({
     fetchData();
   }, []);
 
+  const resetModelMaxBudget = () => {
+    setModelMaxBudget({});
+    setModelMaxBudgetKey((prev) => prev + 1);
+  };
+
   const handleCancel = () => {
     setIsModalVisible(false);
     setApiuser(false);
     form.reset(defaultValues);
+    resetModelMaxBudget();
   };
 
   const handleCreate = async (formValues: CreateUserFormValues) => {
@@ -205,7 +217,10 @@ export const CreateUserButton: React.FC<CreateuserProps> = ({
       if (!isEmbedded) {
         setIsModalVisible(true);
       }
-      const payload = buildCreatePayload(withoutCollapsedModels(formValues, isPersonalKeyOpen));
+      const payload = {
+        ...buildCreatePayload(withoutCollapsedModels(formValues, isPersonalKeyOpen)),
+        ...(Object.keys(modelMaxBudget).length > 0 && { model_max_budget: modelMaxBudget }),
+      };
       const response = await userCreateCall(accessToken, null, payload);
       await queryClient.invalidateQueries({ queryKey: ["userList"] });
       setApiuser(true);
@@ -214,6 +229,7 @@ export const CreateUserButton: React.FC<CreateuserProps> = ({
       if (onUserCreated && isEmbedded) {
         onUserCreated(user_id);
         form.reset(defaultValues);
+        resetModelMaxBudget();
         return;
       }
 
@@ -230,6 +246,7 @@ export const CreateUserButton: React.FC<CreateuserProps> = ({
 
       toast.success("API user Created");
       form.reset(defaultValues);
+      resetModelMaxBudget();
       localStorage.removeItem("userData" + userID);
     } catch (error: any) {
       const errorMessage = error.response?.data?.detail || error?.message || "Error creating the user";
@@ -277,6 +294,17 @@ export const CreateUserButton: React.FC<CreateuserProps> = ({
     </FormField>
   );
 
+  const modelMaxBudgetField = (
+    <ModelMaxBudgetField
+      key={`model-max-budget-${modelMaxBudgetKey}`}
+      premiumUser={premiumUser}
+      value={modelMaxBudget}
+      onChange={setModelMaxBudget}
+      availableModels={userModels}
+      hint="Cap this user's spend on individual models, each with its own reset window."
+    />
+  );
+
   const roleField = (label: React.ReactNode) => (
     <FormField control={form.control} name="user_role" label={label}>
       {({ id, value, onChange }) => (
@@ -312,6 +340,7 @@ export const CreateUserButton: React.FC<CreateuserProps> = ({
             {roleField("User Role")}
             {teamField}
             {metadataField}
+            {modelMaxBudgetField}
             {sendInviteEmailField}
           </FieldGroup>
           <div className="mt-4 text-right">
@@ -386,6 +415,7 @@ export const CreateUserButton: React.FC<CreateuserProps> = ({
                 </FormField>
 
                 {metadataField}
+                {modelMaxBudgetField}
                 {sendInviteEmailField}
 
                 <Collapsible open={isPersonalKeyOpen} onOpenChange={setIsPersonalKeyOpen}>

ui/litellm-dashboard/src/components/organisms/create_key_button.tsx

@@ -1847,6 +1847,7 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey, autoOp
               possibleUIRoles={possibleUIRoles}
               onUserCreated={handleUserCreated}
               isEmbedded={true}
+              premiumUser={premiumUser === true}
             />
           </DialogContent>
         </Dialog>