Merge branch 'main' into codeowners

extension/src/cli/runner.ts

@@ -10,7 +10,6 @@ import { getOptions } from './options'
 import { Config } from '../config'
 import { PseudoTerminal } from '../vscode/pseudoTerminal'
 import { createProcess, Process } from '../processExecution'
-import { setContextValue } from '../vscode/context'
 import { StopWatch } from '../util/time'
 import { sendErrorTelemetryEvent, sendTelemetryEvent } from '../telemetry'
 import { EventName } from '../telemetry/constants'
@@ -70,7 +69,6 @@ export class CliRunner extends Disposable implements ICli {
     this.dispose.track(
       this.onDidCompleteProcess(() => {
         this.pseudoTerminal.setBlocked(false)
-        CliRunner.setRunningContext(false)
         this.processOutput.fire(
           '\r\nTerminal will be reused by DVC, press any key to close it\r\n\n'
         )
@@ -101,9 +99,6 @@ export class CliRunner extends Disposable implements ICli {
     )
   }
 
-  private static setRunningContext = (isRunning: boolean) =>
-    setContextValue('dvc.runner.running', isRunning)
-
   public runExperiment(dvcRoot: string, ...args: Args) {
     return this.run(
       dvcRoot,
@@ -207,7 +202,6 @@ export class CliRunner extends Disposable implements ICli {
   }
 
   private startProcess(cwd: string, args: Args) {
-    CliRunner.setRunningContext(true)
     this.pseudoTerminal.setBlocked(true)
     this.processOutput.fire(`Running: dvc ${args.join(' ')}\r\n\n`)
     this.currentProcess = this.createProcess({

extension/src/context.ts

@@ -0,0 +1,49 @@
+import { Disposable } from './class/dispose'
+import { CliRunner } from './cli/runner'
+import { WorkspaceExperiments } from './experiments/workspace'
+import { setContextValue } from './vscode/context'
+
+export class Context extends Disposable {
+  private readonly experiments: WorkspaceExperiments
+  private readonly cliRunner: CliRunner
+
+  constructor(experiments: WorkspaceExperiments, cliRunner: CliRunner) {
+    super()
+
+    this.experiments = experiments
+    this.cliRunner = cliRunner
+
+    this.dispose.track(
+      this.cliRunner.onDidStartProcess(() => {
+        this.setIsExperimentRunning()
+      })
+    )
+
+    this.dispose.track(
+      this.cliRunner.onDidCompleteProcess(({ cwd }) =>
+        this.experiments.getRepository(cwd).update()
+      )
+    )
+
+    this.dispose.track(
+      this.experiments.onDidChangeExperiments(() => {
+        this.setIsExperimentRunning()
+      })
+    )
+  }
+
+  private setIsExperimentRunning() {
+    if (this.cliRunner.isExperimentRunning()) {
+      setContextValue('dvc.experiment.running', true)
+      return
+    }
+
+    let hasRunning = false
+    for (const dvcRoot of this.experiments.getDvcRoots()) {
+      if (this.experiments.getRepository(dvcRoot).hasRunningExperiment()) {
+        hasRunning = true
+      }
+    }
+    setContextValue('dvc.experiment.running', hasRunning)
+  }
+}

extension/src/experiments/index.ts

@@ -377,6 +377,10 @@ export class Experiments extends BaseRepository<TableData> {
     }
   }
 
+  public hasRunningExperiment() {
+    return this.experiments.hasRunningExperiment()
+  }
+
   private hideTableColumn(path: string) {
     this.toggleColumnStatus(path)
     sendTelemetryEvent(
@@ -438,6 +442,7 @@ export class Experiments extends BaseRepository<TableData> {
       columns: this.columns.getSelected(),
       hasCheckpoints: this.hasCheckpoints(),
       hasColumns: this.columns.hasColumns(),
+      hasRunningExperiment: this.experiments.hasRunningExperiment(),
       rows: this.experiments.getRowData(),
       sorts: this.experiments.getSorts()
     }

extension/src/experiments/model/accumulator.ts

@@ -5,10 +5,12 @@ export class ExperimentsAccumulator {
   public branches: Experiment[] = []
   public checkpointsByTip: Map<string, Experiment[]> = new Map()
   public experimentsByBranch: Map<string, Experiment[]> = new Map()
+  public hasRunning: boolean
 
   constructor(workspace: Experiment | undefined) {
     if (workspace) {
       this.workspace = workspace
     }
+    this.hasRunning = !!workspace?.running
   }
 }

extension/src/experiments/model/collect.ts

@@ -192,6 +192,15 @@ const transformExperimentOrCheckpointData = (
   }
 }
 
+const collectHasRunningExperiment = (
+  acc: ExperimentsAccumulator,
+  experiment: Experiment
+) => {
+  if (experiment.running) {
+    acc.hasRunning = true
+  }
+}
+
 const collectExperimentOrCheckpoint = (
   acc: ExperimentsAccumulator,
   experiment: Experiment,
@@ -227,6 +236,7 @@ const collectFromExperimentsObject = (
     }
 
     collectExperimentOrCheckpoint(acc, experiment, branchName, checkpointTipId)
+    collectHasRunningExperiment(acc, experiment)
   }
 }
 
@@ -248,6 +258,7 @@ const collectFromBranchesObject = (
 
     if (branch) {
       collectFromExperimentsObject(acc, experimentsObject, sha, branch.label)
+      collectHasRunningExperiment(acc, branch)
 
       acc.branches.push(branch)
     }

extension/src/experiments/model/index.ts

@@ -56,6 +56,7 @@ export class ExperimentsModel extends ModelWithPersistence {
   private useFiltersForSelection = false
 
   private currentSorts: SortDefinition[]
+  private hasRunning = false
 
   constructor(dvcRoot: string, workspaceState: Memento) {
     super(dvcRoot, workspaceState)
@@ -85,13 +86,19 @@ export class ExperimentsModel extends ModelWithPersistence {
   }
 
   public transformAndSet(data: ExperimentsOutput) {
-    const { workspace, branches, experimentsByBranch, checkpointsByTip } =
-      collectExperiments(data)
+    const {
+      workspace,
+      branches,
+      experimentsByBranch,
+      checkpointsByTip,
+      hasRunning
+    } = collectExperiments(data)
 
     this.workspace = workspace
     this.branches = branches
     this.experimentsByBranch = experimentsByBranch
     this.checkpointsByTip = checkpointsByTip
+    this.hasRunning = hasRunning
 
     this.setColoredStatus()
   }
@@ -117,6 +124,10 @@ export class ExperimentsModel extends ModelWithPersistence {
     return this.coloredStatus[id]
   }
 
+  public hasRunningExperiment() {
+    return this.hasRunning
+  }
+
   public canSelect() {
     return canSelect(this.coloredStatus)
   }

extension/src/experiments/webview/contract.ts

@@ -57,6 +57,7 @@ export type TableData = {
   columnWidths: Record<string, number>
   hasCheckpoints: boolean
   hasColumns: boolean
+  hasRunningExperiment: boolean
   rows: Row[]
   sorts: SortDefinition[]
 }

extension/src/experiments/workspace.ts

@@ -19,6 +19,8 @@ export class WorkspaceExperiments extends BaseWorkspaceWebviews<
     new EventEmitter<void>()
   )
 
+  public readonly onDidChangeExperiments = this.experimentsChanged.event
+
   public readonly columnsChanged = this.dispose.track(new EventEmitter<void>())
 
   public readonly updatesPaused: EventEmitter<boolean>

extension/src/extension.ts

@@ -1,8 +1,9 @@
 import { commands, env, Event, EventEmitter, ExtensionContext } from 'vscode'
-import { Config } from './config'
 import { CliExecutor } from './cli/executor'
 import { CliRunner } from './cli/runner'
 import { CliReader } from './cli/reader'
+import { Config } from './config'
+import { Context } from './context'
 import { isVersionCompatible } from './cli/version'
 import { isPythonExtensionInstalled } from './extensions/python'
 import { WorkspaceExperiments } from './experiments/workspace'
@@ -133,11 +134,7 @@ export class Extension extends Disposable implements IExtension {
       new WorkspaceRepositories(this.internalCommands)
     )
 
-    this.dispose.track(
-      this.cliRunner.onDidCompleteProcess(({ cwd }) =>
-        this.experiments.getRepository(cwd).update()
-      )
-    )
+    this.dispose.track(new Context(this.experiments, this.cliRunner))
 
     this.dispose.track(
       new ExperimentsColumnsTree(

extension/src/test/fixtures/expShow/deeplyNested.ts

@@ -282,6 +282,7 @@ const deeplyNestedTableData: TableData = {
   columnOrder: [],
   columnWidths: {},
   hasCheckpoints: false,
+  hasRunningExperiment: false,
   sorts: [
     {
       path: 'params:params.yaml:nested1.doubled',

extension/src/test/fixtures/expShow/tableData.ts

@@ -6,6 +6,7 @@ const tableDataFixture: TableData = {
   rows: rowsFixture,
   columns: columnsFixture,
   hasCheckpoints: true,
+  hasRunningExperiment: true,
   hasColumns: true,
   sorts: [],
   changes: [],

extension/src/test/suite/cli/runner.test.ts

@@ -1,7 +1,7 @@
 import { afterEach, beforeEach, describe, it, suite } from 'mocha'
 import { expect } from 'chai'
 import { restore, spy, stub } from 'sinon'
-import { window, commands, Event, EventEmitter } from 'vscode'
+import { window, Event, EventEmitter } from 'vscode'
 import { Disposable, Disposer } from '../../../extension'
 import { Config } from '../../../config'
 import { CliRunner } from '../../../cli/runner'
@@ -38,8 +38,6 @@ suite('CLI Runner Test Suite', () => {
       const cliRunner = disposable.track(new CliRunner({} as Config, 'sleep'))
       const cwd = __dirname
 
-      const executeCommandSpy = spy(commands, 'executeCommand')
-
       const onDidCompleteProcess = (): Promise<void> =>
         new Promise(resolve =>
           disposable.track(cliRunner.onDidCompleteProcess(() => resolve()))
@@ -50,11 +48,6 @@ suite('CLI Runner Test Suite', () => {
       const completed = onDidCompleteProcess()
 
       expect(cliRunner.isExperimentRunning()).to.be.true
-      expect(executeCommandSpy).to.be.calledWith(
-        'setContext',
-        'dvc.runner.running',
-        true
-      )
 
       // eslint-disable-next-line @typescript-eslint/no-explicit-any
       const closeSpy = spy((cliRunner as any).pseudoTerminal, 'close')
@@ -65,11 +58,6 @@ suite('CLI Runner Test Suite', () => {
       await completed
 
       expect(cliRunner.isExperimentRunning()).to.be.false
-      expect(executeCommandSpy).to.be.calledWith(
-        'setContext',
-        'dvc.runner.running',
-        false
-      )
     }).timeout(WEBVIEW_TEST_TIMEOUT)
 
     it('should be able to execute a command and provide the correct events in the correct order', async () => {