Show duplicate revisions when experiment finishes running in the workspace (#3641)

Author: mattseddon

extension/src/experiments/index.ts

@@ -21,6 +21,10 @@ import {
   pickFiltersToRemove
 } from './model/filterBy/quickPick'
 import { Color } from './model/status/colors'
+import {
+  FetchedExperiment,
+  hasFinishedWorkspaceExperiment
+} from './model/status/collect'
 import { UNSELECTED } from './model/status'
 import { starredSort } from './model/sortBy/constants'
 import { pickSortsToRemove, pickSortToAdd } from './model/sortBy/quickPick'
@@ -240,6 +244,17 @@ export class Experiments extends BaseRepository<TableData> {
     return status
   }
 
+  public checkForFinishedWorkspaceExperiment(
+    fetchedExperiments: FetchedExperiment[]
+  ) {
+    if (!hasFinishedWorkspaceExperiment(fetchedExperiments)) {
+      return
+    }
+
+    this.experiments.unselectWorkspace()
+    this.notifyChanged()
+  }
+
   public getSorts() {
     return this.experiments.getSorts()
   }
@@ -462,10 +477,6 @@ export class Experiments extends BaseRepository<TableData> {
     return this.experiments.getRevisions()
   }
 
-  public getSelectedExperiments() {
-    return this.experiments.getSelectedExperiments()
-  }
-
   public async modifyExperimentParamsAndRun(
     commandId: ModifiedExperimentAndRunCommandId,
     experimentId?: string

extension/src/experiments/model/index.ts

@@ -23,11 +23,7 @@ import {
   isRunningInQueue,
   RunningExperiment
 } from '../webview/contract'
-import {
-  definedAndNonEmpty,
-  reorderListSubset,
-  reorderObjectList
-} from '../../util/array'
+import { definedAndNonEmpty, reorderListSubset } from '../../util/array'
 import {
   ExperimentsOutput,
   EXPERIMENT_WORKSPACE_ID
@@ -141,8 +137,8 @@ export class ExperimentsModel extends ModelWithPersistence {
 
     const current = this.coloredStatus[id]
     if (current) {
-      this.unassignColor(current)
       this.coloredStatus[id] = UNSELECTED
+      this.unassignColor(current)
     } else if (this.availableColors.length > 0) {
       this.coloredStatus[id] = this.availableColors.shift() as Color
     }
@@ -151,6 +147,10 @@ export class ExperimentsModel extends ModelWithPersistence {
     return this.coloredStatus[id]
   }
 
+  public unselectWorkspace() {
+    this.coloredStatus[EXPERIMENT_WORKSPACE_ID] = UNSELECTED
+  }
+
   public hasRunningExperiment() {
     return this.running.length > 0
   }
@@ -247,10 +247,6 @@ export class ExperimentsModel extends ModelWithPersistence {
     return this.getSelectedFromList(() => this.getCombinedList())
   }
 
-  public getSelectedExperiments() {
-    return this.getSelectedFromList(() => this.getExperimentsAndQueued())
-  }
-
   public setSelected(selectedExperiments: Experiment[]) {
     if (tooManySelected(selectedExperiments)) {
       selectedExperiments = limitToMaxSelected(selectedExperiments)
@@ -474,6 +470,10 @@ export class ExperimentsModel extends ModelWithPersistence {
   }
 
   private unassignColor(color: Color) {
+    if (new Set(Object.values(this.coloredStatus)).has(color)) {
+      return
+    }
+
     this.availableColors.unshift(color)
     this.availableColors = reorderListSubset(
       this.availableColors,
@@ -536,10 +536,10 @@ export class ExperimentsModel extends ModelWithPersistence {
       }
     }
 
-    return reorderObjectList<SelectedExperimentWithColor>(
-      copyOriginalColors(),
-      acc,
-      'displayColor'
-    )
+    return copyOriginalColors()
+      .flatMap(orderedItem =>
+        acc.filter(item => item.displayColor === orderedItem)
+      )
+      .filter(Boolean)
   }
 }

extension/src/experiments/model/status/collect.test.ts

@@ -246,7 +246,7 @@ describe('collectColoredStatus', () => {
     expect(availableColors).toStrictEqual(colors.slice(1))
   })
 
-  it("should reassign the workspace's color when an experiment finishes running in the workspace", () => {
+  it("should duplicate the workspace's color when an experiment finishes running in the workspace", () => {
     const colors = copyOriginalColors()
     const { availableColors, coloredStatus } = collectColoredStatus(
       [
@@ -269,7 +269,7 @@ describe('collectColoredStatus', () => {
     )
     expect(coloredStatus).toStrictEqual({
       'exp-1': colors[0],
-      workspace: UNSELECTED
+      workspace: colors[0]
     })
 
     expect(availableColors).toStrictEqual(colors.slice(1))

extension/src/experiments/model/status/collect.ts

@@ -77,7 +77,7 @@ const removeUnselectedExperimentRunningInWorkspace = (
   }
 }
 
-const reassignFinishedWorkspaceExperiment = (
+const duplicateFinishedWorkspaceExperiment = (
   coloredStatus: ColoredStatus,
   finishedRunning: { [id: string]: string }
 ): void => {
@@ -87,7 +87,6 @@ const reassignFinishedWorkspaceExperiment = (
     }
 
     coloredStatus[id] = coloredStatus.workspace
-    coloredStatus.workspace = UNSELECTED
   }
 }
 
@@ -143,7 +142,7 @@ export const collectColoredStatus = (
     removeUnselectedExperimentRunningInWorkspace(coloredStatus, experiment)
   }
 
-  reassignFinishedWorkspaceExperiment(coloredStatus, finishedRunning)
+  duplicateFinishedWorkspaceExperiment(coloredStatus, finishedRunning)
 
   return { availableColors, coloredStatus }
 }
@@ -310,3 +309,28 @@ export const collectFinishedRunningExperiments = (
   }
   return acc
 }
+
+export type FetchedExperiment = { id?: string; displayColor: Color }
+
+export const hasFinishedWorkspaceExperiment = (
+  fetchedExperiments: FetchedExperiment[]
+): boolean => {
+  let workspace: FetchedExperiment | undefined
+  const nonWorkspace: FetchedExperiment[] = []
+
+  for (const revision of fetchedExperiments) {
+    if (revision.id === EXPERIMENT_WORKSPACE_ID) {
+      workspace = revision
+      continue
+    }
+    nonWorkspace.push(revision)
+  }
+
+  if (!workspace) {
+    return false
+  }
+
+  return nonWorkspace.some(
+    ({ displayColor }) => displayColor === workspace?.displayColor
+  )
+}

extension/src/plots/model/index.ts

@@ -196,16 +196,17 @@ export class PlotsModel extends ModelWithPersistence {
   }
 
   public getSelectedRevisionDetails() {
-    return this.experiments.getSelectedRevisions().map(exp => {
+    const selectedRevisions: Revision[] = []
+    for (const experiment of this.experiments.getSelectedRevisions()) {
       const { commit, displayName, label, displayColor, logicalGroupName, id } =
-        exp
+        experiment
       const revision: Revision = {
         displayColor,
         errors: this.errors.getRevisionErrors(label),
         fetched: this.fetchedRevs.has(label),
         firstThreeColumns: getRevisionFirstThreeColumns(
           this.experiments.getFirstThreeColumnOrder(),
-          exp
+          experiment
         ),
         group: logicalGroupName,
         id,
@@ -215,8 +216,10 @@ export class PlotsModel extends ModelWithPersistence {
       if (commit) {
         revision.commit = displayName
       }
-      return revision
-    })
+      selectedRevisions.push(revision)
+    }
+
+    return selectedRevisions
   }
 
   public getTemplatePlots(

extension/src/plots/webview/messages.ts

@@ -57,9 +57,10 @@ export class WebviewMessages {
     this.updateData = updateData
   }
 
-  public sendWebviewMessage() {
+  public async sendWebviewMessage() {
     const selectedRevisions = this.plots.getSelectedRevisionDetails()
-    void this.getWebview()?.show({
+
+    await this.getWebview()?.show({
       cliError: this.errors.getCliError()?.error || null,
       comparison: this.getComparisonPlots(),
       custom: this.getCustomPlots(),
@@ -69,6 +70,10 @@ export class WebviewMessages {
       selectedRevisions,
       template: this.getTemplatePlots(selectedRevisions)
     })
+
+    this.experiments.checkForFinishedWorkspaceExperiment(
+      selectedRevisions.filter(({ fetched }) => fetched)
+    )
   }
 
   public handleMessageFromWebview(message: MessageFromWebview) {