Merge pull request #25406 from BerriAI/litellm_regen_key_modal_antd

[Refactor] UI - Virtual Keys: migrate regenerate key modal to AntD

ui/litellm-dashboard/e2e_tests/tests/proxy-admin/keys.spec.ts

@@ -61,10 +61,15 @@ test.describe("Proxy Admin - Keys", () => {
     await expect(page.getByText("Back to Keys")).toBeVisible({ timeout: 10_000 });
 
     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 });
+    // Scope to the modal — the Regenerate button has an icon whose aria-label
+    // ("sync") is concatenated into the button's accessible name, and the
+    // "Regenerate Key" button is still in the DOM behind the modal.
+    const modal = page.locator(".ant-modal:visible");
+    await modal.getByRole("button", { name: /Regenerate/ }).click();
+
+    // Success view shows a Copy button in the footer (text varies between modal versions)
+    await expect(modal.getByRole("button", { name: /Copy.*Key/ })).toBeVisible({ timeout: 20_000 });
   });
 
   test("Update key TPM and RPM limits", async ({ page }) => {

ui/litellm-dashboard/src/components/organisms/RegenerateKeyModal.test.tsx

@@ -0,0 +1,359 @@
+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),
+}));
+
+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 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 Key/ })).toBeInTheDocument();
+    });
+  });
+
+  it("should swap the Copy Key button to 'Copied' after clicking it", 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/ }));
+
+    const copyButton = await screen.findByRole("button", { name: /Copy Key/ });
+    await user.click(copyButton);
+
+    await waitFor(() => {
+      expect(screen.getByRole("button", { name: /Copied/ })).toBeInTheDocument();
+    });
+    expect(screen.queryByRole("button", { name: /Copy Key/ })).not.toBeInTheDocument();
+  });
+
+  it("should display the 'Virtual Key' label above the key in the 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("Virtual Key")).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.each([
+    ["30s", /New expiry:/],
+    ["15m", /New expiry:/],
+    ["2h", /New expiry:/],
+    ["7d", /New expiry:/],
+    ["2w", /New expiry:/],
+    ["1mo", /New expiry:/],
+  ])("should compute a new expiry preview for duration '%s'", async (durationInput, expected) => {
+    const user = userEvent.setup();
+    renderWithProviders(<RegenerateKeyModal {...defaultProps} />);
+
+    const durationField = screen.getByPlaceholderText("e.g. 30s, 30h, 30d");
+    await user.clear(durationField);
+    await user.type(durationField, durationInput);
+
+    await waitFor(() => {
+      expect(screen.getByText(expected)).toBeInTheDocument();
+    });
+  });
+
+  it("should fall back to the previous expiry when duration is unparseable", async () => {
+    // Regression: if calculateNewExpiryTime returns null (unrecognised suffix),
+    // the payload should fall back to the previous expires rather than null.
+    const user = userEvent.setup();
+    const previousExpires = "2026-12-31T00:00:00Z";
+    mockRegenerateKeyCall.mockResolvedValue({
+      key: "sk-new-regenerated-key",
+      token: "new-token-hash",
+    });
+
+    renderWithProviders(
+      <RegenerateKeyModal {...defaultProps} selectedToken={makeToken({ expires: previousExpires })} />,
+    );
+
+    const durationField = screen.getByPlaceholderText("e.g. 30s, 30h, 30d");
+    await user.clear(durationField);
+    await user.type(durationField, "bogus");
+
+    await user.click(screen.getByRole("button", { name: /Regenerate/ }));
+
+    await waitFor(() => {
+      expect(mockOnKeyUpdate).toHaveBeenCalledOnce();
+    });
+
+    const updateCall = mockOnKeyUpdate.mock.calls[0][0];
+    expect(updateCall.expires).toBe(previousExpires);
+  });
+
+  it("should pass form values to onKeyUpdate even when the API echoes back different limits", async () => {
+    // Regression: when the regenerate endpoint returns GenerateKeyResponse, it echoes
+    // back the existing max_budget / tpm_limit / rpm_limit. The modal must prefer the
+    // values the user just submitted, not whatever the server echoes.
+    const user = userEvent.setup();
+    mockRegenerateKeyCall.mockResolvedValue({
+      key: "sk-new-regenerated-key",
+      token: "new-token-hash",
+      // stale values echoed from the server
+      max_budget: 9999,
+      tpm_limit: 9999,
+      rpm_limit: 9999,
+    });
+
+    renderWithProviders(<RegenerateKeyModal {...defaultProps} />);
+    await user.click(screen.getByRole("button", { name: /Regenerate/ }));
+
+    await waitFor(() => {
+      expect(mockOnKeyUpdate).toHaveBeenCalledOnce();
+    });
+
+    const updateCall = mockOnKeyUpdate.mock.calls[0][0];
+    // The form's pre-filled values (from makeToken) must win over the API echo.
+    expect(updateCall.max_budget).toBe(100);
+    expect(updateCall.tpm_limit).toBe(5000);
+    expect(updateCall.rpm_limit).toBe(500);
+  });
+
+  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),
+      );
+    });
+  });
+});

