diff --git a/ui/litellm-dashboard/e2e_tests/fixtures/migratedPages.ts b/ui/litellm-dashboard/e2e_tests/fixtures/migratedPages.ts index a911bbb17c4..00c4982529e 100644 --- a/ui/litellm-dashboard/e2e_tests/fixtures/migratedPages.ts +++ b/ui/litellm-dashboard/e2e_tests/fixtures/migratedPages.ts @@ -40,6 +40,7 @@ export const MIGRATED_E2E_PAGES: Record = { agents: "agents", "router-settings": "router-settings", users: "users", + organizations: "organizations", }; export const MIGRATED_E2E_SEGMENTS: string[] = [...new Set(Object.values(MIGRATED_E2E_PAGES))]; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.test.ts index 2539cc63f95..0bcc37d1389 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.test.ts @@ -8,6 +8,7 @@ import { useModelHub, useModelsInfo, useSelectedTeamModels, + useUserModels, type AllProxyModelsResponse, type PaginatedModelInfoResponse, type ProxyModel, @@ -480,6 +481,70 @@ describe("useAllProxyModels", () => { }); }); +describe("useUserModels", () => { + let queryClient: QueryClient; + + beforeEach(() => { + queryClient = new QueryClient({ + defaultOptions: { + queries: { + retry: false, + }, + }, + }); + + vi.clearAllMocks(); + + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userId: "test-user-id", + userRole: "Admin", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + }); + + const wrapper = ({ children }: { children: ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); + + it("maps the available-models response to a list of model ids", async () => { + (modelAvailableCall as any).mockResolvedValue({ + data: [{ id: "gpt-4" }, { id: "claude-3-opus" }], + }); + + const { result } = renderHook(() => useUserModels(), { wrapper }); + + await waitFor(() => { + expect(result.current.isSuccess).toBe(true); + }); + + expect(result.current.data).toEqual(["gpt-4", "claude-3-opus"]); + expect(modelAvailableCall).toHaveBeenCalledWith("test-access-token", "test-user-id", "Admin"); + expect(modelAvailableCall).toHaveBeenCalledTimes(1); + }); + + it("should not execute query when accessToken is missing", () => { + mockUseAuthorized.mockReturnValue({ + accessToken: null, + userId: "test-user-id", + userRole: "Admin", + token: null, + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useUserModels(), { wrapper }); + + expect(result.current.isFetched).toBe(false); + expect(modelAvailableCall).not.toHaveBeenCalled(); + }); +}); + describe("useSelectedTeamModels", () => { let queryClient: QueryClient; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts index c997f679b2e..113d1616e62 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts @@ -1,4 +1,4 @@ -import { useQuery, useInfiniteQuery } from "@tanstack/react-query"; +import { useQuery, useInfiniteQuery, UseQueryResult } from "@tanstack/react-query"; import { createQueryKeys } from "../common/queryKeysFactory"; import { modelInfoCall, modelHubCall, modelAvailableCall } from "@/components/networking"; import useAuthorized from "../useAuthorized"; @@ -27,6 +27,7 @@ const modelHubKeys = createQueryKeys("modelHub"); const allProxyModelsKeys = createQueryKeys("allProxyModels"); const selectedTeamModelsKeys = createQueryKeys("selectedTeamModels"); const infiniteModelKeys = createQueryKeys("infiniteModels"); +const userModelsKeys = createQueryKeys("userModels"); export const useModelsInfo = ( page: number = 1, @@ -76,6 +77,18 @@ export const useAllProxyModels = () => { }); }; +export const useUserModels = (): UseQueryResult => { + const { accessToken, userId, userRole } = useAuthorized(); + return useQuery({ + queryKey: userModelsKeys.list({}), + queryFn: async () => { + const response = await modelAvailableCall(accessToken!, userId!, userRole!); + return response["data"].map((model: { id: string }) => model.id); + }, + enabled: Boolean(accessToken && userId && userRole), + }); +}; + export const useSelectedTeamModels = (teamID: string | null) => { const { accessToken, userId, userRole } = useAuthorized(); return useQuery({ diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/organizations/useOrganizations.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/organizations/useOrganizations.test.ts index 66c005f37c4..960afe7392c 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/organizations/useOrganizations.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/organizations/useOrganizations.test.ts @@ -2,13 +2,14 @@ import { describe, it, expect, vi, beforeEach } from "vitest"; import { renderHook, waitFor } from "@testing-library/react"; import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; import React, { ReactNode } from "react"; -import { useOrganizations } from "./useOrganizations"; -import { organizationListCall } from "@/components/networking"; +import { organizationKeys, useOrganization, useOrganizations } from "./useOrganizations"; +import { organizationInfoCall, organizationListCall } from "@/components/networking"; import type { Organization } from "@/components/networking"; // Mock the networking function vi.mock("@/components/networking", () => ({ organizationListCall: vi.fn(), + organizationInfoCall: vi.fn(), })); // Mock useAuthorized hook - we can override this in individual tests @@ -107,7 +108,7 @@ describe("useOrganizations", () => { expect(result.current.data).toEqual(mockOrganizations); expect(result.current.error).toBeNull(); - expect(organizationListCall).toHaveBeenCalledWith("test-access-token"); + expect(organizationListCall).toHaveBeenCalledWith("test-access-token", null, null); expect(organizationListCall).toHaveBeenCalledTimes(1); }); @@ -131,10 +132,47 @@ describe("useOrganizations", () => { expect(result.current.error).toEqual(testError); expect(result.current.data).toBeUndefined(); - expect(organizationListCall).toHaveBeenCalledWith("test-access-token"); + expect(organizationListCall).toHaveBeenCalledWith("test-access-token", null, null); expect(organizationListCall).toHaveBeenCalledTimes(1); }); + it("passes org_id and org_alias filters to organizationListCall and caches separately from the unfiltered list", async () => { + (organizationListCall as any).mockResolvedValue(mockOrganizations); + + const { result } = renderHook(() => useOrganizations({ org_id: "org-1", org_alias: "Test Organization 1" }), { + wrapper, + }); + + await waitFor(() => { + expect(result.current.isSuccess).toBe(true); + }); + + expect(organizationListCall).toHaveBeenCalledWith("test-access-token", "org-1", "Test Organization 1"); + + (organizationListCall as any).mockResolvedValue([]); + const { result: unfiltered } = renderHook(() => useOrganizations(), { wrapper }); + + await waitFor(() => { + expect(unfiltered.current.isSuccess).toBe(true); + }); + + expect(organizationListCall).toHaveBeenLastCalledWith("test-access-token", null, null); + expect(organizationListCall).toHaveBeenCalledTimes(2); + }); + + it("treats empty-string filters as no filters, writing to the unfiltered cache entry", async () => { + (organizationListCall as any).mockResolvedValue(mockOrganizations); + + const { result } = renderHook(() => useOrganizations({ org_id: "", org_alias: "" }), { wrapper }); + + await waitFor(() => { + expect(result.current.isSuccess).toBe(true); + }); + + expect(organizationListCall).toHaveBeenCalledWith("test-access-token", null, null); + expect(queryClient.getQueryData(organizationKeys.list({}))).toEqual(mockOrganizations); + }); + it("should not execute query when accessToken is missing", async () => { // Mock missing accessToken mockUseAuthorized.mockReturnValue({ @@ -243,7 +281,7 @@ describe("useOrganizations", () => { expect(result.current.isLoading).toBe(false); }); - expect(organizationListCall).toHaveBeenCalledWith("test-access-token"); + expect(organizationListCall).toHaveBeenCalledWith("test-access-token", null, null); expect(organizationListCall).toHaveBeenCalledTimes(1); }); @@ -260,7 +298,7 @@ describe("useOrganizations", () => { }); expect(result.current.data).toEqual([]); - expect(organizationListCall).toHaveBeenCalledWith("test-access-token"); + expect(organizationListCall).toHaveBeenCalledWith("test-access-token", null, null); }); it("should handle network timeout error", async () => { @@ -280,3 +318,58 @@ describe("useOrganizations", () => { expect(result.current.data).toBeUndefined(); }); }); + +describe("useOrganization", () => { + let queryClient: QueryClient; + + beforeEach(() => { + queryClient = new QueryClient({ + defaultOptions: { queries: { retry: false } }, + }); + + vi.clearAllMocks(); + + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userId: "test-user-id", + userRole: "Admin", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + }); + + const wrapper = ({ children }: { children: ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); + + it("seeds initialData from a filtered list cache entry so the detail renders without a loading state", () => { + (organizationInfoCall as any).mockResolvedValue(mockOrganizations[1]); + // Only a filtered list was ever fetched; the unfiltered list({}) entry stays empty. + queryClient.setQueryData(organizationKeys.list({ filters: { org_id: "org-2" } }), [mockOrganizations[1]]); + + const { result } = renderHook(() => useOrganization("org-2"), { wrapper }); + + // initialData found org-2 in the filtered cache, so data is present on the first render. + expect(result.current.data).toEqual(mockOrganizations[1]); + expect(result.current.isLoading).toBe(false); + }); + + it("falls through to the detail API call when no cached list contains the organization", async () => { + (organizationInfoCall as any).mockResolvedValue(mockOrganizations[0]); + queryClient.setQueryData(organizationKeys.list({ filters: { org_id: "org-2" } }), [mockOrganizations[1]]); + + const { result } = renderHook(() => useOrganization("org-1"), { wrapper }); + + // org-1 is in no cached list, so there is no initialData and it loads via the detail API. + expect(result.current.data).toBeUndefined(); + + await waitFor(() => { + expect(result.current.isSuccess).toBe(true); + }); + + expect(organizationInfoCall).toHaveBeenCalledWith("test-access-token", "org-1"); + expect(result.current.data).toEqual(mockOrganizations[0]); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/organizations/useOrganizations.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/organizations/useOrganizations.ts index 0e7cd8342ec..734c1986f8f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/organizations/useOrganizations.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/organizations/useOrganizations.ts @@ -4,11 +4,23 @@ import { useQuery, useQueryClient, UseQueryResult } from "@tanstack/react-query" import { createQueryKeys } from "../common/queryKeysFactory"; export const organizationKeys = createQueryKeys("organizations"); -export const useOrganizations = (): UseQueryResult => { + +export interface OrganizationListFilters { + org_id?: string | null; + org_alias?: string | null; +} + +export const useOrganizations = (filters?: OrganizationListFilters): UseQueryResult => { const { accessToken, userId, userRole } = useAuthorized(); + const orgId = filters?.org_id || null; + const orgAlias = filters?.org_alias || null; return useQuery({ - queryKey: organizationKeys.list({}), - queryFn: async () => await organizationListCall(accessToken!), + queryKey: organizationKeys.list( + orgId || orgAlias + ? { filters: { ...(orgId && { org_id: orgId }), ...(orgAlias && { org_alias: orgAlias }) } } + : {}, + ), + queryFn: async () => await organizationListCall(accessToken!, orgId, orgAlias), enabled: Boolean(accessToken && userId && userRole), }); }; @@ -31,9 +43,10 @@ export const useOrganization = (organizationID?: string) => { initialData: () => { if (!organizationID) return undefined; - const organizations = queryClient.getQueryData(organizationKeys.list({})); - - return organizations?.find((organization: Organization) => organization.organization_id === organizationID); + return queryClient + .getQueriesData({ queryKey: organizationKeys.lists() }) + .flatMap(([, organizations]) => organizations ?? []) + .find((organization) => organization.organization_id === organizationID); }, }); }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/organizations/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/organizations/page.tsx new file mode 100644 index 00000000000..87e0faf9cce --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/organizations/page.tsx @@ -0,0 +1,9 @@ +"use client"; + +import OrganizationsTable from "@/components/organizations"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; + +export default function OrganizationsPage() { + const { accessToken, userRole, premiumUser } = useAuthorized(); + return ; +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/page.tsx index 2ef7839dadc..2369e130eee 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/page.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/page.tsx @@ -7,8 +7,8 @@ import LoadingScreen from "@/components/common_components/LoadingScreen"; import { Team } from "@/components/key_team_helpers/key_list"; import { Organization, proxyBaseUrl, getInProductNudgesCall } from "@/components/networking"; import OldTeams from "@/components/OldTeams"; -import { fetchUserModels, CreateKeyPrefillData } from "@/components/organisms/create_key_button"; -import Organizations, { fetchOrganizations } from "@/components/organizations"; +import { CreateKeyPrefillData } from "@/components/organisms/create_key_button"; +import { fetchOrganizations } from "@/components/organizations"; import PassThroughSettings from "@/components/pass_through_settings"; import { SurveyPrompt, SurveyModal, ClaudeCodePrompt, ClaudeCodeModal } from "@/components/survey"; import Usage from "@/components/usage"; @@ -32,7 +32,6 @@ function CreateKeyPageContent() { const [teams, setTeams] = useState(null); const [keys, setKeys] = useState([]); const [organizations, setOrganizations] = useState([]); - const [userModels, setUserModels] = useState([]); const router = useRouter(); const searchParams = useSearchParams()!; @@ -172,9 +171,6 @@ function CreateKeyPageContent() { }, [token]); useEffect(() => { - if (accessToken && userID && userRole) { - fetchUserModels(userID, userRole, accessToken, setUserModels); - } if (accessToken && userID && userRole) { v2TeamListCall(accessToken, 1, 100, { userID: userRole !== "Admin" && userRole !== "Admin Viewer" ? userID : null, @@ -333,15 +329,6 @@ function CreateKeyPageContent() { premiumUser={premiumUser} searchParams={searchParams} /> - ) : page == "organizations" ? ( - ) : page == "pass-through-settings" ? ( ({ __esModule: true, @@ -10,22 +11,27 @@ vi.mock("./mcp_server_management/MCPServerSelector", () => ({ __esModule: true, default: () => null, })); +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: () => ({ + accessToken: null, + userId: null, + userRole: null, + }), +})); import OrganizationsTable from "./organizations"; +const renderWithQueryClient = (ui: React.ReactElement) => { + const queryClient = new QueryClient({ + defaultOptions: { queries: { retry: false } }, + }); + return render({ui}); +}; + describe("OrganizationsTable", () => { it("should render the OrganizationsTable component", () => { - const setOrganizations = vi.fn(); - - const { getByText } = render( - , + const { getByText } = renderWithQueryClient( + , ); expect(getByText("+ Create New Organization")).toBeInTheDocument(); diff --git a/ui/litellm-dashboard/src/components/organizations.tsx b/ui/litellm-dashboard/src/components/organizations.tsx index fc825bdf924..d0629bb9cfd 100644 --- a/ui/litellm-dashboard/src/components/organizations.tsx +++ b/ui/litellm-dashboard/src/components/organizations.tsx @@ -1,3 +1,5 @@ +import { organizationKeys, useOrganizations } from "@/app/(dashboard)/hooks/organizations/useOrganizations"; +import { useUserModels } from "@/app/(dashboard)/hooks/models/useModels"; import OrganizationFilters, { FilterState } from "@/app/(dashboard)/organizations/OrganizationFilters"; import { InfoCircleOutlined } from "@ant-design/icons"; import { ChevronDownIcon, ChevronRightIcon, RefreshIcon } from "@heroicons/react/outline"; @@ -23,6 +25,7 @@ import { TextInput, } from "@tremor/react"; import { Form, Input, Modal, Select as Select2, Tooltip } from "antd"; +import { useQueryClient } from "@tanstack/react-query"; import React, { useState } from "react"; import { formatNumberWithCommas } from "../utils/dataUtils"; import DeleteResourceModal from "./common_components/DeleteResourceModal"; @@ -37,15 +40,10 @@ import NumericalInput from "./shared/numerical_input"; import VectorStoreSelector from "./vector_store_management/VectorStoreSelector"; interface OrganizationsTableProps { - organizations: Organization[]; userRole: string; - userModels: string[]; accessToken: string | null; lastRefreshed?: string; handleRefreshClick?: () => void; - currentOrg?: any; - guardrailsList?: string[]; - setOrganizations: (organizations: Organization[]) => void; premiumUser: boolean; } @@ -60,15 +58,10 @@ export const fetchOrganizations = async ( }; const OrganizationsTable: React.FC = ({ - organizations, userRole, - userModels, accessToken, lastRefreshed, handleRefreshClick, - currentOrg, - guardrailsList = [], - setOrganizations, premiumUser, }) => { const [selectedOrgId, setSelectedOrgId] = useState(null); @@ -87,21 +80,14 @@ const OrganizationsTable: React.FC = ({ sort_order: "desc", }); + const queryClient = useQueryClient(); + const { data: organizations = [] } = useOrganizations({ org_id: filters.org_id, org_alias: filters.org_alias }); + const { data: userModels = [] } = useUserModels(); + + const refetchOrganizations = () => queryClient.invalidateQueries({ queryKey: organizationKeys.lists() }); + const handleFilterChange = (key: keyof FilterState, value: string) => { - const newFilters = { ...filters, [key]: value }; - setFilters(newFilters); - // Call organizationListCall with the new filters - if (accessToken) { - organizationListCall(accessToken, newFilters.org_id || null, newFilters.org_alias || null) - .then((response) => { - if (response) { - setOrganizations(response); - } - }) - .catch((error) => { - console.error("Error fetching organizations:", error); - }); - } + setFilters((previousFilters) => ({ ...previousFilters, [key]: value })); }; const handleFilterReset = () => { @@ -111,18 +97,6 @@ const OrganizationsTable: React.FC = ({ sort_by: "created_at", sort_order: "desc", }); - // Reset organizations list - if (accessToken) { - organizationListCall(accessToken, null, null) - .then((response) => { - if (response) { - setOrganizations(response); - } - }) - .catch((error) => { - console.error("Error fetching organizations:", error); - }); - } }; const handleDelete = (orgId: string | null) => { @@ -142,8 +116,7 @@ const OrganizationsTable: React.FC = ({ setIsDeleteModalOpen(false); setOrgToDelete(null); - // Refresh organizations list - await fetchOrganizations(accessToken, setOrganizations, filters.org_id || null, filters.org_alias || null); + await refetchOrganizations(); } catch (error) { console.error("Error deleting organization:", error); } finally { @@ -189,8 +162,7 @@ const OrganizationsTable: React.FC = ({ NotificationsManager.success("Organization created successfully"); setIsOrgModalVisible(false); form.resetFields(); - // Refresh organizations list - fetchOrganizations(accessToken, setOrganizations, filters.org_id || null, filters.org_alias || null); + await refetchOrganizations(); } catch (error) { console.error("Error creating organization:", error); } diff --git a/ui/litellm-dashboard/src/utils/migratedPages.test.ts b/ui/litellm-dashboard/src/utils/migratedPages.test.ts index f2e0f4eb151..3afd5a81a9a 100644 --- a/ui/litellm-dashboard/src/utils/migratedPages.test.ts +++ b/ui/litellm-dashboard/src/utils/migratedPages.test.ts @@ -126,6 +126,13 @@ describe("migratedHref / legacyPageHref", () => { expect(MIGRATED_PAGES.users).toBe("users"); }); + + it("maps the organizations id to its route", async () => { + vi.doMock("@/components/networking", () => ({ serverRootPath: "/" })); + const { MIGRATED_PAGES } = await import("./migratedPages"); + + expect(MIGRATED_PAGES.organizations).toBe("organizations"); + }); }); describe("dev server (NODE_ENV=development)", () => { diff --git a/ui/litellm-dashboard/src/utils/migratedPages.ts b/ui/litellm-dashboard/src/utils/migratedPages.ts index 337018f6ec5..e9e6c527513 100644 --- a/ui/litellm-dashboard/src/utils/migratedPages.ts +++ b/ui/litellm-dashboard/src/utils/migratedPages.ts @@ -43,6 +43,7 @@ export const MIGRATED_PAGES: Record = { agents: "agents", "router-settings": "router-settings", users: "users", + organizations: "organizations", }; function uiBase(): string {