From 4e90a80275cd040cb54c9f5b14b81ffc983aee67 Mon Sep 17 00:00:00 2001
From: jesus <jesus@berri.ai>
Date: Fri, 18 Sep 2026 21:50:11 +0000
Subject: [PATCH] fix(ui): remount the per-model budget editor when the create
 user form resets

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
 .../src/components/CreateUserButton.test.tsx  | 65 +++++++++++++++++++
 .../src/components/CreateUserButton.tsx       | 13 +++-
 2 files changed, 75 insertions(+), 3 deletions(-)

diff --git a/ui/litellm-dashboard/src/components/CreateUserButton.test.tsx b/ui/litellm-dashboard/src/components/CreateUserButton.test.tsx
index 9555dd28e1e..6ab857e6b61 100644
--- a/ui/litellm-dashboard/src/components/CreateUserButton.test.tsx
+++ b/ui/litellm-dashboard/src/components/CreateUserButton.test.tsx
@@ -833,5 +833,70 @@ describe("CreateUserButton", () => {
       });
       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();
+    });
   });
 });
diff --git a/ui/litellm-dashboard/src/components/CreateUserButton.tsx b/ui/litellm-dashboard/src/components/CreateUserButton.tsx
index 8f2ccc1a2b6..88b96e8b540 100644
--- a/ui/litellm-dashboard/src/components/CreateUserButton.tsx
+++ b/ui/litellm-dashboard/src/components/CreateUserButton.tsx
@@ -166,6 +166,7 @@ export const CreateUserButton: React.FC<CreateuserProps> = ({
   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);
@@ -198,11 +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);
-    setModelMaxBudget({});
+    resetModelMaxBudget();
   };
 
   const handleCreate = async (formValues: CreateUserFormValues) => {
@@ -223,7 +229,7 @@ export const CreateUserButton: React.FC<CreateuserProps> = ({
       if (onUserCreated && isEmbedded) {
         onUserCreated(user_id);
         form.reset(defaultValues);
-        setModelMaxBudget({});
+        resetModelMaxBudget();
         return;
       }
 
@@ -240,7 +246,7 @@ export const CreateUserButton: React.FC<CreateuserProps> = ({
 
       toast.success("API user Created");
       form.reset(defaultValues);
-      setModelMaxBudget({});
+      resetModelMaxBudget();
       localStorage.removeItem("userData" + userID);
     } catch (error: any) {
       const errorMessage = error.response?.data?.detail || error?.message || "Error creating the user";
@@ -290,6 +296,7 @@ export const CreateUserButton: React.FC<CreateuserProps> = ({
 
   const modelMaxBudgetField = (
     <ModelMaxBudgetField
+      key={`model-max-budget-${modelMaxBudgetKey}`}
       premiumUser={premiumUser}
       value={modelMaxBudget}
       onChange={setModelMaxBudget}