fix(ui): refactor budget page to React Query hooks and fix crashes

- Migrate budget CRUD from manual state to React Query hooks (useBudgets, useCreateBudget, useUpdateBudget, useDeleteBudget)
- Fix crash when budget list contains null entries by filtering in query hook
- Fix max_budget type from string to number to match DB schema (double precision)
- Disable budget_id field in edit modal to prevent accidental changes
- Use budget_id as React key instead of array index
- Update tests to mock hooks instead of networking functions

litellm/proxy/management_endpoints/budget_management_endpoints.py

@@ -15,6 +15,7 @@ All /budget management endpoints
 from datetime import timedelta
 
 from fastapi import APIRouter, Depends, HTTPException
+from prisma.errors import UniqueViolationError
 
 from litellm.litellm_core_utils.duration_parser import duration_in_seconds
 from litellm.proxy._types import *
@@ -90,13 +91,21 @@ async def new_budget(
 
     budget_obj_json = budget_obj.model_dump(exclude_none=True)
     budget_obj_jsonified = jsonify_object(budget_obj_json)  # json dump any dictionaries
-    response = await prisma_client.db.litellm_budgettable.create(
-        data={
-            **budget_obj_jsonified,  # type: ignore
-            "created_by": user_api_key_dict.user_id or litellm_proxy_admin_name,
-            "updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name,
-        }  # type: ignore
-    )
+    try:
+        response = await prisma_client.db.litellm_budgettable.create(
+            data={
+                **budget_obj_jsonified,  # type: ignore
+                "created_by": user_api_key_dict.user_id or litellm_proxy_admin_name,
+                "updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name,
+            }  # type: ignore
+        )
+    except UniqueViolationError:
+        raise HTTPException(
+            status_code=400,
+            detail={
+                "error": f"Budget with id '{budget_obj.budget_id}' already exists."
+            },
+        )
 
     return response
 

ui/litellm-dashboard/src/app/(dashboard)/hooks/budgets/useBudgets.ts

@@ -0,0 +1,70 @@
+import { useQuery, useMutation, useQueryClient, UseQueryResult } from "@tanstack/react-query";
+import { createQueryKeys } from "../common/queryKeysFactory";
+import { getBudgetList, budgetCreateCall, budgetUpdateCall, budgetDeleteCall } from "@/components/networking";
+import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
+import { budgetItem } from "@/components/budgets/budget_panel";
+
+export const budgetKeys = createQueryKeys("budgets");
+
+export const useBudgets = (): UseQueryResult<budgetItem[]> => {
+  const { accessToken } = useAuthorized();
+  return useQuery<budgetItem[]>({
+    queryKey: budgetKeys.list({}),
+    queryFn: async () => {
+      const data = await getBudgetList(accessToken!);
+      return (data ?? []).filter((item: budgetItem | null): item is budgetItem => item != null);
+    },
+    enabled: Boolean(accessToken),
+  });
+};
+
+export const useCreateBudget = () => {
+  const { accessToken } = useAuthorized();
+  const queryClient = useQueryClient();
+
+  return useMutation<unknown, Error, Record<string, any>>({
+    mutationFn: async (formValues) => {
+      if (!accessToken) {
+        throw new Error("Access token is required");
+      }
+      return budgetCreateCall(accessToken, formValues);
+    },
+    onSuccess: () => {
+      queryClient.invalidateQueries({ queryKey: budgetKeys.all });
+    },
+  });
+};
+
+export const useUpdateBudget = () => {
+  const { accessToken } = useAuthorized();
+  const queryClient = useQueryClient();
+
+  return useMutation<unknown, Error, Record<string, any>>({
+    mutationFn: async (formValues) => {
+      if (!accessToken) {
+        throw new Error("Access token is required");
+      }
+      return budgetUpdateCall(accessToken, formValues);
+    },
+    onSuccess: () => {
+      queryClient.invalidateQueries({ queryKey: budgetKeys.all });
+    },
+  });
+};
+
+export const useDeleteBudget = () => {
+  const { accessToken } = useAuthorized();
+  const queryClient = useQueryClient();
+
+  return useMutation<unknown, Error, string>({
+    mutationFn: async (budgetId) => {
+      if (!accessToken) {
+        throw new Error("Access token is required");
+      }
+      return budgetDeleteCall(accessToken, budgetId);
+    },
+    onSuccess: () => {
+      queryClient.invalidateQueries({ queryKey: budgetKeys.all });
+    },
+  });
+};

ui/litellm-dashboard/src/components/budgets/budget_modal.tsx

@@ -1,17 +1,17 @@
 import React from "react";
 import { TextInput, Accordion, AccordionHeader, AccordionBody } from "@tremor/react";
 import { Button as Button2, Modal, Form, InputNumber, Select } from "antd";
-import { budgetCreateCall } from "../networking";
+import { useCreateBudget } from "@/app/(dashboard)/hooks/budgets/useBudgets";
 import NotificationsManager from "../molecules/notifications_manager";
 
 interface BudgetModalProps {
   isModalVisible: boolean;
-  accessToken: string | null;
   setIsModalVisible: React.Dispatch<React.SetStateAction<boolean>>;
-  setBudgetList: React.Dispatch<React.SetStateAction<any[]>>;
 }
-const BudgetModal: React.FC<BudgetModalProps> = ({ isModalVisible, accessToken, setIsModalVisible, setBudgetList }) => {
+const BudgetModal: React.FC<BudgetModalProps> = ({ isModalVisible, setIsModalVisible }) => {
   const [form] = Form.useForm();
+  const createBudget = useCreateBudget();
+
   const handleOk = () => {
     setIsModalVisible(false);
     form.resetFields();
@@ -23,20 +23,15 @@ const BudgetModal: React.FC<BudgetModalProps> = ({ isModalVisible, accessToken,
   };
 
   const handleCreate = async (formValues: Record<string, any>) => {
-    if (accessToken == null || accessToken == undefined) {
-      return;
-    }
     try {
       NotificationsManager.info("Making API Call");
-      // setIsModalVisible(true);
-      const response = await budgetCreateCall(accessToken, formValues);
-      console.log("key create Response:", response);
-      setBudgetList((prevData) => (prevData ? [...prevData, response] : [response])); // Check if prevData is null
+      await createBudget.mutateAsync(formValues);
       NotificationsManager.success("Budget Created");
       form.resetFields();
+      setIsModalVisible(false);
     } catch (error) {
-      console.error("Error creating the key:", error);
-      NotificationsManager.fromBackend(`Error creating the key: ${error}`);
+      console.error("Error creating the budget:", error);
+      NotificationsManager.fromBackend(`Error creating the budget: ${error}`);
     }
   };
 

ui/litellm-dashboard/src/components/budgets/budget_panel.test.tsx

@@ -1,31 +1,50 @@
-import * as networking from "../networking";
 import { fireEvent, render, waitFor, screen } from "@testing-library/react";
 import { act } from "@testing-library/react";
+import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
 import { afterEach, describe, expect, it, vi } from "vitest";
 import BudgetPanel from "./budget_panel";
 
-vi.mock("../networking", () => ({
-  getBudgetList: vi.fn(),
-  budgetDeleteCall: vi.fn(),
+const mockBudgets = [
+  {
+    budget_id: "budget-1",
+    max_budget: 100,
+    rpm_limit: 10,
+    tpm_limit: 1000,
+    updated_at: "2024-01-01T00:00:00Z",
+  },
+];
+
+vi.mock("@/app/(dashboard)/hooks/budgets/useBudgets", () => ({
+  useBudgets: vi.fn().mockReturnValue({ data: [], isLoading: false }),
+  useDeleteBudget: vi.fn().mockReturnValue({ mutateAsync: vi.fn(), isPending: false }),
+  useCreateBudget: vi.fn().mockReturnValue({ mutateAsync: vi.fn() }),
+  useUpdateBudget: vi.fn().mockReturnValue({ mutateAsync: vi.fn() }),
 }));
 
+import { useBudgets, useDeleteBudget, useCreateBudget, useUpdateBudget } from "@/app/(dashboard)/hooks/budgets/useBudgets";
+
+const createQueryClient = () =>
+  new QueryClient({
+    defaultOptions: { queries: { retry: false, gcTime: 0 } },
+  });
+
+function renderWithProviders(ui: React.ReactElement) {
+  const qc = createQueryClient();
+  return render(<QueryClientProvider client={qc}>{ui}</QueryClientProvider>);
+}
+
 describe("Budget Panel", () => {
   afterEach(() => {
     vi.clearAllMocks();
   });
 
   it("should render the budget panel and load budgets", async () => {
-    vi.mocked(networking.getBudgetList).mockResolvedValue([
-      {
-        budget_id: "budget-1",
-        max_budget: "100",
-        rpm_limit: 10,
-        tpm_limit: 1000,
-        updated_at: "2024-01-01T00:00:00Z",
-      },
-    ]);
+    vi.mocked(useBudgets).mockReturnValue({
+      data: mockBudgets,
+      isLoading: false,
+    } as any);
 
-    render(<BudgetPanel accessToken="token-123" />);
+    renderWithProviders(<BudgetPanel accessToken="token-123" />);
 
     await waitFor(() => {
       expect(screen.getByText("Create a budget to assign to customers.")).toBeInTheDocument();
@@ -34,17 +53,20 @@ describe("Budget Panel", () => {
   });
 
   it("should open delete modal when clicking delete icon", async () => {
-    vi.mocked(networking.getBudgetList).mockResolvedValue([
-      {
-        budget_id: "budget-to-delete",
-        max_budget: "200",
-        rpm_limit: 20,
-        tpm_limit: 2000,
-        updated_at: "2024-01-02T00:00:00Z",
-      },
-    ]);
+    vi.mocked(useBudgets).mockReturnValue({
+      data: [
+        {
+          budget_id: "budget-to-delete",
+          max_budget: 200,
+          rpm_limit: 20,
+          tpm_limit: 2000,
+          updated_at: "2024-01-02T00:00:00Z",
+        },
+      ],
+      isLoading: false,
+    } as any);
 
-    render(<BudgetPanel accessToken="token-123" />);
+    renderWithProviders(<BudgetPanel accessToken="token-123" />);
 
     await waitFor(() => {
       expect(screen.getByText("budget-to-delete")).toBeInTheDocument();
@@ -62,18 +84,25 @@ describe("Budget Panel", () => {
   });
 
   it("should successfully delete a budget", async () => {
-    vi.mocked(networking.getBudgetList).mockResolvedValue([
-      {
-        budget_id: "budget-to-delete",
-        max_budget: "200",
-        rpm_limit: 20,
-        tpm_limit: 2000,
-        updated_at: "2024-01-02T00:00:00Z",
-      },
-    ]);
-    vi.mocked(networking.budgetDeleteCall).mockResolvedValue(undefined);
+    const deleteMutateAsync = vi.fn().mockResolvedValue(undefined);
+    vi.mocked(useBudgets).mockReturnValue({
+      data: [
+        {
+          budget_id: "budget-to-delete",
+          max_budget: 200,
+          rpm_limit: 20,
+          tpm_limit: 2000,
+          updated_at: "2024-01-02T00:00:00Z",
+        },
+      ],
+      isLoading: false,
+    } as any);
+    vi.mocked(useDeleteBudget).mockReturnValue({
+      mutateAsync: deleteMutateAsync,
+      isPending: false,
+    } as any);
 
-    render(<BudgetPanel accessToken="token-123" />);
+    renderWithProviders(<BudgetPanel accessToken="token-123" />);
 
     await waitFor(() => {
       expect(screen.getByText("budget-to-delete")).toBeInTheDocument();
@@ -96,24 +125,43 @@ describe("Budget Panel", () => {
     });
 
     await waitFor(() => {
-      expect(networking.budgetDeleteCall).toHaveBeenCalledWith("token-123", "budget-to-delete");
-      expect(networking.getBudgetList).toHaveBeenCalledTimes(2); // Initial load + refresh after delete
+      expect(deleteMutateAsync).toHaveBeenCalledWith("budget-to-delete");
+    });
+  });
+
+  it("should render empty state without crashing", async () => {
+    vi.mocked(useBudgets).mockReturnValue({
+      data: [],
+      isLoading: false,
+    } as any);
+
+    renderWithProviders(<BudgetPanel accessToken="token-123" />);
+
+    await waitFor(() => {
+      expect(screen.getByText("Create a budget to assign to customers.")).toBeInTheDocument();
     });
   });
 
   it("should handle delete error", async () => {
-    vi.mocked(networking.getBudgetList).mockResolvedValue([
-      {
-        budget_id: "budget-to-delete",
-        max_budget: "200",
-        rpm_limit: 20,
-        tpm_limit: 2000,
-        updated_at: "2024-01-02T00:00:00Z",
-      },
-    ]);
-    vi.mocked(networking.budgetDeleteCall).mockRejectedValue(new Error("Delete failed"));
+    const deleteMutateAsync = vi.fn().mockRejectedValue(new Error("Delete failed"));
+    vi.mocked(useBudgets).mockReturnValue({
+      data: [
+        {
+          budget_id: "budget-to-delete",
+          max_budget: 200,
+          rpm_limit: 20,
+          tpm_limit: 2000,
+          updated_at: "2024-01-02T00:00:00Z",
+        },
+      ],
+      isLoading: false,
+    } as any);
+    vi.mocked(useDeleteBudget).mockReturnValue({
+      mutateAsync: deleteMutateAsync,
+      isPending: false,
+    } as any);
 
-    render(<BudgetPanel accessToken="token-123" />);
+    renderWithProviders(<BudgetPanel accessToken="token-123" />);
 
     await waitFor(() => {
       expect(screen.getByText("budget-to-delete")).toBeInTheDocument();
@@ -136,10 +184,38 @@ describe("Budget Panel", () => {
     });
 
     await waitFor(() => {
-      expect(networking.budgetDeleteCall).toHaveBeenCalledWith("token-123", "budget-to-delete");
+      expect(deleteMutateAsync).toHaveBeenCalledWith("budget-to-delete");
+    });
+  });
+
+  it("should open edit modal when clicking edit icon", async () => {
+    vi.mocked(useBudgets).mockReturnValue({
+      data: [
+        {
+          budget_id: "budget-to-edit",
+          max_budget: 300,
+          rpm_limit: 30,
+          tpm_limit: 3000,
+          updated_at: "2024-01-03T00:00:00Z",
+        },
+      ],
+      isLoading: false,
+    } as any);
+
+    renderWithProviders(<BudgetPanel accessToken="token-123" />);
+
+    await waitFor(() => {
+      expect(screen.getByText("budget-to-edit")).toBeInTheDocument();
     });
 
-    // Modal should still be open (error handling)
-    expect(screen.getByText("Delete Budget?")).toBeInTheDocument();
+    const editButton = screen.getByTestId("edit-budget-button");
+
+    act(() => {
+      fireEvent.click(editButton);
+    });
+
+    await waitFor(() => {
+      expect(screen.getByText("Edit Budget")).toBeInTheDocument();
+    });
   });
 });

ui/litellm-dashboard/src/components/budgets/budget_panel.tsx

@@ -19,12 +19,12 @@ import {
   TabPanels,
   Text,
 } from "@tremor/react";
-import React, { useEffect, useState } from "react";
+import React, { useState } from "react";
 import { Prism as SyntaxHighlighter } from "react-syntax-highlighter";
 import DeleteResourceModal from "../common_components/DeleteResourceModal";
 import TableIconActionButton from "../common_components/IconActionButton/TableIconActionButtons/TableIconActionButton";
 import NotificationsManager from "../molecules/notifications_manager";
-import { budgetDeleteCall, getBudgetList } from "../networking";
+import { useBudgets, useDeleteBudget } from "@/app/(dashboard)/hooks/budgets/useBudgets";
 import BudgetModal from "./budget_modal";
 import EditBudgetModal from "./edit_budget_modal";
 import { CREATE_END_USER_CURL_COMMAND, CHAT_COMPLETIONS_CURL_COMMAND, OPENAI_SDK_PYTHON_CODE } from "./constants";
@@ -35,7 +35,7 @@ interface BudgetSettingsPageProps {
 
 export interface budgetItem {
   budget_id: string;
-  max_budget: string | null;
+  max_budget: number | null;
   rpm_limit: number | null;
   tpm_limit: number | null;
   updated_at: string;
@@ -45,17 +45,10 @@ const BudgetPanel: React.FC<BudgetSettingsPageProps> = ({ accessToken }) => {
   const [isCreateModelVisible, setIsCreateModelVisible] = useState(false);
   const [isEditModalVisible, setIsEditModalVisible] = useState(false);
   const [selectedBudget, setSelectedBudget] = useState<budgetItem | null>(null);
-  const [budgetList, setBudgetList] = useState<budgetItem[]>([]);
-  const [isDeleting, setIsDeleting] = useState(false);
   const [isDeleteModalVisible, setIsDeleteModalVisible] = useState(false);
-  useEffect(() => {
-    if (!accessToken) {
-      return;
-    }
-    getBudgetList(accessToken).then((data) => {
-      setBudgetList(data);
-    });
-  }, [accessToken]);
+
+  const { data: budgetList = [] } = useBudgets();
+  const deleteBudget = useDeleteBudget();
 
   const handleEditCall = async (budget: budgetItem) => {
     if (accessToken == null) {
@@ -74,11 +67,9 @@ const BudgetPanel: React.FC<BudgetSettingsPageProps> = ({ accessToken }) => {
     if (!selectedBudget || accessToken == null) {
       return;
     }
-    setIsDeleting(true);
     try {
-      await budgetDeleteCall(accessToken, selectedBudget.budget_id);
+      await deleteBudget.mutateAsync(selectedBudget.budget_id);
       NotificationsManager.success("Budget deleted.");
-      await handleUpdateCall();
     } catch (error) {
       console.error("Error deleting budget:", error);
       if (typeof NotificationsManager.fromBackend === "function") {
@@ -87,7 +78,6 @@ const BudgetPanel: React.FC<BudgetSettingsPageProps> = ({ accessToken }) => {
         NotificationsManager.info("Failed to delete budget");
       }
     } finally {
-      setIsDeleting(false);
       setIsDeleteModalVisible(false);
       setSelectedBudget(null);
     }
@@ -97,15 +87,6 @@ const BudgetPanel: React.FC<BudgetSettingsPageProps> = ({ accessToken }) => {
     setIsDeleteModalVisible(false);
   };
 
-  const handleUpdateCall = async () => {
-    if (accessToken == null) {
-      return;
-    }
-    getBudgetList(accessToken).then((data) => {
-      setBudgetList(data);
-    });
-  };
-
   return (
     <div className="w-full mx-auto flex-auto overflow-y-auto m-8 p-2">
       <Button size="sm" variant="primary" className="mb-2" onClick={() => setIsCreateModelVisible(true)}>
@@ -120,19 +101,14 @@ const BudgetPanel: React.FC<BudgetSettingsPageProps> = ({ accessToken }) => {
           <TabPanel>
             <div className="mt-6">
               <BudgetModal
-                accessToken={accessToken}
                 isModalVisible={isCreateModelVisible}
                 setIsModalVisible={setIsCreateModelVisible}
-                setBudgetList={setBudgetList}
               />
               {selectedBudget && (
                 <EditBudgetModal
-                  accessToken={accessToken}
                   isModalVisible={isEditModalVisible}
                   setIsModalVisible={setIsEditModalVisible}
-                  setBudgetList={setBudgetList}
                   existingBudget={selectedBudget}
-                  handleUpdateCall={handleUpdateCall}
                 />
               )}
               <Card>
@@ -149,10 +125,10 @@ const BudgetPanel: React.FC<BudgetSettingsPageProps> = ({ accessToken }) => {
 
                   <TableBody>
                     {budgetList
-                      .slice() // Creates a shallow copy to avoid mutating the original array
-                      .sort((a, b) => new Date(b.updated_at).getTime() - new Date(a.updated_at).getTime()) // Sort by updated_at in descending order
-                      .map((value: budgetItem, index: number) => (
-                        <TableRow key={index}>
+                      .slice()
+                      .sort((a, b) => new Date(b.updated_at).getTime() - new Date(a.updated_at).getTime())
+                      .map((value: budgetItem) => (
+                        <TableRow key={value.budget_id}>
                           <TableCell>{value.budget_id}</TableCell>
                           <TableCell>{value.max_budget ? value.max_budget : "n/a"}</TableCell>
                           <TableCell>{value.tpm_limit ? value.tpm_limit : "n/a"}</TableCell>
@@ -187,7 +163,7 @@ const BudgetPanel: React.FC<BudgetSettingsPageProps> = ({ accessToken }) => {
                 ]}
                 onCancel={handleDeleteCancel}
                 onOk={handleDeleteConfirm}
-                confirmLoading={isDeleting}
+                confirmLoading={deleteBudget.isPending}
               />
             </div>
           </TabPanel>

ui/litellm-dashboard/src/components/budgets/edit_budget_modal.tsx

@@ -1,28 +1,22 @@
 import React, { useEffect } from "react";
 import { TextInput, Accordion, AccordionHeader, AccordionBody } from "@tremor/react";
 import { Button as Button2, Modal, Form, InputNumber, Select } from "antd";
-import { budgetUpdateCall } from "../networking";
+import { useUpdateBudget } from "@/app/(dashboard)/hooks/budgets/useBudgets";
 import { budgetItem } from "./budget_panel";
 import NotificationsManager from "../molecules/notifications_manager";
 
-interface BudgetModalProps {
+interface EditBudgetModalProps {
   isModalVisible: boolean;
-  accessToken: string | null;
   setIsModalVisible: React.Dispatch<React.SetStateAction<boolean>>;
-  setBudgetList: React.Dispatch<React.SetStateAction<any[]>>;
   existingBudget: budgetItem;
-  handleUpdateCall: () => void;
 }
-const EditBudgetModal: React.FC<BudgetModalProps> = ({
+const EditBudgetModal: React.FC<EditBudgetModalProps> = ({
   isModalVisible,
-  accessToken,
   setIsModalVisible,
-  setBudgetList,
   existingBudget,
-  handleUpdateCall,
 }) => {
-  console.log("existingBudget", existingBudget);
   const [form] = Form.useForm();
+  const updateBudget = useUpdateBudget();
 
   useEffect(() => {
     form.setFieldsValue(existingBudget);
@@ -38,21 +32,16 @@ const EditBudgetModal: React.FC<BudgetModalProps> = ({
     form.resetFields();
   };
 
-  const handleCreate = async (formValues: Record<string, any>) => {
-    if (accessToken == null || accessToken == undefined) {
-      return;
-    }
+  const handleUpdate = async (formValues: Record<string, any>) => {
     try {
       NotificationsManager.info("Making API Call");
-      setIsModalVisible(true);
-      const response = await budgetUpdateCall(accessToken, formValues);
-      setBudgetList((prevData) => (prevData ? [...prevData, response] : [response])); // Check if prevData is null
+      await updateBudget.mutateAsync(formValues);
       NotificationsManager.success("Budget Updated");
       form.resetFields();
-      handleUpdateCall();
+      setIsModalVisible(false);
     } catch (error) {
-      console.error("Error creating the key:", error);
-      NotificationsManager.fromBackend(`Error creating the key: ${error}`);
+      console.error("Error updating the budget:", error);
+      NotificationsManager.fromBackend(`Error updating the budget: ${error}`);
     }
   };
 
@@ -67,7 +56,7 @@ const EditBudgetModal: React.FC<BudgetModalProps> = ({
     >
       <Form
         form={form}
-        onFinish={handleCreate}
+        onFinish={handleUpdate}
         labelCol={{ span: 8 }}
         wrapperCol={{ span: 16 }}
         labelAlign="left"
@@ -77,15 +66,9 @@ const EditBudgetModal: React.FC<BudgetModalProps> = ({
           <Form.Item
             label="Budget ID"
             name="budget_id"
-            rules={[
-              {
-                required: true,
-                message: "Please input a human-friendly name for the budget",
-              },
-            ]}
-            help="A human-friendly name for the budget"
+            help="Budget ID cannot be changed after creation"
           >
-            <TextInput placeholder="" />
+            <TextInput placeholder="" disabled={true} />
           </Form.Item>
           <Form.Item label="Max Tokens per minute" name="tpm_limit" help="Default is model limit.">
             <InputNumber step={1} precision={2} width={200} />