From 4974290d3f393c434b954693bb47e8679f4a112f Mon Sep 17 00:00:00 2001
From: Yuneng Jiang <yuneng@berri.ai>
Date: Fri, 14 Aug 2026 10:33:27 -0700
Subject: [PATCH] fix(ui): keep the cost tracking removal confirmation open
 until it settles

The discount and margin removal confirmation used AlertDialogAction, which
renders AlertDialogPrimitive.Close and dismisses the dialog on click. The
dialog therefore disappeared while the removal request was still in flight,
leaving the admin with no sign that anything happened and free to fire a
duplicate removal.

Swap the confirm control for a plain destructive Button, track an isRemoving
pending state that disables Cancel and relabels Remove to "Removing...", and
clear the pending removal in a finally block once the request settles.
---
 .../cost_tracking_settings.test.tsx           | 28 +++++++++++++++++-
 .../_components/cost_tracking_settings.tsx    | 29 +++++++++++--------
 2 files changed, 44 insertions(+), 13 deletions(-)

diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/cost_tracking_settings.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/cost_tracking_settings.test.tsx
index c53b7b618b2..716654e914c 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/cost_tracking_settings.test.tsx
+++ b/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,32 @@ 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 });
+      let settleRemoval: () => void = () => {};
+      mockRemoveDiscount.mockReturnValue(
+        new Promise<void>((resolve) => {
+          settleRemoval = resolve;
+        }),
+      );
+
+      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 });
 
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/cost_tracking_settings.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/cost_tracking_settings.tsx
index ba2d830ae7b..7f86bad3028 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/cost_tracking_settings.tsx
+++ b/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>