Merge pull request #36960 from BerriAI/litellm_cost_tracking_removal_pending_state

fix(ui): keep the cost tracking removal confirmation open until it settles
This commit is contained in:
yuneng-jiang 2026-08-14 11:15:53 -07:00 • committed by GitHub
commit 387ad0bb4f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 40 additions and 13 deletions

View file

@ -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 });

View file

@ -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>