Fix row selection (duplicate experiments) (#3921)

* Fix row selection (duplicate experiments)

* Return unique list of ids

Author: sroy3

webview/src/experiments/components/table/RowSelectionContext.tsx

@@ -1,5 +1,6 @@
 import React, { createContext, useState } from 'react'
 import { RowProp } from '../../util/interfaces'
+import { getCompositeId } from '../../util/rows'
 
 interface RowSelectionContextValue {
   selectedRows: Record<string, RowProp | undefined>
@@ -28,14 +29,16 @@ export const RowSelectionProvider: React.FC<{ children: React.ReactNode }> = ({
   const toggleRowSelected = (rowProp: RowProp) => {
     const {
       row: {
-        original: { id }
+        original: { id, branch }
       }
     } = rowProp
+    const compositeId = getCompositeId(id, branch)
+
     setSelectedRows({
       ...selectedRows,
-      [id]: selectedRows[id] ? undefined : rowProp
+      [compositeId]: selectedRows[compositeId] ? undefined : rowProp
     })
-    setSelectionHistory([id, ...selectionHistory])
+    setSelectionHistory([compositeId, ...selectionHistory])
   }
 
   const batchSelection = (batch: RowProp[]) => {
@@ -44,21 +47,22 @@ export const RowSelectionProvider: React.FC<{ children: React.ReactNode }> = ({
     for (const rowProp of batch) {
       const {
         row: {
-          original: { id }
+          original: { id, branch }
         }
       } = rowProp
-      selectedRowsCopy[id] = rowProp
+      const compositeId = getCompositeId(id, branch)
+      selectedRowsCopy[compositeId] = rowProp
     }
 
     setSelectedRows(selectedRowsCopy)
     setSelectionHistory([
       ...batch.map(
         ({
           row: {
-            original: { id }
+            original: { id, branch }
           }
           // eslint-disable-next-line @typescript-eslint/no-unsafe-return
-        }) => id
+        }) => getCompositeId(id, branch)
       ),
       ...selectionHistory
     ])

webview/src/experiments/components/table/body/Row.tsx

@@ -12,6 +12,7 @@ import { ContextMenu } from '../../../../shared/components/contextMenu/ContextMe
 import { HandlerFunc } from '../../../../util/props'
 import { ExperimentsState } from '../../../store'
 import { toggleExperiment, toggleStarred } from '../../../util/messages'
+import { getCompositeId } from '../../../util/rows'
 
 export type BatchSelectionProp = {
   batchRowSelection: (prop: RowProp) => void
@@ -24,14 +25,15 @@ export const RowContent: React.FC<
     (state: ExperimentsState) => state.tableData.changes
   )
   const { getVisibleCells, original, index, getIsExpanded, subRows } = row
-  const { displayColor, error, starred, id, status, selected } = original
+  const { branch, displayColor, error, starred, id, status, selected } =
+    original
   const [firstCell, ...cells] = getVisibleCells()
   const isWorkspace = id === EXPERIMENT_WORKSPACE_ID
   const changesIfWorkspace = isWorkspace ? changes : undefined
 
   const { toggleRowSelected, selectedRows } = useContext(RowSelectionContext)
 
-  const isRowSelected = !!selectedRows[id]
+  const isRowSelected = !!selectedRows[getCompositeId(id, branch)]
 
   const toggleRowSelection = useCallback<HandlerFunc<HTMLElement>>(
     args => {
@@ -52,7 +54,12 @@ export const RowContent: React.FC<
       subRows?.filter(subRow => subRow.original.selected).length ?? 0
 
     const selections =
-      subRows?.filter(subRow => selectedRows[subRow.original.id]).length ?? 0
+      subRows?.filter(
+        subRow =>
+          selectedRows[
+            getCompositeId(subRow.original.id, subRow.original.branch)
+          ]
+      ).length ?? 0
 
     return {
       plotSelections,

webview/src/experiments/components/table/body/RowContextMenu.tsx

@@ -14,6 +14,7 @@ import { MessagesMenu } from '../../../../shared/components/messagesMenu/Message
 import { MessagesMenuOptionProps } from '../../../../shared/components/messagesMenu/MessagesMenuOption'
 import { cond } from '../../../../util/helpers'
 import { ExperimentsState } from '../../../store'
+import { getCompositeId } from '../../../util/rows'
 
 const experimentMenuOption = (
   payload: string | string[] | { id: string; executor?: string | null }[],
@@ -107,47 +108,50 @@ const getMultiSelectMenuOptions = (
     unstarredExperimentIds
   } = collectDisabledOptions(selectedRowsList, hasRunningWorkspaceExperiment)
 
-  const toggleStarOption = (ids: string[], label: string) =>
-    experimentMenuOption(
-      ids,
+  const toggleStarOption = (ids: string[], label: string) => {
+    return experimentMenuOption(
+      [...new Set(ids)],
       label,
       MessageFromWebviewType.TOGGLE_EXPERIMENT_STAR,
       ids.length === 0
     )
+  }
+
+  const ids = [...new Set(selectedIds)]
 
   return [
     toggleStarOption(unstarredExperimentIds, 'Star'),
     toggleStarOption(starredExperimentIds, 'Unstar'),
     experimentMenuOption(
-      selectedIds,
+      ids,
       'Plot',
       MessageFromWebviewType.SET_EXPERIMENTS_FOR_PLOTS,
       disablePlotOption,
       true
     ),
     experimentMenuOption(
-      selectedIds,
+      ids,
       'Plot and Show',
       MessageFromWebviewType.SET_EXPERIMENTS_AND_OPEN_PLOTS,
       disablePlotOption,
       false
     ),
     experimentMenuOption(
-      selectedIds,
+      ids,
       'Stop',
       MessageFromWebviewType.STOP_EXPERIMENTS,
       disableStopOption,
       true
     ),
     experimentMenuOption(
-      selectedIds,
+      ids,
       'Push Selected',
       MessageFromWebviewType.PUSH_EXPERIMENT,
       disableExperimentOnlyOption,
       true
     ),
     experimentMenuOption(
-      selectedIds,
+      ids,
       'Remove Selected',
       MessageFromWebviewType.REMOVE_EXPERIMENT,
       disableExperimentOnlyOption,
@@ -287,6 +291,7 @@ const getSingleSelectMenuOptions = (
 
 const getContextMenuOptions = (
   id: string,
+  branch: string | undefined,
   isWorkspace: boolean,
   projectHasCheckpoints: boolean,
   hasRunningWorkspaceExperiment: boolean,
@@ -296,7 +301,7 @@ const getContextMenuOptions = (
   starred?: boolean,
   executor?: string | null
 ) => {
-  const isFromSelection = !!selectedRows[id]
+  const isFromSelection = !!selectedRows[getCompositeId(id, branch)]
   const selectedRowsList = Object.values(selectedRows).filter(
     value => value !== undefined
   ) as RowProp[]
@@ -324,7 +329,7 @@ const getContextMenuOptions = (
 
 export const RowContextMenu: React.FC<RowProp> = ({
   row: {
-    original: { status, starred, id, executor },
+    original: { branch, status, starred, id, executor },
     depth
   }
 }) => {
@@ -339,6 +344,7 @@ export const RowContextMenu: React.FC<RowProp> = ({
   const contextMenuOptions = useMemo(() => {
     return getContextMenuOptions(
       id,
+      branch,
       isWorkspace,
       projectHasCheckpoints,
       hasRunningWorkspaceExperiment,
@@ -349,6 +355,7 @@ export const RowContextMenu: React.FC<RowProp> = ({
       executor
     )
   }, [
+    branch,
     executor,
     status,
     starred,

webview/src/experiments/util/rows.ts

@@ -0,0 +1 @@
+export const getCompositeId = (id: string, branch = '') => `${id}-${branch}`