Merge pull request #18555 from BerriAI/litellm_ui_usage_budget

[Fix] UI - Usage Page User Max Budget

ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useCurrentUser.test.ts

@@ -0,0 +1,253 @@
+import { describe, it, expect, vi, beforeEach } from "vitest";
+import { renderHook, waitFor } from "@testing-library/react";
+import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
+import React, { ReactNode } from "react";
+import { useCurrentUser } from "./useCurrentUser";
+import { userInfoCall } from "@/components/networking";
+import type { UserInfo } from "@/components/view_users/types";
+
+// Mock the networking function
+vi.mock("@/components/networking", () => ({
+  userInfoCall: vi.fn(),
+}));
+
+// Mock the queryKeysFactory - we'll mock the specific return value
+vi.mock("../common/queryKeysFactory", () => ({
+  createQueryKeys: vi.fn((resource: string) => ({
+    all: [resource],
+    lists: () => [resource, "list"],
+    list: (params?: any) => [resource, "list", { params }],
+    details: () => [resource, "detail"],
+    detail: (uid: string) => [resource, "detail", uid],
+  })),
+}));
+
+// Mock useAuthorized hook - we can override this in individual tests
+const mockUseAuthorized = vi.fn();
+vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({
+  default: () => mockUseAuthorized(),
+}));
+
+// Mock data - response from userInfoCall should have user_info property
+const mockUserInfoResponse = {
+  user_info: {
+    user_id: "test-user-id",
+    user_email: "test@example.com",
+    user_alias: "Test User",
+    user_role: "Admin",
+    spend: 150.75,
+    max_budget: 1000.0,
+    key_count: 5,
+    created_at: "2024-01-01T00:00:00Z",
+    updated_at: "2024-01-01T00:00:00Z",
+    sso_user_id: null,
+    budget_duration: "monthly",
+  } as UserInfo,
+};
+
+describe("useCurrentUser", () => {
+  let queryClient: QueryClient;
+
+  beforeEach(() => {
+    queryClient = new QueryClient({
+      defaultOptions: {
+        queries: {
+          retry: false,
+        },
+      },
+    });
+
+    // Reset all mocks
+    vi.clearAllMocks();
+
+    // Set default mock for useAuthorized (enabled state)
+    mockUseAuthorized.mockReturnValue({
+      accessToken: "test-access-token",
+      userId: "test-user-id",
+      userRole: "Admin",
+      token: "test-token",
+      userEmail: "test@example.com",
+      premiumUser: false,
+      disabledPersonalKeyCreation: null,
+      showSSOBanner: false,
+    });
+  });
+
+  const wrapper = ({ children }: { children: ReactNode }) =>
+    React.createElement(QueryClientProvider, { client: queryClient }, children);
+
+  it("should return user info data when query is successful", async () => {
+    // Mock successful API call
+    (userInfoCall as any).mockResolvedValue(mockUserInfoResponse);
+
+    const { result } = renderHook(() => useCurrentUser(), { wrapper });
+
+    // Initially loading
+    expect(result.current.isLoading).toBe(true);
+    expect(result.current.data).toBeUndefined();
+
+    // Wait for success
+    await waitFor(() => {
+      expect(result.current.isLoading).toBe(false);
+      expect(result.current.isSuccess).toBe(true);
+    });
+
+    expect(result.current.data).toEqual(mockUserInfoResponse.user_info);
+    expect(result.current.error).toBeNull();
+    expect(userInfoCall).toHaveBeenCalledWith("test-access-token", "test-user-id", "Admin", false, null, null);
+    expect(userInfoCall).toHaveBeenCalledTimes(1);
+  });
+
+  it("should handle error when userInfoCall fails", async () => {
+    const errorMessage = "Failed to fetch user info";
+    const testError = new Error(errorMessage);
+
+    // Mock failed API call
+    (userInfoCall as any).mockRejectedValue(testError);
+
+    const { result } = renderHook(() => useCurrentUser(), { wrapper });
+
+    // Initially loading
+    expect(result.current.isLoading).toBe(true);
+
+    // Wait for error
+    await waitFor(() => {
+      expect(result.current.isLoading).toBe(false);
+      expect(result.current.isError).toBe(true);
+    });
+
+    expect(result.current.error).toEqual(testError);
+    expect(result.current.data).toBeUndefined();
+    expect(userInfoCall).toHaveBeenCalledWith("test-access-token", "test-user-id", "Admin", false, null, null);
+    expect(userInfoCall).toHaveBeenCalledTimes(1);
+  });
+
+  it("should not execute query when accessToken is missing", async () => {
+    // Mock missing accessToken
+    mockUseAuthorized.mockReturnValue({
+      accessToken: null,
+      userId: "test-user-id",
+      userRole: "Admin",
+      token: null,
+      userEmail: "test@example.com",
+      premiumUser: false,
+      disabledPersonalKeyCreation: null,
+      showSSOBanner: false,
+    });
+
+    const { result } = renderHook(() => useCurrentUser(), { wrapper });
+
+    // Query should not execute
+    expect(result.current.isLoading).toBe(false);
+    expect(result.current.data).toBeUndefined();
+    expect(result.current.isFetched).toBe(false);
+
+    // API should not be called
+    expect(userInfoCall).not.toHaveBeenCalled();
+  });
+
+  it("should not execute query when userId is missing", async () => {
+    // Mock missing userId
+    mockUseAuthorized.mockReturnValue({
+      accessToken: "test-access-token",
+      userId: null,
+      userRole: "Admin",
+      token: "test-token",
+      userEmail: "test@example.com",
+      premiumUser: false,
+      disabledPersonalKeyCreation: null,
+      showSSOBanner: false,
+    });
+
+    const { result } = renderHook(() => useCurrentUser(), { wrapper });
+
+    // Query should not execute
+    expect(result.current.isLoading).toBe(false);
+    expect(result.current.data).toBeUndefined();
+    expect(result.current.isFetched).toBe(false);
+
+    // API should not be called
+    expect(userInfoCall).not.toHaveBeenCalled();
+  });
+
+  it("should not execute query when userRole is missing", async () => {
+    // Mock missing userRole
+    mockUseAuthorized.mockReturnValue({
+      accessToken: "test-access-token",
+      userId: "test-user-id",
+      userRole: null,
+      token: "test-token",
+      userEmail: "test@example.com",
+      premiumUser: false,
+      disabledPersonalKeyCreation: null,
+      showSSOBanner: false,
+    });
+
+    const { result } = renderHook(() => useCurrentUser(), { wrapper });
+
+    // Query should not execute
+    expect(result.current.isLoading).toBe(false);
+    expect(result.current.data).toBeUndefined();
+    expect(result.current.isFetched).toBe(false);
+
+    // API should not be called
+    expect(userInfoCall).not.toHaveBeenCalled();
+  });
+
+  it("should not execute query when all auth values are missing", async () => {
+    // Mock all auth values missing
+    mockUseAuthorized.mockReturnValue({
+      accessToken: null,
+      userId: null,
+      userRole: null,
+      token: null,
+      userEmail: "test@example.com",
+      premiumUser: false,
+      disabledPersonalKeyCreation: null,
+      showSSOBanner: false,
+    });
+
+    const { result } = renderHook(() => useCurrentUser(), { wrapper });
+
+    // Query should not execute
+    expect(result.current.isLoading).toBe(false);
+    expect(result.current.data).toBeUndefined();
+    expect(result.current.isFetched).toBe(false);
+
+    // API should not be called
+    expect(userInfoCall).not.toHaveBeenCalled();
+  });
+
+  it("should execute query when all auth values are present", async () => {
+    // Mock successful API call
+    (userInfoCall as any).mockResolvedValue(mockUserInfoResponse);
+
+    // Ensure all auth values are present (already set in beforeEach)
+    const { result } = renderHook(() => useCurrentUser(), { wrapper });
+
+    // Wait for query to execute
+    await waitFor(() => {
+      expect(result.current.isLoading).toBe(false);
+    });
+
+    expect(userInfoCall).toHaveBeenCalledWith("test-access-token", "test-user-id", "Admin", false, null, null);
+    expect(userInfoCall).toHaveBeenCalledTimes(1);
+  });
+
+  it("should handle network timeout error", async () => {
+    const timeoutError = new Error("Network timeout");
+
+    // Mock network timeout
+    (userInfoCall as any).mockRejectedValue(timeoutError);
+
+    const { result } = renderHook(() => useCurrentUser(), { wrapper });
+
+    // Wait for error
+    await waitFor(() => {
+      expect(result.current.isError).toBe(true);
+    });
+
+    expect(result.current.error).toEqual(timeoutError);
+    expect(result.current.data).toBeUndefined();
+  });
+});

ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useCurrentUser.ts

