Refresh cached plots data on every update

extension/src/experiments/index.ts

@@ -453,12 +453,6 @@ export class Experiments extends BaseRepository<TableData> {
     return experiment?.name || experiment?.label
   }
 
-  public getMutableRevisions() {
-    return this.experiments.getMutableRevisions(
-      this.checkpoints.hasCheckpoints()
-    )
-  }
-
   public getRevisions() {
     return this.experiments.getRevisions()
   }

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

@@ -1,10 +1,7 @@
-import { collectExperiments, collectMutableRevisions } from './collect'
+import { collectExperiments } from './collect'
 import { Experiment } from '../webview/contract'
 import modifiedFixture from '../../test/fixtures/expShow/modified/output'
-import {
-  ExperimentStatus,
-  EXPERIMENT_WORKSPACE_ID
-} from '../../cli/dvc/contract'
+import { EXPERIMENT_WORKSPACE_ID } from '../../cli/dvc/contract'
 import { COMMITS_SEPARATOR } from '../../cli/git/constants'
 
 describe('collectExperiments', () => {
@@ -233,62 +230,3 @@ describe('collectExperiments', () => {
     ).toStrictEqual(3)
   })
 })
-
-describe('collectMutableRevisions', () => {
-  const baseExperiments = [
-    { label: 'branch-A', selected: false, status: ExperimentStatus.SUCCESS },
-    {
-      label: EXPERIMENT_WORKSPACE_ID,
-      selected: false,
-      status: ExperimentStatus.FAILED
-    }
-  ] as Experiment[]
-
-  it('should return the workspace when there is an unselected running checkpoint experiment', () => {
-    const experiments = [
-      {
-        label: 'exp-123',
-        selected: false,
-        status: ExperimentStatus.RUNNING
-      },
-      ...baseExperiments
-    ] as Experiment[]
-
-    const mutableRevisions = collectMutableRevisions(experiments, true)
-    expect(mutableRevisions).toStrictEqual([EXPERIMENT_WORKSPACE_ID])
-  })
-
-  it('should return the workspace when there are no checkpoints', () => {
-    const experiments = [
-      { label: 'branch-A', selected: false, status: ExperimentStatus.SUCCESS },
-      {
-        label: EXPERIMENT_WORKSPACE_ID,
-        selected: false,
-        status: ExperimentStatus.SUCCESS
-      }
-    ] as Experiment[]
-
-    const mutableRevisions = collectMutableRevisions(experiments, false)
-    expect(mutableRevisions).toStrictEqual([EXPERIMENT_WORKSPACE_ID])
-  })
-
-  it('should return all running experiments when there are checkpoints', () => {
-    const experiments = [
-      { label: 'branch-A', selected: false, status: ExperimentStatus.SUCCESS },
-      {
-        label: EXPERIMENT_WORKSPACE_ID,
-        selected: false,
-        status: ExperimentStatus.SUCCESS
-      },
-      { label: 'running-1', selected: false, status: ExperimentStatus.RUNNING },
-      { label: 'running-2', selected: true, status: ExperimentStatus.RUNNING }
-    ] as Experiment[]
-
-    const mutableRevisions = collectMutableRevisions(experiments, false)
-    expect(mutableRevisions).toStrictEqual([
-      EXPERIMENT_WORKSPACE_ID,
-      'running-1',
-      'running-2'
-    ])
-  })
-})

extension/src/experiments/model/collect.ts

@@ -23,7 +23,6 @@ import {
   ExperimentStatus
 } from '../../cli/dvc/contract'
 import { addToMapArray } from '../../util/map'
-import { uniqueValues } from '../../util/array'
 import { RegisteredCommands } from '../../commands/external'
 import { Resource } from '../../resourceLocator'
 import { shortenForLabel } from '../../util/string'
@@ -408,31 +407,6 @@ export const collectExperiments = (
   return acc
 }
 