ui/litellm-dashboard/src/components/organisms/RegenerateKeyModal.tsx

@@ -0,0 +1,279 @@
+import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
+import { CheckOutlined, CopyOutlined, SyncOutlined } from "@ant-design/icons";
+import { Alert, Button, Col, Flex, 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";
+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;
+  onClose: () => void;
+  onKeyUpdate?: (updatedKeyData: Partial<KeyResponse>) => void;
+}
+
+export function RegenerateKeyModal({ selectedToken, visible, onClose, onKeyUpdate }: RegenerateKeyModalProps) {
+  const { accessToken } = useAuthorized();
+  const [form] = Form.useForm();
+  const [regeneratedKey, setRegeneratedKey] = useState<string | null>(null);
+  const [regenerateFormData, setRegenerateFormData] = useState<any>(null);
+  const [newExpiryTime, setNewExpiryTime] = useState<string | null>(null);
+  const [isRegenerating, setIsRegenerating] = useState(false);
+  const [copied, setCopied] = useState(false);
+
+  useEffect(() => {
+    if (visible && selectedToken && accessToken) {
+      form.setFieldsValue({
+        key_alias: selectedToken.key_alias,
+        max_budget: selectedToken.max_budget,
+        tpm_limit: selectedToken.tpm_limit,
+        rpm_limit: selectedToken.rpm_limit,
+        duration: selectedToken.duration || "",
+        grace_period: "",
+      });
+    }
+  }, [visible, selectedToken, form, accessToken]);
+
+  const calculateNewExpiryTime = (duration: string | undefined): string | null => {
+    if (!duration) return null;
+
+    try {
+      const amount = parseInt(duration);
+      if (Number.isNaN(amount)) {
+        throw new Error("Invalid duration format");
+      }
+      const now = new Date();
+      // Check "mo" before "m" to avoid a false prefix match (e.g. "1mo" → minutes).
+      let newExpiry: Date;
+      if (duration.endsWith("mo")) {
+        newExpiry = add(now, { months: amount });
+      } else if (duration.endsWith("s")) {
+        newExpiry = add(now, { seconds: amount });
+      } else if (duration.endsWith("m")) {
+        newExpiry = add(now, { minutes: amount });
+      } else if (duration.endsWith("h")) {
+        newExpiry = add(now, { hours: amount });
+      } else if (duration.endsWith("d")) {
+        newExpiry = add(now, { days: amount });
+      } else if (duration.endsWith("w")) {
+        newExpiry = add(now, { weeks: amount });
+      } else {
+        throw new Error("Invalid duration format");
+      }
+
+      return newExpiry.toLocaleString();
+    } catch (error) {
+      return null;
+    }
+  };
+
+  useEffect(() => {
+    if (regenerateFormData?.duration) {
+      setNewExpiryTime(calculateNewExpiryTime(regenerateFormData.duration));
+    } else {
+      setNewExpiryTime(null);
+    }
+  }, [regenerateFormData?.duration]);
+
+  const handleRegenerateKey = async () => {
+    if (!selectedToken || !accessToken) return;
+
+    setIsRegenerating(true);
+    try {
+      const formValues = await form.validateFields();
+
+      const response = await regenerateKeyCall(
+        accessToken,
+        selectedToken.token || selectedToken.token_id,
+        formValues,
+      );
+      setRegeneratedKey(response.key);
+      NotificationManager.success("Virtual Key regenerated successfully");
+
+      // Build the update payload. Spread the API response first so any new
+      // fields it returns (new token, timestamps, etc.) are captured, then
+      // override with the explicit form values — the user's just-submitted
+      // edits must win over whatever the API echoes back.
+      const updatedKeyData: Partial<KeyResponse> = {
+        ...response,
+        token: response.token || response.key_id || selectedToken.token,
+        key_name: response.key,
+        max_budget: formValues.max_budget,
+        tpm_limit: formValues.tpm_limit,
+        rpm_limit: formValues.rpm_limit,
+        expires: formValues.duration
+          ? (calculateNewExpiryTime(formValues.duration) ?? selectedToken.expires)
+          : selectedToken.expires,
+      };
+
+      // Update the parent component with new key data
+      if (onKeyUpdate) {
+        onKeyUpdate(updatedKeyData);
+      }
+
+      setIsRegenerating(false);
+    } catch (error) {
+      console.error("Error regenerating key:", error);
+      NotificationManager.fromBackend(error);
+      setIsRegenerating(false); // Reset regenerating state on error
+    }
+  };
+
+  const handleClose = () => {
+    setRegeneratedKey(null);
+    setIsRegenerating(false);
+    setCopied(false);
+    form.resetFields();
+    onClose();
+  };
+
+  const handleCopyKey = () => {
+    setCopied(true);
+  };
+
+  return (
+    <Modal
+      title="Regenerate Virtual Key"
+      open={visible}
+      onCancel={handleClose}
+      width={520}
+      maskClosable={false}
+      footer={
+        regeneratedKey
+          ? [
+              <Space key="footer-actions">
+                <Button onClick={handleClose}>Close</Button>
+                <CopyToClipboard text={regeneratedKey} onCopy={handleCopyKey}>
+                  <Button type="primary" icon={copied ? <CheckOutlined /> : <CopyOutlined />}>
+                    {copied ? "Copied" : "Copy Key"}
+                  </Button>
+                </CopyToClipboard>
+              </Space>,
+            ]
+          : [
+              <Space key="footer-actions">
+                <Button onClick={handleClose}>Cancel</Button>
+                <Button type="primary" icon={<SyncOutlined />} onClick={handleRegenerateKey} loading={isRegenerating}>
+                  Regenerate
+                </Button>
+              </Space>,
+            ]
+      }
+    >
+      {regeneratedKey ? (
+        <Flex vertical gap="middle">
+          <Alert type="warning" showIcon message="Save it now, you will not see it again" />
+
+          <Flex vertical gap={2}>
+            <Text type="secondary" style={{ fontSize: 12 }}>
+              Key Alias
+            </Text>
+            <Text>{selectedToken?.key_alias || "No alias set"}</Text>
+          </Flex>
+
+          <Flex vertical gap={6}>
+            <Text type="secondary" style={{ fontSize: 12 }}>
+              Virtual Key
+            </Text>
+            <div
+              style={{
+                background: "#f5f5f5",
+                border: "1px solid #e8e8e8",
+                borderRadius: 6,
+                padding: "14px 16px",
+                fontFamily: "SFMono-Regular, Consolas, 'Liberation Mono', Menlo, monospace",
+                fontSize: 16,
+                wordBreak: "break-all",
+                color: "#262626",
+              }}
+            >
+              {regeneratedKey}
+            </div>
+          </Flex>
+        </Flex>
+      ) : (
+        <Form
+          form={form}
+          layout="vertical"
+          style={{ marginTop: 4 }}
+          onValuesChange={(changedValues) => {
+            if ("duration" in changedValues) {
+              setRegenerateFormData((prev: { duration?: string }) => ({ ...prev, duration: changedValues.duration }));
+            }
+          }}
+        >
+          <Form.Item name="key_alias" label="Key Alias">
+            <Input disabled />
+          </Form.Item>
+
+          <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={
+                  <Flex vertical gap={2}>
+                    <Text type="secondary" style={{ fontSize: 12 }}>
+                      Current expiry:{" "}
+                      {selectedToken?.expires ? new Date(selectedToken.expires).toLocaleString() : "Never"}
+                    </Text>
+                    {newExpiryTime && (
+                      <Text type="success" style={{ fontSize: 12 }}>
+                        New expiry: {newExpiryTime}
+                      </Text>
+                    )}
+                  </Flex>
+                }
+              >
+                <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>
+  );
+}

ui/litellm-dashboard/src/components/organisms/regenerate_key_modal.tsx

@@ -1,248 +0,0 @@
-import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
-import { Button, Col, Grid, Text, TextInput, Title } from "@tremor/react";
-import { Form, InputNumber, Modal } from "antd";
-import { add } from "date-fns";
-import { useEffect, useState } from "react";
-import { CopyToClipboard } from "react-copy-to-clipboard";
-import { KeyResponse } from "../key_team_helpers/key_list";
-import NotificationManager from "../molecules/notifications_manager";
-import { regenerateKeyCall } from "../networking";
-
-interface RegenerateKeyModalProps {
-  selectedToken: KeyResponse | null;
-  visible: boolean;
-  onClose: () => void;
-  onKeyUpdate?: (updatedKeyData: Partial<KeyResponse>) => void;
-}
-
-export function RegenerateKeyModal({ selectedToken, visible, onClose, onKeyUpdate }: RegenerateKeyModalProps) {
-  const { accessToken } = useAuthorized();
-  const [form] = Form.useForm();
-  const [regeneratedKey, setRegeneratedKey] = useState<string | null>(null);
-  const [regenerateFormData, setRegenerateFormData] = useState<any>(null);
-  const [newExpiryTime, setNewExpiryTime] = useState<string | null>(null);
-  const [isRegenerating, setIsRegenerating] = useState(false);
-
-  // Track whether this is the user's own authentication key
-  const [isOwnKey, setIsOwnKey] = useState<boolean>(false);
-
-  // Keep track of the current valid access token locally
-  const [currentAccessToken, setCurrentAccessToken] = useState<string | null>(null);
-
-  useEffect(() => {
-    if (visible && selectedToken && accessToken) {
-      form.setFieldsValue({
-        key_alias: selectedToken.key_alias,
-        max_budget: selectedToken.max_budget,
-        tpm_limit: selectedToken.tpm_limit,
-        rpm_limit: selectedToken.rpm_limit,
-        duration: selectedToken.duration || "",
-        grace_period: "",
-      });
-
-      // Initialize the current access token
-      setCurrentAccessToken(accessToken);
-
-      // Check if this is the user's own authentication key by comparing the key values
-      const isUserOwnKey = selectedToken.key_name === accessToken;
-      setIsOwnKey(isUserOwnKey);
-    }
-  }, [visible, selectedToken, form, accessToken]);
-
-  useEffect(() => {
-    if (!visible) {
-      // Reset states when modal is closed
-      setRegeneratedKey(null);
-      setIsRegenerating(false);
-      setIsOwnKey(false);
-      setCurrentAccessToken(null);
-      form.resetFields();
-    }
-  }, [visible, form]);
-
-  const calculateNewExpiryTime = (duration: string | undefined): string | null => {
-    if (!duration) return null;
-
-    try {
-      const now = new Date();
-      let newExpiry: Date;
-
-      if (duration.endsWith("s")) {
-        newExpiry = add(now, { seconds: parseInt(duration) });
-      } else if (duration.endsWith("h")) {
-        newExpiry = add(now, { hours: parseInt(duration) });
-      } else if (duration.endsWith("d")) {
-        newExpiry = add(now, { days: parseInt(duration) });
-      } else {
-        throw new Error("Invalid duration format");
-      }
-
-      return newExpiry.toLocaleString();
-    } catch (error) {
-      return null;
-    }
-  };
-
-  useEffect(() => {
-    if (regenerateFormData?.duration) {
-      setNewExpiryTime(calculateNewExpiryTime(regenerateFormData.duration));
-    } else {
-      setNewExpiryTime(null);
-    }
-  }, [regenerateFormData?.duration]);
-
-  const handleRegenerateKey = async () => {
-    if (!selectedToken || !currentAccessToken) return;
-
-    setIsRegenerating(true);
-    try {
-      const formValues = await form.validateFields();
-
-      // Use the current access token for the API call
-      const response = await regenerateKeyCall(
-        currentAccessToken,
-        selectedToken.token || selectedToken.token_id,
-        formValues,
-      );
-      setRegeneratedKey(response.key);
-      NotificationManager.success("Virtual Key regenerated successfully");
-
-      console.log("Full regenerate response:", response); // Debug log to see what's returned
-
-      // Create updated key data with ALL new values from the response
-      const updatedKeyData: Partial<KeyResponse> = {
-        // Use the new token/key ID from the response (this is what was missing!)
-        token: response.token || response.key_id || selectedToken.token, // Try different possible field names
-        key_name: response.key, // This is the new secret key string
-        max_budget: formValues.max_budget,
-        tpm_limit: formValues.tpm_limit,
-        rpm_limit: formValues.rpm_limit,
-        expires: formValues.duration ? calculateNewExpiryTime(formValues.duration) : selectedToken.expires,
-        // Include any other fields that might be returned by the API
-        ...response, // Spread the entire response to capture all updated fields
-      };
-
-      console.log("Updated key data with new token:", updatedKeyData); // Debug log
-
-      // Update the parent component with new key data
-      if (onKeyUpdate) {
-        onKeyUpdate(updatedKeyData);
-      }
-
-      setIsRegenerating(false);
-    } catch (error) {
-      console.error("Error regenerating key:", error);
-      NotificationManager.fromBackend(error);
-      setIsRegenerating(false); // Reset regenerating state on error
-    }
-  };
-
-  const handleClose = () => {
-    setRegeneratedKey(null);
-    setIsRegenerating(false);
-    setIsOwnKey(false);
-    setCurrentAccessToken(null);
-    form.resetFields();
-    onClose();
-  };
-
-  return (
-    <Modal
-      title="Regenerate Virtual Key"
-      open={visible}
-      onCancel={handleClose}
-      footer={
-        regeneratedKey
-          ? [
-              <Button key="close" onClick={handleClose}>
-                Close
-              </Button>,
-            ]
-          : [
-              <Button key="cancel" onClick={handleClose} className="mr-2">
-                Cancel
-              </Button>,
-              <Button key="regenerate" onClick={handleRegenerateKey} disabled={isRegenerating}>
-                {isRegenerating ? "Regenerating..." : "Regenerate"}
-              </Button>,
-            ]
-      }
-    >
-      {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>
-            <CopyToClipboard
-              text={regeneratedKey}
-              onCopy={() => NotificationManager.success("Virtual Key copied to clipboard")}
-            >
-              <Button className="mt-3">Copy Virtual Key</Button>
-            </CopyToClipboard>
-          </Col>
-        </Grid>
-      ) : (
-        <Form
-          form={form}
-          layout="vertical"
-          onValuesChange={(changedValues) => {
-            if ("duration" in changedValues) {
-              setRegenerateFormData((prev: { duration?: string }) => ({ ...prev, duration: changedValues.duration }));
-            }
-          }}
-        >
-          <Form.Item name="key_alias" label="Key Alias">
-            <TextInput disabled={true} />
-          </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>
-        </Form>
-      )}
-    </Modal>
-  );
-}

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;

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";