diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useResetKeySpend.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useResetKeySpend.ts new file mode 100644 index 00000000000..a845fc5881a --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useResetKeySpend.ts @@ -0,0 +1,66 @@ +import { useMutation, useQueryClient } from "@tanstack/react-query"; +import { + getProxyBaseUrl, + getGlobalLitellmHeaderName, + deriveErrorMessage, + handleError, +} from "@/components/networking"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import { keyKeys } from "./useKeys"; + +// ── Types ───────────────────────────────────────────────────────────────────── + +export interface ResetKeySpendResponse { + key_hash: string; + spend: number; + previous_spend: number; + max_budget: number | null; + budget_reset_at: string | null; +} + +// ── Fetch function ──────────────────────────────────────────────────────────── + +export const resetKeySpend = async ( + accessToken: string, + keyToken: string, +): Promise => { + const baseUrl = getProxyBaseUrl(); + const url = `${baseUrl ? `${baseUrl}/key/${keyToken}/reset_spend` : `/key/${keyToken}/reset_spend`}`; + + const response = await fetch(url, { + method: "POST", + headers: { + [getGlobalLitellmHeaderName()]: `Bearer ${accessToken}`, + "Content-Type": "application/json", + }, + body: JSON.stringify({ reset_to: 0 }), + }); + + if (!response.ok) { + const errorData = await response.json(); + const errorMessage = deriveErrorMessage(errorData); + handleError(errorMessage); + throw new Error(errorMessage); + } + + return response.json(); +}; + +// ── Hook ────────────────────────────────────────────────────────────────────── + +export const useResetKeySpend = () => { + const { accessToken } = useAuthorized(); + const queryClient = useQueryClient(); + + return useMutation({ + mutationFn: async (keyToken) => { + if (!accessToken) { + throw new Error("Access token is required"); + } + return resetKeySpend(accessToken, keyToken); + }, + onSuccess: () => { + queryClient.invalidateQueries({ queryKey: keyKeys.all }); + }, + }); +}; diff --git a/ui/litellm-dashboard/src/components/templates/KeyInfoHeader.tsx b/ui/litellm-dashboard/src/components/templates/KeyInfoHeader.tsx index 1befd657843..93ebae9c4be 100644 --- a/ui/litellm-dashboard/src/components/templates/KeyInfoHeader.tsx +++ b/ui/litellm-dashboard/src/components/templates/KeyInfoHeader.tsx @@ -11,6 +11,7 @@ import { ClockCircleOutlined, ThunderboltOutlined, SafetyCertificateOutlined, + TransactionOutlined, } from "@ant-design/icons"; import LabeledField from "../common_components/LabeledField"; @@ -33,6 +34,7 @@ interface KeyInfoHeaderProps { onCreateNew?: () => void; onRegenerate?: () => void; onDelete?: () => void; + onResetSpend?: () => void; canModifyKey?: boolean; backButtonText?: string; regenerateDisabled?: boolean; @@ -45,6 +47,7 @@ export function KeyInfoHeader({ onCreateNew, onRegenerate, onDelete, + onResetSpend, canModifyKey = true, backButtonText = "Back to Keys", regenerateDisabled = false, @@ -84,6 +87,11 @@ export function KeyInfoHeader({ + {onResetSpend && ( + + )} diff --git a/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx b/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx index e58ccfbbfd1..f269ad96a27 100644 --- a/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx @@ -5,6 +5,7 @@ import { renderWithProviders } from "../../../tests/test-utils"; import { screen, waitFor } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { beforeEach, describe, expect, it, vi } from "vitest"; +import { useResetKeySpend } from "@/app/(dashboard)/hooks/keys/useResetKeySpend"; import { KeyResponse, Team } from "../key_team_helpers/key_list"; import KeyInfoView from "./key_info_view"; @@ -28,6 +29,14 @@ vi.mock("../networking", () => ({ }), })); +const mockResetKeySpendMutate = vi.fn(); +vi.mock("@/app/(dashboard)/hooks/keys/useResetKeySpend", () => ({ + useResetKeySpend: vi.fn(() => ({ + mutate: mockResetKeySpendMutate, + isPending: false, + })), +})); + vi.mock("@/utils/dataUtils", () => ({ copyToClipboard: vi.fn().mockResolvedValue(true), formatNumberWithCommas: vi.fn((value: number, decimals?: number) => { @@ -539,4 +548,136 @@ describe("KeyInfoView", () => { expect(screen.getByText("Key not found")).toBeInTheDocument(); }); }); + + describe("Reset Spend button visibility", () => { + it("should show Reset Spend button for proxy admin", async () => { + vi.mocked(useTeams).mockReturnValue({ teams: [], setTeams: vi.fn() }); + vi.mocked(useAuthorized).mockReturnValue({ + ...baseUseAuthorizedMock, + userId: "proxy-admin-user", + userRole: "proxy_admin", + }); + + renderWithProviders( + { }} keyId={"test-key-id"} onKeyDataUpdate={() => { }} teams={[]} />, + ); + + await waitFor(() => { + expect(screen.getByRole("button", { name: /reset spend/i })).toBeInTheDocument(); + }); + }); + + it("should show Reset Spend button for team admin of key's team", async () => { + const teamId = "test-team-id"; + const teamAdminUserId = "team-admin-user"; + const mockTeam: Team = { + team_id: teamId, + team_alias: "Test Team", + models: [], + max_budget: null, + budget_duration: null, + tpm_limit: null, + rpm_limit: null, + organization_id: "org-1", + created_at: "2025-01-01T00:00:00Z", + keys: [], + members_with_roles: [{ user_id: teamAdminUserId, role: "admin" }], + spend: 0, + }; + + vi.mocked(useTeams).mockReturnValue({ teams: [mockTeam], setTeams: vi.fn() }); + vi.mocked(useAuthorized).mockReturnValue({ + ...baseUseAuthorizedMock, + userId: teamAdminUserId, + userRole: "user", + }); + + const keyData = { ...MOCK_KEY_DATA, team_id: teamId, user_id: "other-user-id" }; + renderWithProviders( + { }} keyId={"test-key-id"} onKeyDataUpdate={() => { }} teams={[]} />, + ); + + await waitFor(() => { + expect(screen.getByRole("button", { name: /reset spend/i })).toBeInTheDocument(); + }); + }); + + it("should not show Reset Spend button for regular key owner", async () => { + vi.mocked(useTeams).mockReturnValue({ teams: [], setTeams: vi.fn() }); + vi.mocked(useAuthorized).mockReturnValue({ + ...baseUseAuthorizedMock, + userId: "owner-user-id", + userRole: "user", + }); + + const keyData = { ...MOCK_KEY_DATA, user_id: "owner-user-id" }; + renderWithProviders( + { }} keyId={"test-key-id"} onKeyDataUpdate={() => { }} teams={[]} />, + ); + + await waitFor(() => { + expect(screen.queryByRole("button", { name: /reset spend/i })).not.toBeInTheDocument(); + }); + }); + }); + + describe("Reset Spend modal flow", () => { + it("should open confirmation modal when Reset Spend is clicked", async () => { + vi.mocked(useTeams).mockReturnValue({ teams: [], setTeams: vi.fn() }); + vi.mocked(useAuthorized).mockReturnValue({ + ...baseUseAuthorizedMock, + userId: "proxy-admin-user", + userRole: "proxy_admin", + }); + + renderWithProviders( + { }} keyId={"test-key-id"} onKeyDataUpdate={() => { }} teams={[]} />, + ); + + await waitFor(() => { + expect(screen.getByRole("button", { name: /reset spend/i })).toBeInTheDocument(); + }); + + await userEvent.click(screen.getByRole("button", { name: /reset spend/i })); + + await waitFor(() => { + expect(screen.getByText("Reset Key Spend")).toBeInTheDocument(); + expect(screen.getByRole("button", { name: /^reset$/i })).toBeInTheDocument(); + }); + }); + + it("should call mutate with token on confirm", async () => { + vi.mocked(useTeams).mockReturnValue({ teams: [], setTeams: vi.fn() }); + vi.mocked(useAuthorized).mockReturnValue({ + ...baseUseAuthorizedMock, + userId: "proxy-admin-user", + userRole: "proxy_admin", + }); + + const keyDataWithSpend = { ...MOCK_KEY_DATA, spend: 5.0 }; + renderWithProviders( + { }} keyId={"test-key-id"} onKeyDataUpdate={() => { }} teams={[]} />, + ); + + await waitFor(() => { + expect(screen.getByRole("button", { name: /reset spend/i })).toBeInTheDocument(); + }); + + await userEvent.click(screen.getByRole("button", { name: /reset spend/i })); + + await waitFor(() => { + expect(screen.getByText("Reset Key Spend")).toBeInTheDocument(); + }); + + // Click the confirm button in the modal + await userEvent.click(screen.getByRole("button", { name: /^reset$/i })); + + await waitFor(() => { + expect(mockResetKeySpendMutate).toHaveBeenCalledWith( + MOCK_KEY_DATA.token, + expect.objectContaining({ onSuccess: expect.any(Function), onError: expect.any(Function) }), + ); + }); + }); + }); }); diff --git a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx index 72895c4d617..6733cd6a595 100644 --- a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx @@ -6,7 +6,7 @@ import { formatNumberWithCommas } from "@/utils/dataUtils"; import { mapEmptyStringToNull } from "@/utils/keyUpdateUtils"; import { ArrowLeftIcon } from "@heroicons/react/outline"; import { Badge, Button, Card, Grid, Tab, TabGroup, TabList, TabPanel, TabPanels, Text, Title } from "@tremor/react"; -import { Form, Tag } from "antd"; +import { Form, Modal, Tag } from "antd"; import { KeyInfoHeader } from "./KeyInfoHeader"; import { useEffect, useState } from "react"; import { isProxyAdminRole, isUserTeamAdminForSingleTeam } from "../../utils/roles"; @@ -18,6 +18,7 @@ import { KeyResponse } from "../key_team_helpers/key_list"; import LoggingSettingsView from "../logging_settings_view"; import NotificationManager from "../molecules/notifications_manager"; import { getPolicyInfoWithGuardrails, keyDeleteCall, keyUpdateCall } from "../networking"; +import { useResetKeySpend } from "@/app/(dashboard)/hooks/keys/useResetKeySpend"; import ObjectPermissionsView from "../object_permissions_view"; import { RegenerateKeyModal } from "../organisms/regenerate_key_modal"; import { parseErrorMessage } from "../shared/errorUtils"; @@ -59,6 +60,8 @@ export default function KeyInfoView({ const [deleteLoading, setDeleteLoading] = useState(false); const [deleteConfirmInput, setDeleteConfirmInput] = useState(""); const [isRegenerateModalOpen, setIsRegenerateModalOpen] = useState(false); + const [isResetSpendModalOpen, setIsResetSpendModalOpen] = useState(false); + const { mutate: resetKeySpend, isPending: resetSpendLoading } = useResetKeySpend(); // Add local state to maintain key data and track regeneration const [currentKeyData, setCurrentKeyData] = useState(keyData); const [lastRegeneratedAt, setLastRegeneratedAt] = useState(null); @@ -337,6 +340,31 @@ export default function KeyInfoView({ )) || (userID === currentKeyData.user_id && userRole !== "Internal Viewer"); + const canResetSpend = + isProxyAdminRole(userRole || "") || + (teamsData && + isUserTeamAdminForSingleTeam( + teamsData?.filter((team) => team.team_id === currentKeyData.team_id)[0]?.members_with_roles, + userID || "", + )); + + const handleResetSpend = () => { + resetKeySpend(currentKeyData.token || currentKeyData.token_id, { + onSuccess: () => { + setCurrentKeyData((prevData) => (prevData ? { ...prevData, spend: 0 } : undefined)); + if (onKeyDataUpdate) { + onKeyDataUpdate({ spend: 0 }); + } + NotificationManager.success("Key spend reset to $0"); + setIsResetSpendModalOpen(false); + }, + onError: (error) => { + NotificationManager.fromBackend(parseErrorMessage(error)); + console.error("Error resetting key spend:", error); + }, + }); + }; + return (
setIsRegenerateModalOpen(true)} onDelete={() => setIsDeleteModalOpen(true)} + onResetSpend={canResetSpend ? () => setIsResetSpendModalOpen(true) : undefined} canModifyKey={canModifyKey} backButtonText={backButtonText} regenerateDisabled={!premiumUser} @@ -407,6 +436,26 @@ export default function KeyInfoView({ requiredConfirmation={currentKeyData?.key_alias} /> + {/* Reset Spend Confirmation Modal */} + setIsResetSpendModalOpen(false)} + okText="Reset" + okButtonProps={{ danger: true }} + confirmLoading={resetSpendLoading} + > +

+ Reset spend for {currentKeyData?.key_alias || currentKeyData?.token_id || "this key"} to{" "} + $0? +

+

+ Current spend: ${formatNumberWithCommas(currentKeyData.spend, 4)}. Spend history is + preserved in logs. This resets the current period spend counter, the same as an automatic budget reset. +

+
+ Overview