diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.ts index 8ae4d76ff5..cd3df1a382 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.ts @@ -4,7 +4,7 @@ import { keyListCall } from "@/components/networking"; import { KeyResponse } from "@/components/key_team_helpers/key_list"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; -const keyKeys = createQueryKeys("keys"); +export const keyKeys = createQueryKeys("keys"); export interface KeysResponse { keys: KeyResponse[]; diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.test.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.test.tsx index 08c05f5fad..bd5e0ee1cb 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.test.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.test.tsx @@ -1,5 +1,5 @@ -import { render, screen } from "@testing-library/react"; import { beforeEach, describe, expect, it, vi } from "vitest"; +import { renderWithProviders, screen } from "../../../tests/test-utils"; import CreateKey from "./create_key_button"; const mockKeyCreateCall = vi.fn().mockResolvedValue({ @@ -58,7 +58,7 @@ describe("CreateKey", () => { }); it("should render the CreateKey component", () => { - render(); + renderWithProviders(); expect(screen.getByRole("button", { name: /create new key/i })).toBeInTheDocument(); }); }); diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx index 3cd8a04e06..8ad1555dd8 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx @@ -7,6 +7,8 @@ import { Button as Button2, Form, Input, Modal, Radio, Select, Switch, Tooltip } import debounce from "lodash/debounce"; import React, { useCallback, useEffect, useState } from "react"; import { CopyToClipboard } from "react-copy-to-clipboard"; +import { useQueryClient } from "@tanstack/react-query"; +import { createQueryKeys } from "@/app/(dashboard)/hooks/common/queryKeysFactory"; import { rolesWithWriteAccess } from "../../utils/roles"; import AgentSelector from "../agent_management/AgentSelector"; import { mapDisplayToInternalNames } from "../callback_info_helpers"; @@ -36,6 +38,7 @@ import { } from "../networking"; import NumericalInput from "../shared/numerical_input"; import VectorStoreSelector from "../vector_store_management/VectorStoreSelector"; +import { keyKeys } from "@/app/(dashboard)/hooks/keys/useKeys"; const { Option } = Select; @@ -136,6 +139,7 @@ export const fetchUserModels = async ( */ const CreateKey: React.FC = ({ team, teams, data, addKey }) => { const { accessToken, userId: userID, userRole, premiumUser } = useAuthorized(); + const queryClient = useQueryClient(); const [form] = Form.useForm(); const [isModalVisible, setIsModalVisible] = useState(false); const [apiKey, setApiKey] = useState(null); @@ -392,6 +396,10 @@ const CreateKey: React.FC = ({ team, teams, data, addKey }) => { // Also directly update the keys list in VirtualKeysTable without an API call addKey(response); + // Invalidate and refetch all keys list queries to update the table + // This will trigger a refetch of all key list queries regardless of pagination + queryClient.invalidateQueries({ queryKey: keyKeys.lists() }); + setApiKey(response["key"]); setSoftBudget(response["soft_budget"]); NotificationsManager.success("Virtual Key Created");