-const getDefaultMutableRevision = (): string[] => [EXPERIMENT_WORKSPACE_ID]
-
-const collectMutableRevision = (
-  acc: string[],
-  { label, status }: Experiment,
-  hasCheckpoints: boolean
-) => {
-  if (isRunning(status) && !hasCheckpoints) {
-    acc.push(label)
-  }
-}
-
-export const collectMutableRevisions = (
-  experiments: Experiment[],
-  hasCheckpoints: boolean
-): string[] => {
-  const acc = getDefaultMutableRevision()
-
-  for (const experiment of experiments) {
-    collectMutableRevision(acc, experiment, hasCheckpoints)
-  }
-
-  return uniqueValues(acc)
-}
-
 type DeletableExperimentAccumulator = { [dvcRoot: string]: Set<string> }
 
 const initializeAccumulatorRoot = (

extension/src/experiments/model/index.ts

@@ -6,7 +6,7 @@ import {
   splitExperimentsByFilters,
   getFilterId
 } from './filterBy'
-import { collectExperiments, collectMutableRevisions } from './collect'
+import { collectExperiments } from './collect'
 import {
   collectFiltered,
   collectFilteredCounts,
@@ -264,13 +264,6 @@ export class ExperimentsModel extends ModelWithPersistence {
     return this.getCombinedList().map(({ label }) => label)
   }
 
-  public getMutableRevisions(hasCheckpoints: boolean) {
-    return collectMutableRevisions(
-      this.getRecordsWithoutCheckpoints(),
-      hasCheckpoints
-    )
-  }
-
   public getSelectedRevisions() {
     return this.getSelectedFromList(() => this.getCombinedList())
   }

extension/src/plots/data/collect.ts

@@ -1,8 +1,4 @@
-import {
-  EXPERIMENT_WORKSPACE_ID,
-  ExperimentsOutput,
-  PlotsOutputOrError
-} from '../../cli/dvc/contract'
+import { ExperimentsOutput, PlotsOutputOrError } from '../../cli/dvc/contract'
 import { isDvcError } from '../../cli/dvc/reader'
 import { sortCollectedArray, uniqueValues } from '../../util/array'
 import { isImagePlot, Plot, TemplatePlot } from '../webview/contract'
@@ -76,26 +72,3 @@ export const collectMetricsFiles = (
 
   return sortCollectedArray(metricsFiles)
 }
-
-const collectRev = (acc: string[], rev: string): void => {
-  if (rev !== EXPERIMENT_WORKSPACE_ID && !acc.includes(rev)) {
-    acc.push(rev)
-  }
-}
-
-export const collectRevs = (
-  missingRevisions: string[],
-  mutableRevisions: string[]
-): string[] => {
-  const acc: string[] = []
-
-  for (const missingRevision of missingRevisions) {
-    collectRev(acc, missingRevision)
-  }
-
-  for (const mutableRevision of mutableRevisions) {
-    collectRev(acc, mutableRevision)
-  }
-
-  return [EXPERIMENT_WORKSPACE_ID, ...sortCollectedArray(acc)]
-}

extension/src/plots/data/index.ts

@@ -1,5 +1,5 @@
-import { EventEmitter } from 'vscode'
-import { collectMetricsFiles, collectFiles, collectRevs } from './collect'
+import { Event, EventEmitter } from 'vscode'
+import { collectMetricsFiles, collectFiles } from './collect'
 import {
   EXPERIMENT_WORKSPACE_ID,
   ExperimentsOutput,
@@ -14,10 +14,16 @@ export class PlotsData extends BaseData<{
   data: PlotsOutputOrError
   revs: string[]
 }> {
+  public readonly onDidTrigger: Event<void>
+
   private readonly model: PlotsModel
 
   private metricFiles: string[] = []
 
+  private readonly triggered: EventEmitter<void> = this.dispose.track(
+    new EventEmitter()
+  )
+
   constructor(
     dvcRoot: string,
     internalCommands: InternalCommands,
@@ -37,13 +43,12 @@ export class PlotsData extends BaseData<{
       ['dvc.yaml', 'dvc.lock']
     )
     this.model = model
+    this.onDidTrigger = this.triggered.event
   }
 
   public async update(): Promise<void> {
-    const revs = collectRevs(
-      this.model.getMissingRevisions(),
-      this.model.getMutableRevisions()
-    )
+    this.notifyTriggered()
+    const revs = this.model.getSelectedOrderedCliIds()
 
     const args = this.getArgs(revs)
     const data = await this.internalCommands.executeCommand<PlotsOutputOrError>(
@@ -76,6 +81,10 @@ export class PlotsData extends BaseData<{
     this.collectedFiles = collectFiles(data, this.collectedFiles)
   }
 
+  private notifyTriggered() {
+    this.triggered.fire()
+  }
+
   private getArgs(revs: string[]) {
     const cliWillThrowError = sameContents(revs, [EXPERIMENT_WORKSPACE_ID])
     if (this.model && cliWillThrowError) {

extension/src/plots/index.ts

@@ -14,7 +14,6 @@ import { BaseRepository } from '../webview/repository'
 import { Experiments } from '../experiments'
 import { Resource } from '../resourceLocator'
 import { InternalCommands } from '../commands/internal'
-import { definedAndNonEmpty } from '../util/array'
 import { TEMP_PLOTS_DIR } from '../cli/dvc/constants'
 import { removeDir } from '../fileSystem'
 import { Toast } from '../vscode/toast'
@@ -69,6 +68,7 @@ export class Plots extends BaseRepository<TPlotsData> {
       new PlotsData(dvcRoot, internalCommands, this.plots, updatesPaused)
     )
 
+    this.onDidTriggerDataUpdate()
     this.onDidUpdateData()
     this.waitForInitialData(experiments)
 
@@ -105,10 +105,7 @@ export class Plots extends BaseRepository<TPlotsData> {
     void Toast.infoWithOptions(
       'Attempting to refresh plots for selected experiments.'
     )
-    for (const { revision } of this.plots.getSelectedRevisionDetails()) {
-      this.plots.setupManualRefresh(revision)
-    }
-    void this.data.managedUpdate()
+    this.triggerDataUpdate()
   }
 
   public getChildPaths(path: string | undefined) {
@@ -130,7 +127,7 @@ export class Plots extends BaseRepository<TPlotsData> {
   }
 
   protected sendInitialWebviewData() {
-    return this.fetchMissingOrSendPlots()
+    return this.sendPlots()
   }
 
   private notifyChanged() {
@@ -140,19 +137,18 @@ export class Plots extends BaseRepository<TPlotsData> {
       this.errors.getErrorPaths(selectedRevisions)
     )
     this.pathsChanged.fire()
-    void this.fetchMissingOrSendPlots()
+
+    if (this.plots.requiresUpdate()) {
+      this.triggerDataUpdate()
+      return
+    }
+
+    void this.sendPlots()
   }
 
-  private async fetchMissingOrSendPlots() {
+  private async sendPlots() {
     await this.isReady()
 
-    if (
-      this.paths.hasPaths() &&
-      definedAndNonEmpty(this.plots.getUnfetchedRevisions())
-    ) {
-      void this.data.managedUpdate()
-    }
-
     return this.webviewMessages.sendWebviewMessage()
   }
 
@@ -214,7 +210,7 @@ export class Plots extends BaseRepository<TPlotsData> {
 
   private async initializeData() {
     await this.plots.transformAndSetExperiments()
-    void this.data.managedUpdate()
+    this.triggerDataUpdate()
     await Promise.all([
       this.data.isReady(),
       this.plots.isReady(),
@@ -223,6 +219,19 @@ export class Plots extends BaseRepository<TPlotsData> {
     this.deferred.resolve()
   }
 
+  private triggerDataUpdate() {
+    void this.data.managedUpdate()
+  }
+
+  private onDidTriggerDataUpdate() {
+    const sendCachedDataToWebview = () => {
+      this.plots.resetFetched()
+      this.plots.setComparisonOrder()
+      return this.sendPlots()
+    }
+    this.dispose.track(this.data.onDidTrigger(() => sendCachedDataToWebview()))
+  }
+
   private onDidUpdateData() {
     this.dispose.track(
       this.data.onDidUpdate(async ({ data, revs }) => {