@@ -0,0 +1,19 @@
+import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
+import { UserInfo, userInfoCall } from "@/components/networking";
+import { useQuery, UseQueryResult } from "@tanstack/react-query";
+import { createQueryKeys } from "../common/queryKeysFactory";
+
+const userKeys = createQueryKeys("users");
+
+export const useCurrentUser = (): UseQueryResult<UserInfo> => {
+  const { accessToken, userId, userRole } = useAuthorized();
+  return useQuery<UserInfo>({
+    queryKey: userKeys.detail(userId!),
+    queryFn: async () => {
+      const data = await userInfoCall(accessToken!, userId!, userRole!, false, null, null);
+      console.log(`userInfo: ${JSON.stringify(data)}`);
+      return data.user_info;
+    },
+    enabled: Boolean(accessToken && userId && userRole),
+  });
+};

ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.test.tsx

@@ -1,8 +1,9 @@
 import { useAgents } from "@/app/(dashboard)/hooks/agents/useAgents";
 import { useCustomers } from "@/app/(dashboard)/hooks/customers/useCustomers";
 import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
-import { act, fireEvent, render, screen, waitFor } from "@testing-library/react";
+import { act, fireEvent, screen, waitFor } from "@testing-library/react";
 import { beforeAll, beforeEach, describe, expect, it, vi } from "vitest";
