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:
ryan-crabbe-berri 2026-06-29 13:27:31 -07:00 • committed by GitHub
parent 0e5aee1838
commit 5e5b09709c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 143 additions and 18 deletions

View file

@ -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));
});
});

View file

@ -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();

View file

@ -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 () => {

View file

@ -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