mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-06 02:48:13 +00:00
perf(ui): load virtual-keys team filter from the fast v2 endpoint (#31638)
* perf(ui): load virtual-keys team filter from the fast v2 endpoint The virtual-keys table sourced all teams through fetchAllTeams, which hits the unpaginated /team/list. On a proxy with 125 teams that call takes ~9.5s, so the Team ID filter and the team-alias/budget columns sat empty for that whole window. The key list itself does not carry team_alias or team_max_budget, so the table genuinely needs a team lookup and cannot just drop the fetch. Add useAllTeams, which pages the fast /v2/team/list to completion (~0.6s per 100-team page, so ~1.2s for 125 vs ~9.5s), and point VirtualKeysTable at it instead of fetchAllTeams. The allTeams shape, the filter searchFn, the column lookups, and the loading indicator are all unchanged; only the source endpoint changes. fetchAllTeams stays for its other callers. * test(ui): tighten team-filter test readability and robustness Address adversarial review of the added tests. Rename the single-page useAllTeams test to match what it asserts (one request for a one-page result) rather than implying it guards the early-return, and drop the unread, misleading total: 125 from the mock page response. Scope the created_by alias-over-email assertion to the key's table row so it checks the visible cell value; the hover popover that also holds the email is portaled out of the row, so the previous document-wide negative assertion was relying on antd's lazy popover mounting. * fix(ui): scope useAllTeams cache by access token The previous /team/list query keyed on accessToken, so a user switch in the same SPA session produced a distinct cache entry. useAllTeams dropped that, so team IDs and aliases could be briefly reused across users until the staleTime expired. Put accessToken back in the query key to restore per-identity isolation, and add a regression test that a token switch triggers a refetch rather than serving the cached list.
This commit is contained in:
parent
0e5aee1838
commit
5e5b09709c
4 changed files with 143 additions and 18 deletions
|
|
@ -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<typeof vi.fn>;
|
||||
|
||||
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));
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -95,6 +95,31 @@ export const useTeams = (): UseQueryResult<Team[]> => {
|
|||
});
|
||||
};
|
||||
|
||||
const ALL_TEAMS_PAGE_SIZE = 100;
|
||||
|
||||
const fetchAllTeamsPaged = async (accessToken: string): Promise<Team[]> => {
|
||||
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<Team[]> => {
|
||||
const { accessToken } = useAuthorized();
|
||||
return useQuery<Team[]>({
|
||||
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();
|
||||
|
|
|
|||
|
|
@ -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(<VirtualKeysTable />);
|
||||
|
||||
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 () => {
|
||||
|
|
|
|||
|
|
@ -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<KeyResponse | null>(null);
|
||||
|
|
@ -98,12 +95,7 @@ export function VirtualKeysTable() {
|
|||
|
||||
const keyList = useMemo(() => keys?.keys ?? [], [keys]);
|
||||
|
||||
const { data: fetchedTeams, isLoading: isTeamsLoading } = useQuery<Team[]>({
|
||||
queryKey: ["allTeamsForKeyFilters", accessToken],
|
||||
queryFn: async () => (accessToken ? await fetchAllTeams(accessToken) : []),
|
||||
enabled: !!accessToken,
|
||||
staleTime: 30000,
|
||||
});
|
||||
const { data: fetchedTeams, isLoading: isTeamsLoading } = useAllTeams();
|
||||
const allTeams = useMemo<Team[]>(() => fetchedTeams ?? [], [fetchedTeams]);
|
||||
|
||||
// Defer the transition so the button stays in loading state until the table
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue