use usual pattern for registering commands against the experiments tree

extension/package.json

@@ -429,6 +429,18 @@
         "command": "dvc.showSetup",
         "category": "DVC"
       },
+      {
+        "title": "Stop",
+        "command": "dvc.views.experimentsTree.stopExperiment",
+        "category": "DVC",
+        "icon": "$(debug-stop)"
+      },
+      {
+        "title": "Stop",
+        "command": "dvc.views.experiments.stopExperiment",
+        "category": "DVC",
+        "icon": "$(debug-stop)"
+      },
       {
         "title": "Stop All Running Experiments",
         "command": "dvc.stopAllRunningExperiments",
@@ -914,6 +926,14 @@
           "command": "dvc.views.experiments.showLogs",
           "when": "false"
         },
+        {
+          "command": "dvc.views.experimentsTree.stopExperiment",
+          "when": "false"
+        },
+        {
+          "command": "dvc.views.experiments.stopExperiment",
+          "when": "false"
+        },
         {
           "command": "dvc.views.experimentsFilterByTree.removeAllFilters",
           "when": "false"
@@ -1160,6 +1180,11 @@
           "group": "inline",
           "when": "view == dvc.views.experimentsFilterByTree && dvc.commands.available && viewItem != dvcRoot"
         },
+        {
+          "command": "dvc.views.experimentsTree.stopExperiment",
+          "group": "inline@0",
+          "when": "view == dvc.views.experimentsTree && dvc.commands.available && viewItem == running"
+        },
         {
           "command": "dvc.views.experiments.applyExperiment",
           "group": "inline@1",

extension/src/commands/external.ts

@@ -17,7 +17,6 @@ export enum RegisteredCliCommands {
   EXPERIMENT_VIEW_PUSH = 'dvc.views.experiments.pushExperiment',
   EXPERIMENT_VIEW_REMOVE = 'dvc.views.experiments.removeExperiment',
   EXPERIMENT_VIEW_SHOW_LOGS = 'dvc.views.experiments.showLogs',
-  EXPERIMENT_VIEW_STOP = 'dvc.views.experiments.stopQueueExperiment',
 
   EXPERIMENT_VIEW_QUEUE = 'dvc.views.experiments.queueExperiment',
   EXPERIMENT_VIEW_RESUME = 'dvc.views.experiments.resumeCheckpointExperiment',
@@ -63,6 +62,7 @@ export enum RegisteredCommands {
   EXPERIMENT_SORTS_REMOVE_ALL = 'dvc.views.experimentsSortByTree.removeAllSorts',
   EXPERIMENT_STOP = 'dvc.stopExperiments',
   EXPERIMENT_TOGGLE = 'dvc.views.experiments.toggleStatus',
+  EXPERIMENT_VIEW_STOP = 'dvc.views.experiments.stopExperiment',
   STOP_EXPERIMENTS = 'dvc.stopAllRunningExperiments',
 
   PLOTS_PATH_TOGGLE = 'dvc.views.plotsPathsTree.toggleStatus',

extension/src/experiments/commands/register.ts

@@ -115,6 +115,12 @@ const registerExperimentNameCommands = (
     ({ dvcRoot, ids }: { dvcRoot: string; ids: string[] }) =>
       experiments.runCommand(AvailableCommands.EXP_REMOVE, dvcRoot, ...ids)
   )
+
+  internalCommands.registerExternalCommand(
+    RegisteredCommands.EXPERIMENT_VIEW_STOP,
+    ({ dvcRoot, ids }: { dvcRoot: string; ids: string[] }) =>
+      experiments.stopExperiments(dvcRoot, ...ids)
+  )
 }
 
 const registerExperimentInputCommands = (

extension/src/experiments/index.ts

@@ -42,11 +42,7 @@ import { DecorationProvider } from './model/decorationProvider'
 import { starredFilter } from './model/filterBy/constants'
 import { ResourceLocator } from '../resourceLocator'
 import { AvailableCommands, InternalCommands } from '../commands/internal'
-import {
-  EXPERIMENT_WORKSPACE_ID,
-  Executor,
-  ExpShowOutput
-} from '../cli/dvc/contract'
+import { ExpShowOutput } from '../cli/dvc/contract'
 import { ViewKey } from '../webview/constants'
 import { BaseRepository } from '../webview/repository'
 import { Title } from '../vscode/title'
@@ -491,34 +487,30 @@ export class Experiments extends BaseRepository<TableData> {
     return this.notifyChanged()
   }
 
-  // eslint-disable-next-line sonarjs/cognitive-complexity
   public stopExperiments(ids: string[]) {
-    const experiments = this.experiments.getExperiments()
-    const idSet = new Set(ids)
-    let pidRunningInWorkspace
-    const runningInQueueIds = new Set<string>()
-    for (const { executor, id, executorPid } of experiments) {
-      if (!idSet.has(id)) {
-        continue
-      }
-      if (executor === EXPERIMENT_WORKSPACE_ID && executorPid) {
-        pidRunningInWorkspace = executorPid
-      }
-      if (executor === Executor.DVC_TASK) {
-        runningInQueueIds.add(id)
-      }
-    }
+    const { runningInQueueIds, workspaceDetails } =
+      this.experiments.getStopDetails(ids)
 
-    if (runningInQueueIds.size > 0) {
-      void this.internalCommands.executeCommand(
-        AvailableCommands.QUEUE_KILL,
-        this.dvcRoot,
-        ...runningInQueueIds
+    const promises: Promise<string | void>[] = []
+
+    if (runningInQueueIds) {
+      promises.push(
+        this.internalCommands.executeCommand(
+          AvailableCommands.QUEUE_KILL,
+          this.dvcRoot,
+          ...runningInQueueIds
+        )
+      )
+    }
+    if (workspaceDetails) {
+      promises.push(
+        stopWorkspaceExperiment(workspaceDetails.id, workspaceDetails.pid)
       )
-      if (pidRunningInWorkspace) {
-        void stopWorkspaceExperiment(pidRunningInWorkspace)
-      }
     }
+
+    return Toast.showOutput(
+      Promise.all(promises).then(output => output.filter(Boolean).join('\n'))
+    )
   }
 
   public hasRunningExperiment() {
@@ -597,7 +589,6 @@ export class Experiments extends BaseRepository<TableData> {
       () => this.getWebview(),
       () => this.notifyChanged(),
       () => this.selectColumns(),
-      (ids: string[]) => this.stopExperiments(ids),
       () =>
         this.internalCommands.executeCommand(
           AvailableCommands.STAGE_LIST,

extension/src/experiments/model/collect.ts

@@ -10,7 +10,8 @@ import {
   Experiment,
   CommitData,
   RunningExperiment,
-  isQueued
+  isQueued,
+  isRunning
 } from '../webview/contract'
 import {
   EXPERIMENT_WORKSPACE_ID,
@@ -28,6 +29,7 @@ import { RegisteredCommands } from '../../commands/external'
 import { Resource } from '../../resourceLocator'
 import { shortenForLabel } from '../../util/string'
 import { COMMITS_SEPARATOR } from '../../cli/git/constants'
+import { createValidInteger } from '../../util/number'
 
 export type ExperimentItem = {
   command?: {
@@ -205,33 +207,36 @@ const getExecutor = (experiment: Experiment): Executor => {
 }
 
 const collectExecutorInfo = (
-  acc: ExperimentsAccumulator,
   experiment: Experiment,
   executor: ExecutorState
-  // eslint-disable-next-line sonarjs/cognitive-complexity
 ): void => {
   if (!executor) {
     return
   }
 
-  const { name, state, local } = executor
+  const { name, state } = executor
 
   if (name && state === ExperimentStatus.RUNNING) {
     experiment.executor = name
   }
   if (state && state !== ExperimentStatus.SUCCESS) {
     experiment.status = state
   }
-  if (local?.pid) {
-    experiment.executorPid = local.pid
-  }
+}
 
-  if (experiment.status === ExperimentStatus.RUNNING) {
-    acc.runningExperiments.push({
-      executor: getExecutor(experiment),
-      id: experiment.id
-    })
+const collectRunningExperiment = (
+  acc: ExperimentsAccumulator,
+  experiment: Experiment,
+  executor: ExecutorState
+): void => {
+  if (!isRunning(experiment.status)) {
+    return
   }
+  acc.runningExperiments.push({
+    executor: getExecutor(experiment),
+    id: experiment.id,
+    pid: createValidInteger(executor?.local?.pid)
+  })
 }
 
 const collectExpRange = (
@@ -265,7 +270,8 @@ const collectExpRange = (
     experiment.description = `[${name}]`
   }
 
-  collectExecutorInfo(acc, experiment, executor)
+  collectExecutorInfo(experiment, executor)
+  collectRunningExperiment(acc, experiment, executor)
 
   addToMapArray(acc.experimentsByCommit, baseline.id, experiment)
 }
@@ -406,3 +412,33 @@ export const collectOrderedCommitsAndExperiments = (
   }
   return acc
 }
+
+export const collectRunningQueueTaskIds = (
+  ids: Set<string>,
+  runningExperiments: RunningExperiment[]
+): string[] | undefined => {
+  const runningInQueueIds = new Set<string>()
+  for (const { executor, id } of runningExperiments) {
+    if (!ids.has(id)) {
+      continue
+    }
+    if (executor === Executor.DVC_TASK) {
+      runningInQueueIds.add(id)
+    }
+  }
+  return runningInQueueIds.size > 0 ? [...runningInQueueIds] : undefined
+}
+
+export const collectWorkspaceExecutorPid = (
+  ids: Set<string>,
+  runningExperiments: RunningExperiment[]
+): { pid: number; id: string } | undefined => {
+  for (const { executor, id, pid } of runningExperiments) {
+    if (!ids.has(id)) {
+      continue
+    }
+    if (executor === EXPERIMENT_WORKSPACE_ID && pid) {
+      return { id, pid }
+    }
+  }
+}

extension/src/experiments/model/index.ts

@@ -3,7 +3,9 @@ import { SortDefinition, sortExperiments } from './sortBy'
 import { FilterDefinition, filterExperiment, getFilterId } from './filterBy'
 import {
   collectExperiments,
-  collectOrderedCommitsAndExperiments
+  collectOrderedCommitsAndExperiments,
+  collectRunningQueueTaskIds,
+  collectWorkspaceExecutorPid as collectWorkspaceRunDetails
 } from './collect'
 import {
   collectColoredStatus,
@@ -340,6 +342,15 @@ export class ExperimentsModel extends ModelWithPersistence {
     )
   }
 
+  public getStopDetails(idsToStop: string[]) {
+    const running = [...this.running]
+    const ids = new Set(idsToStop)
+    return {
+      runningInQueueIds: collectRunningQueueTaskIds(ids, running),
+      workspaceDetails: collectWorkspaceRunDetails(ids, running)
+    }
+  }
+
   public getRowData() {
     return [
       this.addDetails(this.workspace),

extension/src/experiments/model/tree.ts

@@ -127,7 +127,8 @@ export class ExperimentsTree
     const callCommandWithSelected = async (
       command:
         | RegisteredCliCommands.EXPERIMENT_VIEW_REMOVE
-        | RegisteredCliCommands.EXPERIMENT_VIEW_PUSH,
+        | RegisteredCliCommands.EXPERIMENT_VIEW_PUSH
+        | RegisteredCommands.EXPERIMENT_VIEW_STOP,
       experimentItem: ExperimentItem | string,
       types: ExperimentType[]
     ) => {
@@ -162,6 +163,16 @@ export class ExperimentsTree
           [ExperimentType.EXPERIMENT]
         )
     )
+
+    commands.registerCommand(
+      'dvc.views.experimentsTree.stopExperiment',
+      (experimentItem: ExperimentItem) =>
+        callCommandWithSelected(
+          RegisteredCommands.EXPERIMENT_VIEW_STOP,
+          experimentItem,
+          [ExperimentType.RUNNING]
+        )
+    )
   }
 
   private async getRootElements() {

extension/src/experiments/processExecution/index.ts

@@ -1,16 +1,16 @@
-import { Toast } from '../../vscode/toast'
 import { processExists, stopProcesses } from '../../process/execution'
 
-export const stopWorkspaceExperiment = async (pid: number) => {
+export const stopWorkspaceExperiment = async (id: string, pid: number) => {
   if (!(await processExists(pid))) {
-    return
+    return `process executing ${id} was not found.`
   }
 
-  void Toast.showOutput(
-    stopProcesses([pid]).then(stopped =>
-      stopped
-        ? 'The experiment running in the workspace was stopped.'
-        : 'Failed to stop the experiment running in the workspace.'
-    )
-  )
+  const failedToStop = `failed to kill ${id}.`
+
+  try {
+    const stopped = await stopProcesses([pid])
+    return stopped ? `${id} has been killed.` : failedToStop
+  } catch {
+    return failedToStop
+  }
 }

extension/src/experiments/webview/contract.ts

@@ -16,7 +16,11 @@ export interface DepColumns {
   [path: string]: ValueWithChanges
 }
 
-export type RunningExperiment = { executor: Executor; id: string }
+export type RunningExperiment = {
+  executor: Executor
+  id: string
+  pid?: number | null
+}
 
 export type CommitData = {
   author: string
@@ -33,7 +37,6 @@ export type Experiment = {
   description?: string
   error?: string
   executor?: Executor
-  executorPid?: number
   id: string
   label: string
   metrics?: MetricOrParamColumns