diff --git a/ui/litellm-dashboard/e2e_tests/tests/proxy-admin/keys.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/proxy-admin/keys.spec.ts index aba37e25be..24e8d4f4b3 100644 --- a/ui/litellm-dashboard/e2e_tests/tests/proxy-admin/keys.spec.ts +++ b/ui/litellm-dashboard/e2e_tests/tests/proxy-admin/keys.spec.ts @@ -63,8 +63,9 @@ test.describe("Proxy Admin - Keys", () => { await page.getByRole("button", { name: "Regenerate Key" }).click(); await page.getByRole("button", { name: "Regenerate", exact: true }).click(); - // Success shows "Copy Virtual Key" button in the regenerated key dialog - await expect(page.getByText("Copy Virtual Key")).toBeVisible({ timeout: 10_000 }); + // Success view shows the warning banner and a Copy button for the regenerated key + await expect(page.getByText("Save it now, you will not see it again")).toBeVisible({ timeout: 10_000 }); + await expect(page.getByRole("button", { name: /Copy/ })).toBeVisible({ timeout: 10_000 }); }); test("Update key TPM and RPM limits", async ({ page }) => { diff --git a/ui/litellm-dashboard/src/components/organisms/RegenerateKeyModal.test.tsx b/ui/litellm-dashboard/src/components/organisms/RegenerateKeyModal.test.tsx new file mode 100644 index 0000000000..1cb77bb9af --- /dev/null +++ b/ui/litellm-dashboard/src/components/organisms/RegenerateKeyModal.test.tsx @@ -0,0 +1,271 @@ +import { describe, it, expect, vi, beforeEach } from "vitest"; +import userEvent from "@testing-library/user-event"; +import { renderWithProviders, screen, waitFor } from "../../../tests/test-utils"; +import { RegenerateKeyModal } from "./RegenerateKeyModal"; +import { KeyResponse } from "../key_team_helpers/key_list"; + +// Mock the networking call +const mockRegenerateKeyCall = vi.fn(); +vi.mock("../networking", () => ({ + regenerateKeyCall: (...args: unknown[]) => mockRegenerateKeyCall(...args), +})); + +// Mock CopyToClipboard to render a simple button +vi.mock("react-copy-to-clipboard", () => ({ + CopyToClipboard: ({ children, onCopy }: { children: React.ReactElement; onCopy: () => void }) => { + const React = require("react"); + return React.cloneElement(children, { onClick: onCopy }); + }, +})); + +const makeToken = (overrides: Partial = {}): KeyResponse => + ({ + token: "token-hash-123", + token_id: "token-id-123", + key_name: "sk-test-key", + key_alias: "my-test-key", + max_budget: 100, + tpm_limit: 5000, + rpm_limit: 500, + duration: "30d", + expires: "2026-12-31T00:00:00Z", + ...overrides, + }) as KeyResponse; + +describe("RegenerateKeyModal", () => { + const mockOnClose = vi.fn(); + const mockOnKeyUpdate = vi.fn(); + + const defaultProps = { + selectedToken: makeToken(), + visible: true, + onClose: mockOnClose, + onKeyUpdate: mockOnKeyUpdate, + }; + + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("should render the modal with correct title", () => { + renderWithProviders(); + expect(screen.getByText("Regenerate Virtual Key")).toBeInTheDocument(); + }); + + it("should not render the modal when visible is false", () => { + renderWithProviders(); + expect(screen.queryByText("Regenerate Virtual Key")).not.toBeInTheDocument(); + }); + + it("should display the form with pre-filled values", () => { + renderWithProviders(); + + const keyAliasInput = screen.getByLabelText("Key Alias") as HTMLInputElement; + expect(keyAliasInput).toBeDisabled(); + expect(keyAliasInput).toHaveValue("my-test-key"); + }); + + it("should display the current expiry when token has expires", () => { + renderWithProviders(); + expect(screen.getByText(/Current expiry:/)).toBeInTheDocument(); + }); + + it("should display 'Never' when token has no expires", () => { + renderWithProviders( + , + ); + expect(screen.getByText("Current expiry: Never")).toBeInTheDocument(); + }); + + it("should show Cancel and Regenerate buttons in form view", () => { + renderWithProviders(); + expect(screen.getByRole("button", { name: "Cancel" })).toBeInTheDocument(); + expect(screen.getByRole("button", { name: /Regenerate/ })).toBeInTheDocument(); + }); + + it("should call onClose when Cancel is clicked", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + await user.click(screen.getByRole("button", { name: "Cancel" })); + expect(mockOnClose).toHaveBeenCalledOnce(); + }); + + it("should call onClose when the X close button is clicked", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + await user.click(screen.getByRole("button", { name: "Close" })); + expect(mockOnClose).toHaveBeenCalledOnce(); + }); + + it("should render form fields for budget and rate limits", () => { + renderWithProviders(); + + expect(screen.getByText("Max Budget (USD)")).toBeInTheDocument(); + expect(screen.getByText("TPM Limit")).toBeInTheDocument(); + expect(screen.getByText("RPM Limit")).toBeInTheDocument(); + }); + + it("should render duration and grace period fields", () => { + renderWithProviders(); + + expect(screen.getByText("Expire Key")).toBeInTheDocument(); + expect(screen.getByText("Grace Period")).toBeInTheDocument(); + }); + + it("should display grace period recommendation text", () => { + renderWithProviders(); + expect( + screen.getByText("Recommended: 24h to 72h for production keys"), + ).toBeInTheDocument(); + }); + + it("should call regenerateKeyCall and show success view on successful regeneration", async () => { + const user = userEvent.setup(); + mockRegenerateKeyCall.mockResolvedValue({ + key: "sk-new-regenerated-key", + token: "new-token-hash", + }); + + renderWithProviders(); + + await user.click(screen.getByRole("button", { name: /Regenerate/ })); + + await waitFor(() => { + expect(mockRegenerateKeyCall).toHaveBeenCalledOnce(); + }); + + await waitFor(() => { + expect(screen.getByText("sk-new-regenerated-key")).toBeInTheDocument(); + }); + + expect(screen.getByText(/will not see it again/)).toBeInTheDocument(); + }); + + it("should show Close button after successful regeneration", async () => { + const user = userEvent.setup(); + mockRegenerateKeyCall.mockResolvedValue({ + key: "sk-new-regenerated-key", + token: "new-token-hash", + }); + + renderWithProviders(); + await user.click(screen.getByRole("button", { name: /Regenerate/ })); + + await waitFor(() => { + expect(screen.getByText("sk-new-regenerated-key")).toBeInTheDocument(); + }); + + // Should show Close buttons (footer + modal X), not Cancel/Regenerate + const closeButtons = screen.getAllByRole("button", { name: "Close" }); + expect(closeButtons.length).toBeGreaterThanOrEqual(1); + expect(screen.queryByRole("button", { name: "Cancel" })).not.toBeInTheDocument(); + expect(screen.queryByRole("button", { name: /Regenerate/ })).not.toBeInTheDocument(); + }); + + it("should show Copy Virtual Key button after successful regeneration", async () => { + const user = userEvent.setup(); + mockRegenerateKeyCall.mockResolvedValue({ + key: "sk-new-regenerated-key", + token: "new-token-hash", + }); + + renderWithProviders(); + await user.click(screen.getByRole("button", { name: /Regenerate/ })); + + await waitFor(() => { + expect(screen.getByRole("button", { name: /Copy/ })).toBeInTheDocument(); + }); + }); + + it("should call onKeyUpdate with updated data after successful regeneration", async () => { + const user = userEvent.setup(); + mockRegenerateKeyCall.mockResolvedValue({ + key: "sk-new-regenerated-key", + token: "new-token-hash", + }); + + renderWithProviders(); + await user.click(screen.getByRole("button", { name: /Regenerate/ })); + + await waitFor(() => { + expect(mockOnKeyUpdate).toHaveBeenCalledOnce(); + }); + + const updateCall = mockOnKeyUpdate.mock.calls[0][0]; + expect(updateCall.key_name).toBe("sk-new-regenerated-key"); + }); + + it("should display key alias in success view", async () => { + const user = userEvent.setup(); + mockRegenerateKeyCall.mockResolvedValue({ + key: "sk-new-regenerated-key", + token: "new-token-hash", + }); + + renderWithProviders(); + await user.click(screen.getByRole("button", { name: /Regenerate/ })); + + await waitFor(() => { + expect(screen.getByText("my-test-key")).toBeInTheDocument(); + }); + }); + + it("should display 'No alias set' when key has no alias", async () => { + const user = userEvent.setup(); + mockRegenerateKeyCall.mockResolvedValue({ + key: "sk-new-regenerated-key", + token: "new-token-hash", + }); + + renderWithProviders( + , + ); + await user.click(screen.getByRole("button", { name: /Regenerate/ })); + + await waitFor(() => { + expect(screen.getByText("No alias set")).toBeInTheDocument(); + }); + }); + + it("should not call regenerateKeyCall when selectedToken is null", async () => { + const user = userEvent.setup(); + renderWithProviders( + , + ); + + // The form shouldn't even be populated, but we check the button doesn't trigger a call + const regenerateBtn = screen.queryByRole("button", { name: /Regenerate/ }); + if (regenerateBtn) { + await user.click(regenerateBtn); + } + + expect(mockRegenerateKeyCall).not.toHaveBeenCalled(); + }); + + it("should pass the correct token identifier to regenerateKeyCall", async () => { + const user = userEvent.setup(); + mockRegenerateKeyCall.mockResolvedValue({ + key: "sk-new-key", + token: "new-hash", + }); + + renderWithProviders(); + await user.click(screen.getByRole("button", { name: /Regenerate/ })); + + await waitFor(() => { + expect(mockRegenerateKeyCall).toHaveBeenCalledWith( + "123", // accessToken from mocked useAuthorized + "token-hash-123", // selectedToken.token + expect.any(Object), + ); + }); + }); +}); diff --git a/ui/litellm-dashboard/src/components/organisms/regenerate_key_modal.tsx b/ui/litellm-dashboard/src/components/organisms/RegenerateKeyModal.tsx similarity index 60% rename from ui/litellm-dashboard/src/components/organisms/regenerate_key_modal.tsx rename to ui/litellm-dashboard/src/components/organisms/RegenerateKeyModal.tsx index 2fad101c20..e888713fe0 100644 --- a/ui/litellm-dashboard/src/components/organisms/regenerate_key_modal.tsx +++ b/ui/litellm-dashboard/src/components/organisms/RegenerateKeyModal.tsx @@ -1,6 +1,6 @@ import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; -import { Button, Col, Grid, Text, TextInput, Title } from "@tremor/react"; -import { Form, InputNumber, Modal } from "antd"; +import { CopyOutlined, SyncOutlined } from "@ant-design/icons"; +import { Alert, Button, Col, Form, Input, InputNumber, Modal, Row, Space, Typography } from "antd"; import { add } from "date-fns"; import { useEffect, useState } from "react"; import { CopyToClipboard } from "react-copy-to-clipboard"; @@ -8,6 +8,10 @@ import { KeyResponse } from "../key_team_helpers/key_list"; import NotificationManager from "../molecules/notifications_manager"; import { regenerateKeyCall } from "../networking"; +const { Text } = Typography; + + + interface RegenerateKeyModalProps { selectedToken: KeyResponse | null; visible: boolean; @@ -151,6 +155,7 @@ export function RegenerateKeyModal({ selectedToken, visible, onClose, onKeyUpdat title="Regenerate Virtual Key" open={visible} onCancel={handleClose} + width={520} footer={ regeneratedKey ? [ @@ -159,46 +164,69 @@ export function RegenerateKeyModal({ selectedToken, visible, onClose, onKeyUpdat , ] : [ - , - , + + + + , ] } > {regeneratedKey ? ( - - Regenerated Key - -

- Please replace your old key with the new key generated. For security reasons,{" "} - you will not be able to view it again through your LiteLLM account. If you lose this secret key, - you will need to generate a new one. -

- - - Key Alias: -
-
{selectedToken?.key_alias || "No alias set"}
-
- New Virtual Key: -
-
{regeneratedKey}
+
+ + +
+
Key Alias
+
+ {selectedToken?.key_alias || "No alias set"}
+
+ +
+ + {regeneratedKey} + NotificationManager.success("Virtual Key copied to clipboard")} > - + - - +
+
) : (
{ if ("duration" in changedValues) { setRegenerateFormData((prev: { duration?: string }) => ({ ...prev, duration: changedValues.duration })); @@ -206,41 +234,69 @@ export function RegenerateKeyModal({ selectedToken, visible, onClose, onKeyUpdat }} > - + - - - - - - - - - - - - -
- Current expiry: {selectedToken?.expires ? new Date(selectedToken.expires).toLocaleString() : "Never"} -
- {newExpiryTime &&
New expiry: {newExpiryTime}
} - - - -
- Recommended: 24h to 72h for production keys to allow seamless client migration. -
+ + + + + + + + + + + + + + + + + + + + + + + + Current expiry: {selectedToken?.expires ? new Date(selectedToken.expires).toLocaleString() : "Never"} + + {newExpiryTime && ( +
+ New expiry: {newExpiryTime} +
+ )} + + } + > + +
+ + + + Recommended: 24h to 72h for production keys + + } + rules={[ + { + pattern: /^(\d+(s|m|h|d|w|mo))?$/, + message: "Must be a duration like 30s, 30m, 24h, 2d, 1w, or 1mo", + }, + ]} + > + + + +
)} diff --git a/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx b/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx index 590864637a..abbc92f021 100644 --- a/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx @@ -196,7 +196,7 @@ vi.mock("lucide-react", async () => { }); // Heavy children -> async factories & local React -vi.mock("../organisms/regenerate_key_modal", async () => { +vi.mock("../organisms/RegenerateKeyModal", async () => { const React = await import("react"); function RegenerateKeyModal() { return null; 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 24e3e18b93..5b5e7722c0 100644 --- a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx @@ -20,7 +20,7 @@ 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 { RegenerateKeyModal } from "../organisms/RegenerateKeyModal"; import { parseErrorMessage } from "../shared/errorUtils"; import { KeyEditView } from "./key_edit_view";