Merge pull request #22715 from BerriAI/litellm_reset_key_spend

[Feature] UI - Virtual Keys: Manual Spend Reset

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<ResetKeySpendResponse> => {
+  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<ResetKeySpendResponse, Error, string>({
+    mutationFn: async (keyToken) => {
+      if (!accessToken) {
+        throw new Error("Access token is required");
+      }
+      return resetKeySpend(accessToken, keyToken);
+    },
+    onSuccess: () => {
+      queryClient.invalidateQueries({ queryKey: keyKeys.all });
+    },
+  });
+};

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({
                 </Button>
               </span>
             </Tooltip>
+            {onResetSpend && (
+              <Button danger icon={<TransactionOutlined />} onClick={onResetSpend}>
+                Reset Spend
+              </Button>
+            )}
             <Button danger icon={<DeleteOutlined />} onClick={onDelete}>
               Delete Key
             </Button>

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(
+        <KeyInfoView keyData={MOCK_KEY_DATA} onClose={() => { }} 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(
+        <KeyInfoView keyData={keyData} onClose={() => { }} 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(
+        <KeyInfoView keyData={keyData} onClose={() => { }} 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(
+        <KeyInfoView keyData={MOCK_KEY_DATA} onClose={() => { }} 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(
+        <KeyInfoView keyData={keyDataWithSpend} onClose={() => { }} 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) }),
+        );
+      });
+    });
+  });
 });

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<KeyResponse | undefined>(keyData);
   const [lastRegeneratedAt, setLastRegeneratedAt] = useState<Date | null>(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 (
     <div className="w-full h-screen p-4">
       <KeyInfoHeader
@@ -353,6 +381,7 @@ export default function KeyInfoView({
         onBack={onClose}
         onRegenerate={() => 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 */}
+      <Modal
+        title="Reset Key Spend"
+        open={isResetSpendModalOpen}
+        onOk={handleResetSpend}
+        onCancel={() => setIsResetSpendModalOpen(false)}
+        okText="Reset"
+        okButtonProps={{ danger: true }}
+        confirmLoading={resetSpendLoading}
+      >
+        <p>
+          Reset spend for <strong>{currentKeyData?.key_alias || currentKeyData?.token_id || "this key"}</strong> to{" "}
+          <strong>$0</strong>?
+        </p>
+        <p style={{ color: "#666", fontSize: "0.875rem", marginTop: 8 }}>
+          Current spend: <strong>${formatNumberWithCommas(currentKeyData.spend, 4)}</strong>. Spend history is
+          preserved in logs. This resets the current period spend counter, the same as an automatic budget reset.
+        </p>
+      </Modal>
+
       <TabGroup>
         <TabList className="mb-4">
           <Tab>Overview</Tab>