[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.
2026-08-07 16:21:45 +00:00 · 2026-04-08 23:52:25 -07:00
parent 97f722f558
commit e42baeb5ab
5 changed files with 394 additions and 66 deletions
@@ -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 }) => {
@@ -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> = {}): 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(<RegenerateKeyModal {...defaultProps} />);
+    expect(screen.getByText("Regenerate Virtual Key")).toBeInTheDocument();
+  });
+
+  it("should not render the modal when visible is false", () => {
+    renderWithProviders(<RegenerateKeyModal {...defaultProps} visible={false} />);
+    expect(screen.queryByText("Regenerate Virtual Key")).not.toBeInTheDocument();
+  });
+
+  it("should display the form with pre-filled values", () => {
+    renderWithProviders(<RegenerateKeyModal {...defaultProps} />);
+
+    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(<RegenerateKeyModal {...defaultProps} />);
+    expect(screen.getByText(/Current expiry:/)).toBeInTheDocument();
+  });
+
+  it("should display 'Never' when token has no expires", () => {
+    renderWithProviders(
+      <RegenerateKeyModal
+        {...defaultProps}
+        selectedToken={makeToken({ expires: undefined })}
+      />,
+    );
+    expect(screen.getByText("Current expiry: Never")).toBeInTheDocument();
+  });
+
+  it("should show Cancel and Regenerate buttons in form view", () => {
+    renderWithProviders(<RegenerateKeyModal {...defaultProps} />);
+    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(<RegenerateKeyModal {...defaultProps} />);
+
+    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(<RegenerateKeyModal {...defaultProps} />);
+
+    await user.click(screen.getByRole("button", { name: "Close" }));
+    expect(mockOnClose).toHaveBeenCalledOnce();
+  });
+
+  it("should render form fields for budget and rate limits", () => {
+    renderWithProviders(<RegenerateKeyModal {...defaultProps} />);
+
+    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(<RegenerateKeyModal {...defaultProps} />);
+
+    expect(screen.getByText("Expire Key")).toBeInTheDocument();
+    expect(screen.getByText("Grace Period")).toBeInTheDocument();
+  });
+
+  it("should display grace period recommendation text", () => {
+    renderWithProviders(<RegenerateKeyModal {...defaultProps} />);
+    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(<RegenerateKeyModal {...defaultProps} />);
+
+    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(<RegenerateKeyModal {...defaultProps} />);
+    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(<RegenerateKeyModal {...defaultProps} />);
+    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(<RegenerateKeyModal {...defaultProps} />);
+    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(<RegenerateKeyModal {...defaultProps} />);
+    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(
+      <RegenerateKeyModal
+        {...defaultProps}
+        selectedToken={makeToken({ key_alias: undefined })}
+      />,
+    );
+    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(
+      <RegenerateKeyModal {...defaultProps} selectedToken={null} />,
+    );
+
+    // 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(<RegenerateKeyModal {...defaultProps} />);
+    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),
+      );
+    });
+  });
+});
@@ -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
              </Button>,
            ]
          : [
-              <Button key="cancel" onClick={handleClose} className="mr-2">
-                Cancel
-              </Button>,
-              <Button key="regenerate" onClick={handleRegenerateKey} disabled={isRegenerating}>
-                {isRegenerating ? "Regenerating..." : "Regenerate"}
-              </Button>,
+              <Space key="footer-actions">
+                <Button onClick={handleClose}>Cancel</Button>
+                <Button type="primary" icon={<SyncOutlined />} onClick={handleRegenerateKey} loading={isRegenerating}>
+                  Regenerate
+                </Button>
+              </Space>,
            ]
      }
    >
      {regeneratedKey ? (
-        <Grid numItems={1} className="gap-2 w-full">
-          <Title>Regenerated Key</Title>
-          <Col numColSpan={1}>
-            <p>
-              Please replace your old key with the new key generated. For security reasons,{" "}
-              <b>you will not be able to view it again</b> through your LiteLLM account. If you lose this secret key,
-              you will need to generate a new one.
-            </p>
-          </Col>
-          <Col numColSpan={1}>
-            <Text className="mt-3">Key Alias:</Text>
-            <div className="bg-gray-100 p-2 rounded mb-2">
-              <pre className="break-words whitespace-normal">{selectedToken?.key_alias || "No alias set"}</pre>
-            </div>
-            <Text className="mt-3">New Virtual Key:</Text>
-            <div className="bg-gray-100 p-2 rounded mb-2">
-              <pre className="break-words whitespace-normal">{regeneratedKey}</pre>
+        <div style={{ display: "flex", flexDirection: "column", gap: 16 }}>
+          <Alert
+            type="warning"
+            showIcon
+            banner
+            message="Save it now, you will not see it again"
+          />
+
+          <div>
+            <div style={{ fontSize: 12, color: "#8c8c8c", marginBottom: 2 }}>Key Alias</div>
+            <div style={{ fontSize: 14, color: "#595959" }}>
+              {selectedToken?.key_alias || "No alias set"}
            </div>
+          </div>
+
+          <div
+            style={{
+              display: "flex",
+              alignItems: "center",
+              gap: 8,
+              background: "#f5f5f5",
+              border: "1px solid #d9d9d9",
+              borderRadius: 6,
+              padding: "10px 12px",
+            }}
+          >
+            <code
+              style={{
+                flex: 1,
+                fontFamily: "SFMono-Regular, Consolas, 'Liberation Mono', Menlo, monospace",
+                fontSize: 13,
+                color: "#262626",
+                wordBreak: "break-all",
+                lineHeight: 1.5,
+              }}
+            >
+              {regeneratedKey}
+            </code>
            <CopyToClipboard
              text={regeneratedKey}
              onCopy={() => NotificationManager.success("Virtual Key copied to clipboard")}
            >
-              <Button className="mt-3">Copy Virtual Key</Button>
+              <Button type="primary" icon={<CopyOutlined />} size="small">
+                Copy
+              </Button>
            </CopyToClipboard>
-          </Col>
-        </Grid>
+          </div>
+        </div>
      ) : (
        <Form
          form={form}
          layout="vertical"
+          style={{ marginTop: 4 }}
          onValuesChange={(changedValues) => {
            if ("duration" in changedValues) {
              setRegenerateFormData((prev: { duration?: string }) => ({ ...prev, duration: changedValues.duration }));
@@ -206,41 +234,69 @@ export function RegenerateKeyModal({ selectedToken, visible, onClose, onKeyUpdat
          }}
        >
          <Form.Item name="key_alias" label="Key Alias">
-            <TextInput disabled={true} />
+            <Input disabled />
          </Form.Item>
-          <Form.Item name="max_budget" label="Max Budget (USD)">
-            <InputNumber step={0.01} precision={2} style={{ width: "100%" }} />
-          </Form.Item>
-          <Form.Item name="tpm_limit" label="TPM Limit">
-            <InputNumber style={{ width: "100%" }} />
-          </Form.Item>
-          <Form.Item name="rpm_limit" label="RPM Limit">
-            <InputNumber style={{ width: "100%" }} />
-          </Form.Item>
-          <Form.Item name="duration" label="Expire Key (eg: 30s, 30h, 30d)" className="mt-8">
-            <TextInput placeholder="" />
-          </Form.Item>
-          <div className="mt-2 text-sm text-gray-500">
-            Current expiry: {selectedToken?.expires ? new Date(selectedToken.expires).toLocaleString() : "Never"}
-          </div>
-          {newExpiryTime && <div className="mt-2 text-sm text-green-600">New expiry: {newExpiryTime}</div>}
-          <Form.Item
-            name="grace_period"
-            label="Grace Period (eg: 24h, 2d)"
-            tooltip="Keep the old key valid for this duration after rotation. Both keys work during this period for seamless cutover. Empty = immediate revoke."
-            className="mt-8"
-            rules={[
-              {
-                pattern: /^(\d+(s|m|h|d|w|mo))?$/,
-                message: "Must be a duration like 30s, 30m, 24h, 2d, 1w, or 1mo",
-              },
-            ]}
-          >
-            <TextInput placeholder="e.g. 24h, 2d (empty = immediate revoke)" />
-          </Form.Item>
-          <div className="mt-2 text-sm text-gray-500">
-            Recommended: 24h to 72h for production keys to allow seamless client migration.
-          </div>
+
+          <Row gutter={12}>
+            <Col span={8}>
+              <Form.Item name="max_budget" label="Max Budget (USD)">
+                <InputNumber step={0.01} precision={2} style={{ width: "100%" }} />
+              </Form.Item>
+            </Col>
+            <Col span={8}>
+              <Form.Item name="tpm_limit" label="TPM Limit">
+                <InputNumber style={{ width: "100%" }} />
+              </Form.Item>
+            </Col>
+            <Col span={8}>
+              <Form.Item name="rpm_limit" label="RPM Limit">
+                <InputNumber style={{ width: "100%" }} />
+              </Form.Item>
+            </Col>
+          </Row>
+
+          <Row gutter={12}>
+            <Col span={12}>
+              <Form.Item
+                name="duration"
+                label="Expire Key"
+                extra={
+                  <>
+                    <Text type="secondary" style={{ fontSize: 12 }}>
+                      Current expiry: {selectedToken?.expires ? new Date(selectedToken.expires).toLocaleString() : "Never"}
+                    </Text>
+                    {newExpiryTime && (
+                      <div>
+                        <Text style={{ fontSize: 12, color: "#52c41a" }}>New expiry: {newExpiryTime}</Text>
+                      </div>
+                    )}
+                  </>
+                }
+              >
+                <Input placeholder="e.g. 30s, 30h, 30d" />
+              </Form.Item>
+            </Col>
+            <Col span={12}>
+              <Form.Item
+                name="grace_period"
+                label="Grace Period"
+                tooltip="Keep the old key valid for this duration after rotation. Both keys work during this period for seamless cutover. Empty = immediate revoke."
+                extra={
+                  <Text type="secondary" style={{ fontSize: 12 }}>
+                    Recommended: 24h to 72h for production keys
+                  </Text>
+                }
+                rules={[
+                  {
+                    pattern: /^(\d+(s|m|h|d|w|mo))?$/,
+                    message: "Must be a duration like 30s, 30m, 24h, 2d, 1w, or 1mo",
+                  },
+                ]}
+              >
+                <Input placeholder="e.g. 24h, 2d" />
+              </Form.Item>
+            </Col>
+          </Row>
        </Form>
      )}
    </Modal>
@@ -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;
@@ -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";