From e42baeb5abf2ae37d64117dc70a9dc0138daf4d3 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 8 Apr 2026 23:52:25 -0700 Subject: [PATCH] [Refactor] UI - Virtual Keys: migrate regenerate key modal to AntD Replace Tremor components in the regenerate key modal with Ant Design equivalents and move the component to a new PascalCase file. The form layout now uses Row/Col to place Max Budget, TPM Limit, and RPM Limit on one row and Expire Key with Grace Period on another, reducing the vertical footprint. The success view shows an Alert banner, the key alias as secondary context, and the regenerated key in a monospace block with an inline primary Copy button. Also adds unit tests for the new component and updates the existing Playwright spec to match the new banner and button text. --- .../e2e_tests/tests/proxy-admin/keys.spec.ts | 5 +- .../organisms/RegenerateKeyModal.test.tsx | 271 ++++++++++++++++++ ...e_key_modal.tsx => RegenerateKeyModal.tsx} | 180 ++++++++---- .../KeyInfoView.handleKeyUpdate.test.tsx | 2 +- .../components/templates/key_info_view.tsx | 2 +- 5 files changed, 394 insertions(+), 66 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/organisms/RegenerateKeyModal.test.tsx rename ui/litellm-dashboard/src/components/organisms/{regenerate_key_modal.tsx => RegenerateKeyModal.tsx} (60%) 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";