Add context menu to experiment table rows (#1267)

Add context menu payload with item details from webview

Add non-wired context menu

Revert broken label

Add rough functional wiring

Refactor context menu

Create context menu component

Open context menu on row

Use original implementation of the context-menu

Create MessagesMenu component

Update context-menu story

Use messages-menu in the table rows

Process experiment row ctx-menu messages

Remove context-menu webview message

Refactor Row.tsx

Refactor Experiments constructor

Remove commented line

Add tests for apply to workspace and create branch

Fis option labels

Add test for Modify params and queue option

Add test for Experiment Removed message

Prevent default menu from showing up when we close ours

Make runCommand private in extension Experiments repository

Stub internalCommand instance method instead of the prototype

Refactor mock input box utility

Disable row context-menu if the experiment is running

Hide context menu on click

Disable the 'Apply to workspace' and 'Create branch' options when row is workspace or already a branch

Extract row context-menu into a separate component

Do not use the constructor shortcut in order to keep consistency

Refactor option label and contextMenuOptions function

Refactor Experiments repository constructor

Refactor Experiments modifyExperimentParamsAndQueue

Remove unused interface

Disable context-menu if an y experiment is running

Report the command id when throwng a not found error

Undo workspace refactor

Refactor contextMenuOptions function in Row.tsx

extension/src/commands/internal.ts

@@ -41,7 +41,7 @@ export class InternalCommands extends Disposable {
   ): Promise<T> {
     const command = this.commands.get(commandId)
     if (!command) {
-      throw new Error('Unknown command')
+      throw new Error(`Unknown command: ${commandId}`)
     }
 
     return command(...args) as Promise<T>

extension/src/experiments/index.ts

@@ -16,7 +16,11 @@ import { askToDisableAutoApplyFilters } from './toast'
 import { Experiment, MetricOrParamType, TableData } from './webview/contract'
 import { SortDefinition } from './model/sortBy'
 import { ResourceLocator } from '../resourceLocator'
-import { InternalCommands } from '../commands/internal'
+import {
+  AvailableCommands,
+  CommandId,
+  InternalCommands
+} from '../commands/internal'
 import { ExperimentsOutput } from '../cli/reader'
 import { ViewKey } from '../webview/constants'
 import { BaseRepository } from '../webview/repository'
@@ -27,6 +31,8 @@ import { Response } from '../vscode/response'
 import { Title } from '../vscode/title'
 import { sendTelemetryEvent } from '../telemetry'
 import { EventName } from '../telemetry/constants'
+import { Toast } from '../vscode/toast'
+import { getInput } from '../vscode/inputBox'
 import { createTypedAccumulator } from '../util/object'
 
 export const ExperimentsScale = {
@@ -56,6 +62,8 @@ export class Experiments extends BaseRepository<TableData> {
     new EventEmitter<void>()
   )
 
+  private readonly internalCommands: InternalCommands
+
   constructor(
     dvcRoot: string,
     internalCommands: InternalCommands,
@@ -67,6 +75,8 @@ export class Experiments extends BaseRepository<TableData> {
   ) {
     super(dvcRoot, resourceLocator.beaker)
 
+    this.internalCommands = internalCommands
+
     this.onDidChangeExperiments = this.experimentsChanged.event
     this.onDidChangeMetricsOrParams = this.metricsOrParamsChanged.event
 
@@ -96,14 +106,7 @@ export class Experiments extends BaseRepository<TableData> {
     )
 
     this.handleMessageFromWebview()
-
-    const waitForInitialData = this.dispose.track(
-      this.onDidChangeExperiments(() => {
-        this.deferred.resolve()
-        this.dispose.untrack(waitForInitialData)
-        waitForInitialData.dispose()
-      })
-    )
+    this.setupInitialData()
   }
 
   public update() {
@@ -320,6 +323,34 @@ export class Experiments extends BaseRepository<TableData> {
     return this.experiments.getSelectedExperiments()
   }
 
+  public async modifyExperimentParamsAndQueue(experimentId?: string) {
+    const paramsToQueue = await this.pickParamsToQueue(experimentId)
+    if (paramsToQueue) {
+      return this.runCommand(
+        AvailableCommands.EXPERIMENT_QUEUE,
+        ...paramsToQueue
+      )
+    }
+  }
+
+  private setupInitialData() {
+    const waitForInitialData = this.dispose.track(
+      this.onDidChangeExperiments(() => {
+        this.deferred.resolve()
+        this.dispose.untrack(waitForInitialData)
+        waitForInitialData.dispose()
+      })
+    )
+  }
+
+  private async runCommand(commandId: CommandId, ...args: string[]) {
+    await Toast.showOutput(
+      this.internalCommands.executeCommand(commandId, this.dvcRoot, ...args)
+    )
+
+    return this.notifyChanged()
+  }
+
   private notifyChanged(data?: ExperimentsOutput) {
     this.experimentsChanged.fire(data)
     this.notifyMetricsOrParamsChanged()
@@ -362,6 +393,14 @@ export class Experiments extends BaseRepository<TableData> {
             return this.addColumnSort(message.payload)
           case MessageFromWebviewType.COLUMN_SORT_REMOVED:
             return this.removeColumnSort(message.payload)
+          case MessageFromWebviewType.EXPERIMENT_APPLIED_TO_WORKSPACE:
+            return this.applyExperimentToWorkspace(message.payload)
+          case MessageFromWebviewType.BRANCH_CREATED_FROM_EXPERIMENT:
+            return this.createBranchFromExperiment(message.payload)
+          case MessageFromWebviewType.EXPERIMENT_QUEUE_AND_PARAMS_VARIED:
+            return this.modifyExperimentParamsAndQueue(message.payload)
+          case MessageFromWebviewType.EXPERIMENT_REMOVED:
+            return this.removeExperiment(message.payload)
           default:
             Logger.error(`Unexpected message: ${JSON.stringify(message)}`)
         }
@@ -416,6 +455,25 @@ export class Experiments extends BaseRepository<TableData> {
     return this.notifyChanged()
   }
 
+  private async createBranchFromExperiment(experimentId: string) {
+    const input = await getInput(Title.ENTER_BRANCH_NAME)
+    if (input) {
+      return this.runCommand(
+        AvailableCommands.EXPERIMENT_BRANCH,
+        experimentId,
+        input
+      )
+    }
+  }
+
+  private applyExperimentToWorkspace(experimentId: string) {
+    return this.runCommand(AvailableCommands.EXPERIMENT_APPLY, experimentId)
+  }
+
+  private removeExperiment(experimentId: string) {
+    return this.runCommand(AvailableCommands.EXPERIMENT_REMOVE, experimentId)
+  }
+
   private async checkAutoApplyFilters(...filterIdsToRemove: string[]) {
     if (this.experiments.canAutoApplyFilters(...filterIdsToRemove)) {
       return

extension/src/experiments/workspace.ts

@@ -2,8 +2,8 @@ import { EventEmitter, Memento } from 'vscode'
 import { Experiments } from '.'
 import { TableData } from './webview/contract'
 import {
-  CommandId,
   AvailableCommands,
+  CommandId,
   InternalCommands
 } from '../commands/internal'
 import { ResourceLocator } from '../resourceLocator'

extension/src/test/suite/experiments/index.test.ts

@@ -23,6 +23,7 @@ import {
   closeAllEditors,
   experimentsUpdatedEvent,
   extensionUri,
+  getInputBoxEvent,
   getMessageReceivedEmitter
 } from '../util'
 import { buildMockMemento, dvcDemoPath } from '../../util'
@@ -40,7 +41,7 @@ import { MetricsAndParamsModel } from '../../../experiments/metricsAndParams/mod
 import { MessageFromWebviewType } from '../../../webview/contract'
 import { ExperimentsModel } from '../../../experiments/model'
 import { copyOriginalColors } from '../../../experiments/model/status/colors'
-import { InternalCommands } from '../../../commands/internal'
+import { AvailableCommands, InternalCommands } from '../../../commands/internal'
 import { FileSystemData } from '../../../fileSystem/data'
 import { ExperimentsData } from '../../../experiments/data'
 import { WEBVIEW_TEST_TIMEOUT } from '../timeouts'
@@ -327,6 +328,116 @@ suite('Experiments Test Suite', () => {
       )
     })
 
+    describe('Handling webview messages that run extension commands', () => {
+      const setupExperimentsAndMockCommands = () => {
+        const { experiments, internalCommands } = buildExperiments(
+          disposable,
+          expShowFixture
+        )
+        const mockExecuteCommand = stub(
+          internalCommands,
+          'executeCommand'
+        ).resolves(undefined)
+
+        return { experiments, mockExecuteCommand }
+      }
+
+      it('should be able to handle a message to apply an experiment to workspace', async () => {
+        const { experiments, mockExecuteCommand } =
+          setupExperimentsAndMockCommands()
+
+        const webview = await experiments.showWebview()
+        const mockMessageReceived = getMessageReceivedEmitter(webview)
+        const mockExperimentId = 'mock-experiment-id'
+
+        mockMessageReceived.fire({
+          payload: mockExperimentId,
+          type: MessageFromWebviewType.EXPERIMENT_APPLIED_TO_WORKSPACE
+        })
+
+        expect(mockExecuteCommand).to.be.calledOnce
+        expect(mockExecuteCommand).to.be.calledWithExactly(
+          AvailableCommands.EXPERIMENT_APPLY,
+          dvcDemoPath,
+          mockExperimentId
+        )
+      })
+
+      it('should be able to handle a message to create a branch from an experiment', async () => {
+        const { experiments, mockExecuteCommand } =
+          setupExperimentsAndMockCommands()
+
+        const mockBranch = 'mock-branch-input'
+        const inputEvent = getInputBoxEvent(mockBranch)
+
+        const webview = await experiments.showWebview()
+        const mockMessageReceived = getMessageReceivedEmitter(webview)
+        const mockExperimentId = 'mock-experiment-id'
+
+        mockMessageReceived.fire({
+          payload: mockExperimentId,
+          type: MessageFromWebviewType.BRANCH_CREATED_FROM_EXPERIMENT
+        })
+
+        await inputEvent
+        expect(mockExecuteCommand).to.be.calledOnce
+        expect(mockExecuteCommand).to.be.calledWithExactly(
+          AvailableCommands.EXPERIMENT_BRANCH,
+          dvcDemoPath,
+          mockExperimentId,
+          mockBranch
+        )
+      })
+
+      it('should be able to handle a message to modify the params and queue of an experiment', async () => {
+        const { experiments, mockExecuteCommand } =
+          setupExperimentsAndMockCommands()
+
+        const mockParams = ['param1', 'param2']
+
+        stub(Experiments.prototype, 'pickParamsToQueue').resolves(mockParams)
+
+        const webview = await experiments.showWebview()
+        const mockMessageReceived = getMessageReceivedEmitter(webview)
+        const mockExperimentId = 'mock-experiment-id'
+        const tableChangePromise = experimentsUpdatedEvent(experiments)
+
+        mockMessageReceived.fire({
+          payload: mockExperimentId,
+          type: MessageFromWebviewType.EXPERIMENT_QUEUE_AND_PARAMS_VARIED
+        })
+
+        await tableChangePromise
+        expect(mockExecuteCommand).to.be.calledOnce
+        expect(mockExecuteCommand).to.be.calledWithExactly(
+          AvailableCommands.EXPERIMENT_QUEUE,
+          dvcDemoPath,
+          ...mockParams
+        )
+      })
+
+      it('should be able to handle a message to remove an experiment', async () => {
+        const { experiments, mockExecuteCommand } =
+          setupExperimentsAndMockCommands()
+
+        const webview = await experiments.showWebview()
+        const mockMessageReceived = getMessageReceivedEmitter(webview)
+        const mockExperimentId = 'mock-experiment-id'
+
+        mockMessageReceived.fire({
+          payload: mockExperimentId,
+          type: MessageFromWebviewType.EXPERIMENT_REMOVED
+        })
+
+        expect(mockExecuteCommand).to.be.calledOnce
+        expect(mockExecuteCommand).to.be.calledWithExactly(
+          AvailableCommands.EXPERIMENT_REMOVE,
+          dvcDemoPath,
+          mockExperimentId
+        )
+      })
+    })
+
     it('should be able to sort', async () => {
       const { internalCommands } = buildInternalCommands(disposable)
 

extension/src/test/suite/util.ts

@@ -187,3 +187,13 @@ export const getMessageReceivedEmitter = (
   webview: BaseWebview<PlotsData | TableData>
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
 ): EventEmitter<MessageFromWebview> => (webview as any).messageReceived
+
+export const getInputBoxEvent = (mockInputValue: string) => {
+  const mockInput = stub(window, 'showInputBox')
+  return new Promise(resolve =>
+    mockInput.callsFake(() => {
+      resolve(mockInputValue)
+      return Promise.resolve(mockInputValue)
+    })
+  )
+}

extension/src/webview/contract.ts

@@ -17,6 +17,10 @@ export enum MessageFromWebviewType {
   COLUMN_SORTED = 'column-sorted',
   COLUMN_SORT_REMOVED = 'column-sort-removed',
   EXPERIMENT_TOGGLED = 'experiment-toggled',
+  EXPERIMENT_APPLIED_TO_WORKSPACE = 'experiment-applied-to-workspace',
+  BRANCH_CREATED_FROM_EXPERIMENT = 'branch-created-from-experiment',
+  EXPERIMENT_QUEUE_AND_PARAMS_VARIED = 'experiment-queue-and-params-varied',
+  EXPERIMENT_REMOVED = 'experiment-removed',
   METRIC_TOGGLED = 'metric-toggled',
   PLOTS_COMPARISON_REORDERED = 'plots-comparison-reordered',
   PLOTS_METRICS_REORDERED = 'plots-metrics-reordered',
@@ -53,6 +57,22 @@ export type MessageFromWebview =
       type: MessageFromWebviewType.EXPERIMENT_TOGGLED
       payload: string
     }
+  | {
+      type: MessageFromWebviewType.EXPERIMENT_APPLIED_TO_WORKSPACE
+      payload: string
+    }
+  | {
+      type: MessageFromWebviewType.BRANCH_CREATED_FROM_EXPERIMENT
+      payload: string
+    }
+  | {
+      type: MessageFromWebviewType.EXPERIMENT_QUEUE_AND_PARAMS_VARIED
+      payload: string
+    }
+  | {
+      type: MessageFromWebviewType.EXPERIMENT_REMOVED
+      payload: string
+    }
   | {
       type: MessageFromWebviewType.COLUMN_SORTED
       payload: SortDefinition