+import { renderWithProviders } from "../../../../tests/test-utils";
 import type { Organization } from "../../networking";
 import * as networking from "../../networking";
 import NewUsagePage from "./UsagePageView";
@@ -332,7 +333,7 @@ describe("NewUsage", () => {
   });
 
   it("should render and fetch usage data on mount", async () => {
-    render(<NewUsagePage {...defaultProps} />);
+    renderWithProviders(<NewUsagePage {...defaultProps} />);
 
     // Wait for data to be fetched
     await waitFor(() => {
@@ -349,7 +350,7 @@ describe("NewUsage", () => {
   });
 
   it("should display usage metrics and charts", async () => {
-    render(<NewUsagePage {...defaultProps} />);
+    renderWithProviders(<NewUsagePage {...defaultProps} />);
 
     await waitFor(() => {
       expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalled();
@@ -367,7 +368,7 @@ describe("NewUsage", () => {
   });
 
   it("should switch between usage views correctly", async () => {
-    render(<NewUsagePage {...defaultProps} />);
+    renderWithProviders(<NewUsagePage {...defaultProps} />);
 
     await waitFor(() => {
       expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalled();
@@ -401,7 +402,7 @@ describe("NewUsage", () => {
   });
 
   it("should show organization usage banner and view for admins", async () => {
-    render(<NewUsagePage {...defaultProps} organizations={mockOrganizations} />);
+    renderWithProviders(<NewUsagePage {...defaultProps} organizations={mockOrganizations} />);
 
     await waitFor(() => {
       expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalled();
@@ -426,7 +427,7 @@ describe("NewUsage", () => {
       error: null,
     } as any);
 
-    render(<NewUsagePage {...defaultProps} />);
+    renderWithProviders(<NewUsagePage {...defaultProps} />);
 
     await waitFor(() => {
       expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalled();
@@ -450,7 +451,7 @@ describe("NewUsage", () => {
       error: null,
     } as any);
 
-    render(<NewUsagePage {...defaultProps} />);
+    renderWithProviders(<NewUsagePage {...defaultProps} />);
 
     await waitFor(() => {
       expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalled();

ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.tsx

@@ -52,6 +52,7 @@ import { valueFormatterSpend } from "../utils/value_formatters";
 import EntityUsage, { EntityList } from "./EntityUsage/EntityUsage";
 import TopKeyView from "./EntityUsage/TopKeyView";
 import { UsageOption, UsageViewSelect } from "./UsageViewSelect/UsageViewSelect";
+import { useCurrentUser } from "@/app/(dashboard)/hooks/users/useCurrentUser";
 
 interface UsagePageProps {
   teams: Team[];
@@ -82,6 +83,9 @@ const UsagePage: React.FC<UsagePageProps> = ({ teams, organizations }) => {
   const [allTags, setAllTags] = useState<EntityList[]>([]);
   const { data: customers = [] } = useCustomers();
   const { data: agentsResponse } = useAgents();
+  const { data: currentUser } = useCurrentUser();
+  console.log(`currentUser: ${JSON.stringify(currentUser)}`);
+  console.log(`currentUser max budget: ${currentUser?.max_budget}`);
   const [modelViewType, setModelViewType] = useState<"groups" | "individual">("groups");
   const [isCloudZeroModalOpen, setIsCloudZeroModalOpen] = useState(false);
   const [isGlobalExportModalOpen, setIsGlobalExportModalOpen] = useState(false);
@@ -477,7 +481,11 @@ const UsagePage: React.FC<UsagePageProps> = ({ teams, organizations }) => {
                         )}
                       </Text>
 
-                      <ViewUserSpend userSpend={totalSpend} selectedTeam={null} userMaxBudget={null} />
+                      <ViewUserSpend
+                        userSpend={totalSpend}
+                        selectedTeam={null}
+                        userMaxBudget={currentUser?.max_budget || null}
+                      />
                     </Col>
 
                     <Col numColSpan={2}>

ui/litellm-dashboard/src/components/view_user_spend.tsx

@@ -17,7 +17,6 @@ interface ViewUserSpendProps {
 }
 const ViewUserSpend: React.FC<ViewUserSpendProps> = ({ userSpend, userMaxBudget, selectedTeam }) => {
   const { accessToken, userRole, userId: userID } = useAuthorized();
-  console.log(`userSpend: ${userSpend}`);
   let [spend, setSpend] = useState(userSpend !== null ? userSpend : 0.0);
   const [maxBudget, setMaxBudget] = useState(
     selectedTeam ? Number(formatNumberWithCommas(selectedTeam.max_budget, 4)) : null,
@@ -61,6 +60,8 @@ const ViewUserSpend: React.FC<ViewUserSpendProps> = ({ userSpend, userMaxBudget,
           setMaxBudget(selectedTeam.max_budget);
         }
       }
+    } else {
+      setMaxBudget(userMaxBudget);
     }
   }, [selectedTeam, userMaxBudget]);
   const [userModels, setUserModels] = useState([]);