diff --git a/web/packages/studio/src/components/dataViews/ExperimentGroupDataView/index.tsx b/web/packages/studio/src/components/dataViews/ExperimentGroupDataView/index.tsx index 5e6cc41e17..7dc7fddf7b 100644 --- a/web/packages/studio/src/components/dataViews/ExperimentGroupDataView/index.tsx +++ b/web/packages/studio/src/components/dataViews/ExperimentGroupDataView/index.tsx @@ -12,23 +12,22 @@ import { RelativeTime } from '@nemo/common/src/components/RelativeTime'; import { useStudioDataViewState } from '@nemo/common/src/hooks/useStudioDataViewState'; import { snakeCaseToTitleCase } from '@nemo/common/src/utils/formatters'; import { getSortParamWithWhitelist } from '@nemo/common/src/utils/query'; -import { useGetExperimentGroup, useListExperiments } from '@nemo/sdk/generated/platform/api'; -import type { - ExperimentFilter, - ExperimentResponse, - ListExperimentsSort, -} from '@nemo/sdk/generated/platform/schema'; -import { Text, Tooltip } from '@nvidia/foundations-react-core'; +import { useGetExperimentGroup } from '@nemo/sdk/generated/platform/api'; +import type { ExperimentFilter } from '@nemo/sdk/generated/platform/schema'; +import { Button, Text, Tooltip } from '@nvidia/foundations-react-core'; import { Empty } from '@studio/components/dataViews/ExperimentGroupDataView/Empty'; +import { + type ExperimentRow, + useExperimentGroupExperiments, +} from '@studio/components/dataViews/ExperimentGroupDataView/useExperimentGroupExperiments'; import { useWorkspaceFromPath } from '@studio/hooks/useWorkspaceFromPath'; import { getExperimentDetailRoute } from '@studio/routes/utils'; import { tooltipClassName } from '@studio/styles/common'; -import { keepPreviousData } from '@tanstack/react-query'; -import { Columns3 } from 'lucide-react'; +import { Columns3, Pin } from 'lucide-react'; import { type ComponentProps, type FC, useCallback, useMemo } from 'react'; import { useNavigate } from 'react-router-dom'; -export type ExperimentRow = ExperimentResponse & { id: string }; +export type { ExperimentRow }; const SORTABLE_FIELDS = ['name', 'created_at'] as const; const DEFAULT_SORT = '-created_at'; @@ -64,6 +63,8 @@ export const ExperimentGroupDataView: FC = ({ defaultSort: { id: 'created_at', desc: true }, // created_by isn't returned by the API and updated_at isn't shown; both are filter-only. columnVisibility: { created_by: false, updated_at: false }, + // Keep the pin toggle reachable while horizontally scrolling this wide table. + columnPinning: { left: ['pin'] }, }); const page = dataViewState.pagination.state.pageIndex + 1; @@ -75,42 +76,26 @@ export const ExperimentGroupDataView: FC = ({ ); const { - data: experimentsResponse, - isLoading, + rows: orderedData, + togglePin, + totalCount, error, - } = useListExperiments( + isLoading, + } = useExperimentGroupExperiments({ workspace, - { - page, - page_size: pageSize, - sort: sortParam as ListExperimentsSort, - // User filters merge under the group scope, which always wins so it can't be overridden. - filter: { - ...dataViewState.apiFilter.filter, - ...(dataViewState.searchBar.state && { name: { $like: dataViewState.searchBar.state } }), - experiment_group_id: experimentGroupId, - } as ExperimentFilter, - }, - { query: { placeholderData: keepPreviousData, enabled: !!experimentGroupId } } - ); - - const experimentsData = experimentsResponse?.data; - const totalCount = experimentsResponse?.pagination?.total_results ?? experimentsData?.length ?? 0; - - const tableData = useMemo( - () => - (experimentsData ?? []).map((experiment) => ({ - ...experiment, - id: experiment.id ?? experiment.name ?? '', - })), - [experimentsData] - ); + experimentGroupId, + filter: dataViewState.apiFilter.filter, + search: dataViewState.debouncedSearchBar, + page, + pageSize, + sort: sortParam, + }); // One score column per evaluator: the union of evaluator names across the loaded rows, // sorted for a deterministic column order across renders and page changes. const evaluatorNames = useMemo( - () => [...new Set(tableData.flatMap((e) => Object.keys(e.aggregate_scores ?? {})))].sort(), - [tableData] + () => [...new Set(orderedData.flatMap((e) => Object.keys(e.aggregate_scores ?? {})))].sort(), + [orderedData] ); // One column per metadata key: keys are lowercased so case variants (e.g. "status" @@ -119,16 +104,46 @@ export const ExperimentGroupDataView: FC = ({ () => [ ...new Set( - tableData.flatMap((e) => Object.keys(e.metadata ?? {}).map((k) => k.toLowerCase())) + orderedData.flatMap((e) => Object.keys(e.metadata ?? {}).map((k) => k.toLowerCase())) ), ].sort(), - [tableData] + [orderedData] ); const makeColumns = useCallback< ComponentProps>['makeColumns'] >( - ({ accessor }) => [ + ({ accessor, display }) => [ + display({ + id: 'pin', + header: () => Pinned, + enableSorting: false, + enableHiding: false, + enableResizing: false, + size: 48, + minSize: 48, + maxSize: 48, + meta: { alignment: 'center', _isPrebuiltColumn: true, _isSizeInitialized: true }, + cell: ({ row }) => { + const { pinned_at } = row.original; + const isPinned = pinned_at != null; + return ( + + ); + }, + }), accessor('name', { header: 'Name', enableSorting: true, @@ -282,7 +297,7 @@ export const ExperimentGroupDataView: FC = ({ }, }), ], - [evaluatorNames, metadataKeys] + [evaluatorNames, togglePin, metadataKeys] ); if (groupError) { @@ -317,9 +332,9 @@ export const ExperimentGroupDataView: FC = ({ } attributes={{ DataViewRoot: { - data: tableData, + data: orderedData, totalCount, - requestStatus: isGroupLoading || (isLoading && !experimentsData) ? 'loading' : undefined, + requestStatus: isGroupLoading || isLoading ? 'loading' : undefined, }, DataViewTableContent: { renderEmptyState: ({ hasFiltersApplied, hasSearchApplied }) => diff --git a/web/packages/studio/src/components/dataViews/ExperimentGroupDataView/useExperimentGroupExperiments.test.ts b/web/packages/studio/src/components/dataViews/ExperimentGroupDataView/useExperimentGroupExperiments.test.ts new file mode 100644 index 0000000000..917b59384c --- /dev/null +++ b/web/packages/studio/src/components/dataViews/ExperimentGroupDataView/useExperimentGroupExperiments.test.ts @@ -0,0 +1,180 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +import { useToast } from '@nemo/common/src/providers/toast/useToast'; +import { + getListExperimentsQueryKey, + useListExperiments, + usePinExperiment, + useUnpinExperiment, +} from '@nemo/sdk/generated/platform/api'; +import { + useExperimentGroupExperiments, + type UseExperimentGroupExperimentsParams, +} from '@studio/components/dataViews/ExperimentGroupDataView/useExperimentGroupExperiments'; +import { renderHook } from '@testing-library/react'; + +vi.mock('@nemo/common/src/providers/toast/useToast'); +vi.mock('@nemo/sdk/generated/platform/api'); +// useQueryClient needs a provider; the queries themselves are mocked, so a stub client is enough. +const { invalidateQueries } = vi.hoisted(() => ({ invalidateQueries: vi.fn() })); +vi.mock('@tanstack/react-query', async (importOriginal) => ({ + ...(await importOriginal()), + useQueryClient: () => ({ invalidateQueries }), +})); + +const mockUseToast = vi.mocked(useToast); +const mockUseListExperiments = vi.mocked(useListExperiments); +const mockUsePinExperiment = vi.mocked(usePinExperiment); +const mockUseUnpinExperiment = vi.mocked(useUnpinExperiment); +const mockGetListExperimentsQueryKey = vi.mocked(getListExperimentsQueryKey); + +interface Row { + id: string; + name: string; + pinned_at?: string | null; +} + +const pin = (name: string): Row => ({ id: name, name, pinned_at: '2026-01-01T00:00:00Z' }); +const unp = (name: string): Row => ({ id: name, name, pinned_at: null }); + +// Minimal stand-in for the SDK query result; the hook only reads `data`, `pagination.total_results`, +// and the loading/fetching/error flags. +const queryResult = (rows: Row[], total: number) => + ({ + data: { data: rows, pagination: { total_results: total } }, + isLoading: false, + isFetching: false, + error: null, + }) as unknown as ReturnType; + +const mockLists = ( + pinned: { rows: Row[]; total: number }, + unpinned: { rows: Row[]; total: number } +) => { + mockUseListExperiments.mockImplementation(((_workspace, params) => + (params?.filter as { is_pinned?: boolean } | undefined)?.is_pinned + ? queryResult(pinned.rows, pinned.total) + : queryResult(unpinned.rows, unpinned.total)) as typeof useListExperiments); +}; + +const baseParams: UseExperimentGroupExperimentsParams = { + workspace: 'ws', + experimentGroupId: 'grp', + filter: undefined, + search: '', + page: 1, + pageSize: 50, + sort: '-created_at', +}; + +describe('useExperimentGroupExperiments', () => { + beforeEach(() => { + vi.clearAllMocks(); + mockUseToast.mockReturnValue({ error: vi.fn() } as unknown as ReturnType); + mockUsePinExperiment.mockReturnValue({ + mutate: vi.fn(), + } as unknown as ReturnType); + mockUseUnpinExperiment.mockReturnValue({ + mutate: vi.fn(), + } as unknown as ReturnType); + }); + + it('paginates over the unpinned set only, so pinned rows do not inflate the page count', () => { + // 3 pinned + exactly one page (50) of unpinned. Summing the two (the old behavior) gives 53, + // which would make the table render a phantom 2nd page containing only the pinned rows again. + mockLists( + { rows: [pin('p1'), pin('p2'), pin('p3')], total: 3 }, + { rows: Array.from({ length: 50 }, (_unused, i) => unp(`u${i}`)), total: 50 } + ); + + const { result } = renderHook(() => useExperimentGroupExperiments(baseParams)); + + expect(result.current.totalCount).toBe(50); + }); + + it('falls back to the pinned count when nothing is unpinned so a fully-pinned group is not empty', () => { + mockLists({ rows: [pin('p1'), pin('p2')], total: 2 }, { rows: [], total: 0 }); + + const { result } = renderHook(() => useExperimentGroupExperiments(baseParams)); + + expect(result.current.totalCount).toBe(2); + }); + + it('reports a zero count only when both lists are empty', () => { + mockLists({ rows: [], total: 0 }, { rows: [], total: 0 }); + + const { result } = renderHook(() => useExperimentGroupExperiments(baseParams)); + + expect(result.current.totalCount).toBe(0); + }); + + it('lists pinned rows first, then unpinned, dropping an unpinned row already shown as pinned', () => { + // 'b' appears in both lists (the brief window where the two queries refetch out of step). + mockLists({ rows: [pin('a'), pin('b')], total: 2 }, { rows: [unp('b'), unp('c')], total: 2 }); + + const { result } = renderHook(() => useExperimentGroupExperiments(baseParams)); + + expect(result.current.rows.map((row) => row.name)).toEqual(['a', 'b', 'c']); + }); + + it('fetches the full pinned set in one page and paginates the unpinned set by the caller page size', () => { + mockLists({ rows: [], total: 0 }, { rows: [], total: 0 }); + + renderHook(() => useExperimentGroupExperiments({ ...baseParams, page: 2, pageSize: 25 })); + + const calls = mockUseListExperiments.mock.calls; + const pinnedParams = calls.find( + (call) => (call[1]?.filter as { is_pinned?: boolean } | undefined)?.is_pinned === true + )?.[1]; + const unpinnedParams = calls.find( + (call) => (call[1]?.filter as { is_pinned?: boolean } | undefined)?.is_pinned === false + )?.[1]; + + // Pinned: a single large page (MAX_PINNED_ROWS), pinned-recency order, independent of caller page. + expect(pinnedParams).toMatchObject({ page: 1, page_size: 100, sort: '-pinned_at' }); + // Unpinned: the caller's page/page_size and sort. + expect(unpinnedParams).toMatchObject({ page: 2, page_size: 25, sort: '-created_at' }); + }); + + it('stays loading until both queries have loaded, not just the faster one', () => { + // Pinned has returned; unpinned is still on its initial load (no data yet). + mockUseListExperiments.mockImplementation(((_workspace, params) => + (params?.filter as { is_pinned?: boolean } | undefined)?.is_pinned + ? queryResult([pin('p')], 1) + : ({ + data: undefined, + isLoading: true, + isFetching: true, + error: null, + } as unknown as ReturnType)) as typeof useListExperiments); + + const { result } = renderHook(() => useExperimentGroupExperiments(baseParams)); + + expect(result.current.isLoading).toBe(true); + }); + + it('clears loading once both queries have responded', () => { + mockLists({ rows: [pin('p')], total: 1 }, { rows: [unp('u')], total: 1 }); + + const { result } = renderHook(() => useExperimentGroupExperiments(baseParams)); + + expect(result.current.isLoading).toBe(false); + }); + + it('scopes pin/unpin invalidation to this group, not the whole workspace', () => { + mockLists({ rows: [], total: 0 }, { rows: [], total: 0 }); + + renderHook(() => useExperimentGroupExperiments(baseParams)); + // onSuccess is the group-scoped invalidate the hook wires into both mutations. + const onSuccess = mockUsePinExperiment.mock.calls[0]?.[0]?.mutation?.onSuccess as + | (() => void) + | undefined; + onSuccess?.(); + + expect(mockGetListExperimentsQueryKey).toHaveBeenCalledWith('ws', { + filter: { experiment_group_id: 'grp' }, + }); + expect(invalidateQueries).toHaveBeenCalledTimes(1); + }); +}); diff --git a/web/packages/studio/src/components/dataViews/ExperimentGroupDataView/useExperimentGroupExperiments.ts b/web/packages/studio/src/components/dataViews/ExperimentGroupDataView/useExperimentGroupExperiments.ts new file mode 100644 index 0000000000..a3e680a2cc --- /dev/null +++ b/web/packages/studio/src/components/dataViews/ExperimentGroupDataView/useExperimentGroupExperiments.ts @@ -0,0 +1,200 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +import { useToast } from '@nemo/common/src/providers/toast/useToast'; +import { + getListExperimentsQueryKey, + type ListExperimentsQueryError, + useListExperiments, + usePinExperiment, + useUnpinExperiment, +} from '@nemo/sdk/generated/platform/api'; +import type { + ExperimentFilter, + ExperimentResponse, + ListExperimentsSort, +} from '@nemo/sdk/generated/platform/schema'; +import { keepPreviousData, useQueryClient } from '@tanstack/react-query'; +import { useCallback, useMemo, useRef } from 'react'; + +/** An API response plus the stable `id` the data view needs. */ +export type ExperimentRow = ExperimentResponse & { id: string }; + +const toRows = (experiments: ExperimentResponse[] | undefined): ExperimentRow[] => + (experiments ?? []).map((experiment) => ({ + ...experiment, + id: experiment.id ?? experiment.name ?? '', + })); + +/** + * Pinned rows are shown in full atop every page rather than paginated, so the whole pin set is + * fetched in one request. A group's pin set is a small curated list; this bound keeps that request + * well under the API's max page size (1000) while covering any realistic number of pins. + */ +const MAX_PINNED_ROWS = 100; + +export interface UseExperimentGroupExperimentsParams { + workspace: string; + experimentGroupId: string; + filter: Partial | undefined; + search: string; + page: number; + pageSize: number; + sort: string; +} + +export interface ExperimentGroupExperiments { + /** Pinned rows first (newest-pinned first), then the current page of unpinned rows. */ + rows: ExperimentRow[]; + /** Pins the row if unpinned, unpins it otherwise, then refetches both lists. */ + togglePin: (row: ExperimentRow) => void; + /** + * Row count that drives pagination: the unpinned total. Pinned rows ride atop every page and are + * not paginated, so they're excluded; falls back to the pinned count when nothing is unpinned so + * a fully-pinned group still renders one page instead of reading as empty. + */ + totalCount: number; + error: ListExperimentsQueryError | null; + /** True until both queries have loaded once; ignores background refetches (page changes, pins). */ + isLoading: boolean; + /** True when either query is fetching */ + isFetching: boolean; +} + +/** + * Loads one experiment group's experiments as two queries — the full pinned set (`is_pinned=true`, + * sorted `-pinned_at`) and the current page of unpinned (`is_pinned=false`) — and concatenates them + * so pinned rows top every page. Pagination covers only the unpinned set; the pinned set is fetched + * once and repeated atop every page rather than paginated. A pin/unpin persists through the API, + * then invalidates both lists so the new state is refetched (no optimistic update). + */ +export function useExperimentGroupExperiments({ + workspace, + experimentGroupId, + filter, + search, + page, + pageSize, + sort, +}: UseExperimentGroupExperimentsParams): ExperimentGroupExperiments { + const queryClient = useQueryClient(); + const toast = useToast(); + + const baseFilter = { + ...filter, + ...(search && { name: { $like: search } }), + // Spread last so the group scope can't be overridden by a user filter. + experiment_group_id: experimentGroupId, + }; + const listQueryOptions = { + query: { + placeholderData: keepPreviousData, + enabled: !!experimentGroupId, + }, + }; + + const { + data: pinnedResponse, + isLoading: isPinnedLoading, + isFetching: isPinnedFetching, + error: pinnedError, + } = useListExperiments( + workspace, + { + page: 1, + page_size: MAX_PINNED_ROWS, + sort: '-pinned_at', + filter: { ...baseFilter, is_pinned: true } as ExperimentFilter, + }, + listQueryOptions + ); + + const { + data: unpinnedResponse, + isLoading: isUnpinnedLoading, + isFetching: isUnpinnedFetching, + error: unpinnedError, + } = useListExperiments( + workspace, + { + page, + page_size: pageSize, + sort: sort as ListExperimentsSort, + filter: { ...baseFilter, is_pinned: false } as ExperimentFilter, + }, + listQueryOptions + ); + + const pinned = useMemo(() => toRows(pinnedResponse?.data), [pinnedResponse]); + const unpinned = useMemo(() => toRows(unpinnedResponse?.data), [unpinnedResponse]); + + // Synchronous guard so a row can't fire a second pin/unpin while one is already in flight; the + // name is cleared once the mutation settles. + const pendingRef = useRef>(new Set()); + + // Scope invalidation to this group's experiment lists (any page/sort/filter, pinned and unpinned) + // via a partial key match on experiment_group_id, so a pin/unpin doesn't refetch other groups' + // lists. Returned (not voided) so the mutation's onSuccess awaits the refetch before onSettled + // re-enables the row. + const invalidateList = useCallback( + () => + queryClient.invalidateQueries({ + queryKey: getListExperimentsQueryKey(workspace, { + filter: { experiment_group_id: experimentGroupId }, + }), + }), + [queryClient, workspace, experimentGroupId] + ); + + const { mutate: pinExperiment } = usePinExperiment({ + mutation: { + onSuccess: invalidateList, + onError: () => toast.error('Failed to pin experiment.'), + onSettled: (_data, _error, { name }) => { + pendingRef.current.delete(name); + }, + }, + }); + const { mutate: unpinExperiment } = useUnpinExperiment({ + mutation: { + onSuccess: invalidateList, + onError: () => toast.error('Failed to unpin experiment.'), + onSettled: (_data, _error, { name }) => { + pendingRef.current.delete(name); + }, + }, + }); + + const togglePin = useCallback( + (row: ExperimentRow) => { + const { name } = row; + if (pendingRef.current.has(name)) return; + pendingRef.current.add(name); + if (row.pinned_at != null) unpinExperiment({ workspace, name }); + else pinExperiment({ workspace, name }); + }, + [workspace, pinExperiment, unpinExperiment] + ); + + // Pinned first, then the current page of unpinned. Drop any unpinned row already shown as pinned — + // it can appear in both server lists during the brief window where the two queries refetch out of step. + const rows = useMemo(() => { + const pinnedNames = new Set(pinned.map((row) => row.name)); + return [...pinned, ...unpinned.filter((row) => !pinnedNames.has(row.name))]; + }, [pinned, unpinned]); + + // Pagination covers the unpinned set only — pinned rows are a fixed header repeated atop every + // page, so counting them would inflate the page count and add trailing pages that show nothing + // but the (already-visible) pinned rows. Fall back to the pinned count when nothing is unpinned so + // a fully-pinned group still renders a single page instead of reading as empty. + const unpinnedTotal = unpinnedResponse?.pagination?.total_results ?? unpinned.length; + const totalCount = unpinnedTotal > 0 ? unpinnedTotal : pinned.length; + const error = pinnedError ?? unpinnedError; + // react-query's per-query isLoading is true only on first load (keepPreviousData keeps data across + // refetches, so isPending stays false). OR-ing the two keeps the table in its loading state until + // both lists have loaded once, rather than clearing as soon as the faster query returns. + const isLoading = isPinnedLoading || isUnpinnedLoading; + const isFetching = isPinnedFetching || isUnpinnedFetching; + + return { rows, togglePin, totalCount, error, isLoading, isFetching }; +}