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..03cef2a66b8 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,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(); + 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 }); 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 = ({ userID, use const [fixedAmountValue, setFixedAmountValue] = useState(""); const [models, setModels] = useState([]); const [pendingRemoval, setPendingRemoval] = useState(null); + const [isRemoving, setIsRemoving] = useState(false); const [form] = Form.useForm(); const [marginForm] = Form.useForm(); @@ -131,14 +131,19 @@ const CostTrackingSettings: React.FC = ({ 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 = ({ userID, use {pendingRemoval && ( - !open && setPendingRemoval(null)}> + !open && !isRemoving && setPendingRemoval(null)}> {REMOVAL_COPY[pendingRemoval.kind].title} @@ -321,10 +326,10 @@ const CostTrackingSettings: React.FC = ({ userID, use - Cancel - - Remove - + Cancel +