Merge pull request #36960 from BerriAI/litellm_cost_tracking_removal_pending_state

fix(ui): keep the cost tracking removal confirmation open until it settles

Author: yuneng-berri

ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/cost_tracking_settings.test.tsx

@@ -1,6 +1,6 @@
 import React from "react";
 import { describe, it, expect, vi, beforeEach } from "vitest";
-import { screen } from "@testing-library/react";
+import { act, screen, waitFor } from "@testing-library/react";
 import userEvent from "@testing-library/user-event";
 import { renderWithProviders } from "../../../../../tests/test-utils";
 import CostTrackingSettings from "./cost_tracking_settings";
@@ -197,6 +197,28 @@ describe("CostTrackingSettings", () => {
       expect(screen.queryByRole("button", { name: "Remove" })).not.toBeInTheDocument();
     });
 
+    it("should hold the confirmation open while the removal is still in flight", async () => {
+      mockDiscountConfig.mockReturnValue({ openai: 0.05 });
+      const { promise, resolve: settleRemoval } = Promise.withResolvers<void>();
+      mockRemoveDiscount.mockReturnValue(promise);
+
+      const user = await expandAndRemove("Provider Discounts", "Remove discount for openai");
+      await user.click(await screen.findByRole("button", { name: "Remove" }));
+
+      const removing = await screen.findByRole("button", { name: "Removing…" });
+      expect(removing).toBeDisabled();
+      expect(screen.getByRole("button", { name: "Cancel" })).toBeDisabled();
+
+      await act(async () => {
+        settleRemoval();
+      });
+
+      await waitFor(() => {
+        expect(screen.queryByRole("alertdialog")).not.toBeInTheDocument();
+      });
+      expect(mockRemoveDiscount).toHaveBeenCalledWith("openai");
+    });
+
     it("should remove the margin once removal is confirmed", async () => {
       mockMarginConfig.mockReturnValue({ openai: 0.1 });
 

ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/cost_tracking_settings.tsx

@@ -3,7 +3,6 @@ import { ChevronDown } from "lucide-react";
 import { Modal, Form } from "antd";
 import {
   AlertDialog,
-  AlertDialogAction,
   AlertDialogCancel,
   AlertDialogContent,
   AlertDialogDescription,
@@ -66,6 +65,7 @@ const CostTrackingSettings: React.FC<CostTrackingSettingsProps> = ({ userID, use
   const [fixedAmountValue, setFixedAmountValue] = useState<string>("");
   const [models, setModels] = useState<string[]>([]);
   const [pendingRemoval, setPendingRemoval] = useState<PendingRemoval | null>(null);
+  const [isRemoving, setIsRemoving] = useState(false);
   const [form] = Form.useForm();
   const [marginForm] = Form.useForm();
 
@@ -131,14 +131,19 @@ const CostTrackingSettings: React.FC<CostTrackingSettingsProps> = ({ userID, use
     setPendingRemoval({ kind: "discount", provider, displayName: providerDisplayName });
   };
 
-  const handleConfirmRemoval = () => {
+  const handleConfirmRemoval = async () => {
     if (!pendingRemoval) return;
-    if (pendingRemoval.kind === "discount") {
-      removeProvider(pendingRemoval.provider);
-    } else {
-      removeMargin(pendingRemoval.provider);
+    setIsRemoving(true);
+    try {
+      if (pendingRemoval.kind === "discount") {
+        await removeProvider(pendingRemoval.provider);
+      } else {
+        await removeMargin(pendingRemoval.provider);
+      }
+    } finally {
+      setIsRemoving(false);
+      setPendingRemoval(null);
     }
-    setPendingRemoval(null);
   };
 
   const handleAddMargin = async () => {
@@ -311,7 +316,7 @@ const CostTrackingSettings: React.FC<CostTrackingSettingsProps> = ({ userID, use
       </div>
 
       {pendingRemoval && (
-        <AlertDialog open onOpenChange={(open) => !open && setPendingRemoval(null)}>
+        <AlertDialog open onOpenChange={(open) => !open && !isRemoving && setPendingRemoval(null)}>
           <AlertDialogContent>
             <AlertDialogHeader>
               <AlertDialogTitle>{REMOVAL_COPY[pendingRemoval.kind].title}</AlertDialogTitle>
@@ -321,10 +326,10 @@ const CostTrackingSettings: React.FC<CostTrackingSettingsProps> = ({ userID, use
               </AlertDialogDescription>
             </AlertDialogHeader>
             <AlertDialogFooter>
-              <AlertDialogCancel>Cancel</AlertDialogCancel>
-              <AlertDialogAction variant="destructive" onClick={handleConfirmRemoval}>
-                Remove
-              </AlertDialogAction>
+              <AlertDialogCancel disabled={isRemoving}>Cancel</AlertDialogCancel>
+              <Button variant="destructive" onClick={handleConfirmRemoval} disabled={isRemoving}>
+                {isRemoving ? "Removing…" : "Remove"}
+              </Button>
             </AlertDialogFooter>
           </AlertDialogContent>
         </AlertDialog>