diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.test.ts index 20f034ada36..28320b76597 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.test.ts @@ -2,7 +2,7 @@ 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 { useTeams, useTeam, useDeletedTeams, DeletedTeam, teamListCall } from "./useTeams"; +import { useTeams, useTeam, useAllTeams, useDeletedTeams, DeletedTeam, teamListCall } from "./useTeams"; import { fetchTeams } from "@/app/(dashboard)/networking"; import { teamInfoCall } from "@/components/networking"; import type { Team } from "@/components/key_team_helpers/key_list"; @@ -792,3 +792,107 @@ describe("useDeletedTeams", () => { expect(result.current.error).toBeNull(); }); }); + +describe("useAllTeams", () => { + let queryClient: QueryClient; + let fetchMock: ReturnType; + + 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, + }); + fetchMock = vi.fn(); + global.fetch = fetchMock as unknown as typeof fetch; + }); + + const wrapper = ({ children }: { children: ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); + + const pageResponse = (teams: Team[], page: number, totalPages: number) => ({ + ok: true, + json: async () => ({ teams, page, page_size: 100, total_pages: totalPages }), + }); + + const requestedPage = (url: string) => new URLSearchParams(url.split("?")[1]).get("page"); + + it("paginates /v2/team/list to completion and concatenates every page", async () => { + fetchMock.mockImplementation((url: string) => + Promise.resolve( + requestedPage(url) === "1" ? pageResponse([mockTeams[0]], 1, 2) : pageResponse([mockTeams[1]], 2, 2), + ), + ); + + const { result } = renderHook(() => useAllTeams(), { wrapper }); + + await waitFor(() => expect(result.current.isSuccess).toBe(true)); + + expect(result.current.data).toEqual(mockTeams); + expect(fetchMock).toHaveBeenCalledTimes(2); + const requestedPages = fetchMock.mock.calls.map((call) => requestedPage(call[0] as string)).sort(); + expect(requestedPages).toEqual(["1", "2"]); + const firstUrl = fetchMock.mock.calls[0][0] as string; + expect(firstUrl).toContain("/v2/team/list"); + expect(firstUrl).toContain("page_size=100"); + }); + + it("issues exactly one request for a single-page result", async () => { + fetchMock.mockResolvedValue(pageResponse(mockTeams, 1, 1)); + + const { result } = renderHook(() => useAllTeams(), { wrapper }); + + await waitFor(() => expect(result.current.isSuccess).toBe(true)); + + expect(result.current.data).toEqual(mockTeams); + expect(fetchMock).toHaveBeenCalledTimes(1); + }); + + it("does not execute 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(() => useAllTeams(), { wrapper }); + + expect(result.current.isLoading).toBe(false); + expect(result.current.isFetched).toBe(false); + expect(fetchMock).not.toHaveBeenCalled(); + }); + + it("scopes the cache per access token so a switch of identity refetches", async () => { + fetchMock.mockResolvedValue(pageResponse(mockTeams, 1, 1)); + + const { result, rerender } = renderHook(() => useAllTeams(), { wrapper }); + await waitFor(() => expect(result.current.isSuccess).toBe(true)); + expect(fetchMock).toHaveBeenCalledTimes(1); + + mockUseAuthorized.mockReturnValue({ + accessToken: "a-different-users-token", + userId: "other-user-id", + userRole: "Admin", + token: "a-different-users-token", + userEmail: "other@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + rerender(); + + await waitFor(() => expect(fetchMock).toHaveBeenCalledTimes(2)); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.ts index c356434ba04..81d2b84e9d9 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.ts @@ -95,6 +95,31 @@ export const useTeams = (): UseQueryResult => { }); }; +const ALL_TEAMS_PAGE_SIZE = 100; + +const fetchAllTeamsPaged = async (accessToken: string): Promise => { + const firstPage: TeamsResponse = await teamListCall(accessToken, 1, ALL_TEAMS_PAGE_SIZE); + const totalPages = firstPage.total_pages ?? 1; + if (totalPages <= 1) return firstPage.teams; + + const remainingPages: TeamsResponse[] = await Promise.all( + Array.from({ length: totalPages - 1 }, (_, i) => teamListCall(accessToken, i + 2, ALL_TEAMS_PAGE_SIZE)), + ); + return [firstPage, ...remainingPages].flatMap((page) => page.teams); +}; + +export const useAllTeams = (): UseQueryResult => { + const { accessToken } = useAuthorized(); + return useQuery({ + queryKey: teamKeys.list({ + filters: { scope: "all", pageSize: ALL_TEAMS_PAGE_SIZE, accessToken: accessToken ?? "" }, + }), + queryFn: async () => await fetchAllTeamsPaged(accessToken!), + enabled: Boolean(accessToken), + staleTime: 30000, + }); +}; + export const useTeam = (teamId?: string) => { const { accessToken } = useAuthorized(); const queryClient = useQueryClient(); diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx index 6389b82c9fd..ba616a03fd9 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx @@ -1,4 +1,4 @@ -import { act, screen, waitFor, fireEvent } from "@testing-library/react"; +import { act, screen, waitFor, within, fireEvent } from "@testing-library/react"; import { vi, it, expect, beforeEach, describe, MockedFunction } from "vitest"; import { renderWithProviders } from "../../../tests/test-utils"; import { VirtualKeysTable } from "./VirtualKeysTable"; @@ -28,8 +28,11 @@ vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ })), })); -vi.mock("../key_team_helpers/filter_helpers", () => ({ - fetchAllTeams: vi.fn().mockResolvedValue([{ team_id: "team-1", team_alias: "Test Team" }]), +vi.mock("@/app/(dashboard)/hooks/teams/useTeams", () => ({ + useAllTeams: vi.fn(() => ({ + data: [{ team_id: "team-1", team_alias: "Test Team" }], + isLoading: false, + })), })); vi.mock("@/app/(dashboard)/hooks/keys/useKeys", () => ({ @@ -309,10 +312,11 @@ it("should display created_by_user alias over email when both are available", as renderWithProviders(); - await waitFor(() => { - expect(screen.getByText("The Creator")).toBeInTheDocument(); - }); - expect(screen.queryByText("creator@example.com")).not.toBeInTheDocument(); + // Scope to the key's row so we assert the visible cell value: the hover popover that + // also holds the email is portaled out of the row, not the displayed "Created By" text. + const row = (await screen.findByText("Test Key Alias")).closest("tr") as HTMLElement; + expect(within(row).getByText("The Creator")).toBeInTheDocument(); + expect(within(row).queryByText("creator@example.com")).not.toBeInTheDocument(); }); it("should render table without crashing when models is null", async () => { diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx index 0c71c4526ad..a803d78ed57 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx @@ -1,8 +1,7 @@ "use client"; import { useKeys, KeyListCallOptions } from "@/app/(dashboard)/hooks/keys/useKeys"; import { useOrganizations } from "@/app/(dashboard)/hooks/organizations/useOrganizations"; -import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; -import { useQuery } from "@tanstack/react-query"; +import { useAllTeams } from "@/app/(dashboard)/hooks/teams/useTeams"; import { useDebouncedValue } from "@tanstack/react-pacer/debouncer"; import { formatNumberWithCommas } from "@/utils/dataUtils"; import { ChevronDownIcon, ChevronRightIcon, ChevronUpIcon, SwitchVerticalIcon } from "@heroicons/react/outline"; @@ -30,7 +29,6 @@ import { InfoCircleOutlined, SyncOutlined } from "@ant-design/icons"; import { Button as AntButton, Popover, Skeleton, Tag, Tooltip, Typography } from "antd"; import React, { useDeferredValue, useMemo, useState } from "react"; import { getModelDisplayName } from "../key_team_helpers/fetch_available_models_team_key"; -import { fetchAllTeams } from "../key_team_helpers/filter_helpers"; import { PaginatedKeyAliasSelect } from "../KeyAliasSelect/PaginatedKeyAliasSelect/PaginatedKeyAliasSelect"; import { KeyResponse, Team } from "../key_team_helpers/key_list"; import FilterComponent, { FilterOption } from "../molecules/filter"; @@ -67,7 +65,6 @@ const toKeyListFilters = (filters: KeyFilterState): KeyListFilterOptions => ({ }); export function VirtualKeysTable() { - const { accessToken } = useAuthorized(); const { data: fetchedOrganizations, isLoading: isOrgsLoading } = useOrganizations(); const resolvedOrganizations = useMemo(() => fetchedOrganizations ?? [], [fetchedOrganizations]); const [selectedKey, setSelectedKey] = useState(null); @@ -98,12 +95,7 @@ export function VirtualKeysTable() { const keyList = useMemo(() => keys?.keys ?? [], [keys]); - const { data: fetchedTeams, isLoading: isTeamsLoading } = useQuery({ - queryKey: ["allTeamsForKeyFilters", accessToken], - queryFn: async () => (accessToken ? await fetchAllTeams(accessToken) : []), - enabled: !!accessToken, - staleTime: 30000, - }); + const { data: fetchedTeams, isLoading: isTeamsLoading } = useAllTeams(); const allTeams = useMemo(() => fetchedTeams ?? [], [fetchedTeams]); // Defer the transition so the button stays in loading state until the table