Integrate remaining share experiment command with the extension (exp …

…push) (#3781)

* integrate remaining share function with the extension (exp push)

* self review

* rename share to push

* remove custom error handling (use offer to show error in output channel)

README.md

@@ -154,6 +154,7 @@ These are the VS Code [settings] available for the Extension:
 | `dvc.studio.shareExperimentsLive`      | Automatically share all new experiment metrics and plots logged with DVCLive to Studio. This option will only take effect once Studio is connected.      |
 | `dvc.focusedProjects`                  | A subset of paths to the workspace's available DVC projects. Using this option will override project auto-discovery.                                     |
 | `dvc.doNotShowWalkthroughAfterInstall` | Do not prompt to show the Get Started page after installing. Useful for pre-configured development environments                                          |
+| `dvc.doNotRecommendAddStudioToken`     | Do not prompt to add a [studio.token] to the global DVC config, which enables automatic sharing of experiments to [Studio].                              |
 | `dvc.doNotRecommendRedHatExtension`    | Do not prompt to install the Red Hat YAML extension, which helps with DVC YAML schema validation (`dvc.yaml` and `.dvc` files).                          |
 | `dvc.doNotShowCliUnavailable`          | Do not warn when the workspace contains a DVC project but the DVC binary is unavailable.                                                                 |
 
@@ -162,6 +163,9 @@ These are the VS Code [settings] available for the Extension:
 
 [python extension]:
   https://marketplace.visualstudio.com/items?itemName=ms-python.python
+[studio.token]:
+  https://dvc.org/doc/user-guide/project-structure/configuration#studio
+[Studio]: https://studio.iterative.ai
 [workspace level]:
   https://code.visualstudio.com/docs/getstarted/settings#_workspace-settings
 

extension/package.json

@@ -506,8 +506,8 @@
         "icon": "$(play)"
       },
       {
-        "title": "Share to Studio",
-        "command": "dvc.views.experiments.shareExperimentToStudio",
+        "title": "Push",
+        "command": "dvc.views.experiments.pushExperiment",
         "category": "DVC",
         "icon": "$(repo-push)"
       },
@@ -564,6 +564,11 @@
           "type": "boolean",
           "default": null
         },
+        "dvc.doNotRecommendAddStudioToken": {
+          "description": "Do not prompt to add a studio.token to the global DVC config, which enables automatic sharing of experiments to Studio.",
+          "type": "boolean",
+          "default": null
+        },
         "dvc.doNotShowCliUnavailable": {
           "description": "Do not warn when the workspace contains a DVC project but the DVC binary is unavailable.",
           "type": "boolean",
@@ -599,7 +604,7 @@
           "default": null
         },
         "dvc.studio.shareExperimentsLive": {
-          "description": "Automatically share all new experiment metrics and plots logged with DVCLive to Studio. This option will only take effect once Studio is connected.",
+          "description": "Automatically share all new experiment metrics and plots logged with DVCLive to Studio. This option will only take effect once Studio is connected (studio.token is set).",
           "type": "boolean",
           "default": false
         }
@@ -883,7 +888,7 @@
           "when": "false"
         },
         {
-          "command": "dvc.views.experiments.shareExperimentToStudio",
+          "command": "dvc.views.experiments.pushExperiment",
           "when": "false"
         },
         {
@@ -1147,7 +1152,7 @@
           "when": "view == dvc.views.experimentsTree && dvc.commands.available && viewItem =~ /^(experiment|queued)$/ && !dvc.experiment.running"
         },
         {
-          "command": "dvc.views.experiments.shareExperimentToStudio",
+          "command": "dvc.views.experiments.pushExperiment",
           "group": "1_share@0",
           "when": "view == dvc.views.experimentsTree && dvc.commands.available && viewItem == experiment && !dvc.experiment.running"
         },

extension/src/commands/external.ts

@@ -100,7 +100,6 @@ export enum RegisteredCommands {
   ADD_STUDIO_ACCESS_TOKEN = 'dvc.addStudioAccessToken',
   UPDATE_STUDIO_ACCESS_TOKEN = 'dvc.updateStudioAccessToken',
   REMOVE_STUDIO_ACCESS_TOKEN = 'dvc.removeStudioAccessToken',
-  EXPERIMENT_VIEW_SHARE_TO_STUDIO = 'dvc.views.experiments.shareExperimentToStudio',
   SETUP_SHOW_STUDIO_CONNECT = 'dvc.showStudioConnect',
   SETUP_SHOW_STUDIO_SETTINGS = 'dvc.showStudioSettings',
 

extension/src/experiments/commands/index.ts

@@ -3,37 +3,76 @@ import { AvailableCommands, InternalCommands } from '../../commands/internal'
 import { Toast } from '../../vscode/toast'
 import { WorkspaceExperiments } from '../workspace'
 import { Setup } from '../../setup'
+import { Response } from '../../vscode/response'
+import {
+  ConfigKey,
+  getConfigValue,
+  setUserConfigValue
+} from '../../vscode/config'
+import { STUDIO_URL } from '../../setup/webview/contract'
 import { RegisteredCommands } from '../../commands/external'
 
 export const getBranchExperimentCommand =
   (experiments: WorkspaceExperiments) =>
   (cwd: string, name: string, input: string) =>
     experiments.runCommand(AvailableCommands.EXP_BRANCH, cwd, name, input)
 
-export const getShareExperimentToStudioCommand =
+const promptToAddStudioToken = async () => {
+  const response = await Toast.askShowOrCloseOrNever(
+    `Experiments can be automatically shared to [Studio](${STUDIO_URL}) by setting the studio.token in your config.`
+  )
+
+  if (!response || response === Response.CLOSE) {
+    return
+  }
+  if (response === Response.SHOW) {
+    return commands.executeCommand(RegisteredCommands.SETUP_SHOW_STUDIO_CONNECT)
+  }
+  if (response === Response.NEVER) {
+    return setUserConfigValue(ConfigKey.DO_NOT_RECOMMEND_ADD_STUDIO_TOKEN, true)
+  }
+}
+
+const convertUrlTextToLink = (stdout: string) => {
+  const experimentAtRegex = /\sat\s+(https:\/\/studio\.iterative\.ai\/.*$)/
+  const match = stdout.match(experimentAtRegex)
+  if (!(match?.[0] && match?.[1])) {
+    return stdout
+  }
+  return stdout.replace(match[0], ` in [Studio](${match[1]})`)
+}
+
+export const getPushExperimentCommand =
   (internalCommands: InternalCommands, setup: Setup) =>
   ({ dvcRoot, id }: { dvcRoot: string; id: string }) => {
     const studioAccessToken = setup.getStudioAccessToken()
-    if (!studioAccessToken) {
-      return commands.executeCommand(RegisteredCommands.SETUP_SHOW)
+    if (
+      !(
+        getConfigValue(ConfigKey.DO_NOT_RECOMMEND_ADD_STUDIO_TOKEN) ||
+        studioAccessToken
+      )
+    ) {
+      void promptToAddStudioToken()
     }
 
-    return Toast.showProgress('Sharing', async progress => {
+    return Toast.showProgress('exp push', async progress => {
       progress.report({ increment: 0 })
 
-      progress.report({ increment: 25, message: 'Running exp push...' })
-
-      await Toast.runCommandAndIncrementProgress(
-        () =>
-          internalCommands.executeCommand(
-            AvailableCommands.EXP_PUSH,
-            dvcRoot,
-            id
-          ),
-        progress,
-        75
+      progress.report({ increment: 25, message: `Pushing ${id}...` })
+
+      const remainingProgress = 75
+
+      const stdout = await internalCommands.executeCommand(
+        AvailableCommands.EXP_PUSH,
+        dvcRoot,
+        id
       )
 
+      progress.report({
+        increment: remainingProgress,
+        message: convertUrlTextToLink(stdout)
+      })
+
       return Toast.delayProgressClosing(15000)
     })
   }

extension/src/experiments/commands/register.ts

@@ -1,7 +1,4 @@
-import {
-  getBranchExperimentCommand,
-  getShareExperimentToStudioCommand
-} from '.'
+import { getBranchExperimentCommand, getPushExperimentCommand } from '.'
 import { pickGarbageCollectionFlags } from '../quickPick'
 import { WorkspaceExperiments } from '../workspace'
 import { AvailableCommands, InternalCommands } from '../../commands/internal'
@@ -292,9 +289,9 @@ export const registerExperimentCommands = (
       experiments.getRepository(dvcRoot).toggleExperimentStatus(id)
   )
 
-  internalCommands.registerExternalCommand(
-    RegisteredCommands.EXPERIMENT_VIEW_SHARE_TO_STUDIO,
-    getShareExperimentToStudioCommand(internalCommands, setup)
+  internalCommands.registerExternalCliCommand(
+    RegisteredCliCommands.EXPERIMENT_VIEW_PUSH,
+    getPushExperimentCommand(internalCommands, setup)
   )
 
   internalCommands.registerExternalCliCommand(

extension/src/experiments/webview/messages.ts

@@ -194,9 +194,9 @@ export class WebviewMessages {
       case MessageFromWebviewType.ADD_CONFIGURATION: {
         return this.addConfiguration()
       }
-      case MessageFromWebviewType.SHARE_EXPERIMENT_TO_STUDIO:
+      case MessageFromWebviewType.PUSH_EXPERIMENT:
         return commands.executeCommand(
-          RegisteredCommands.EXPERIMENT_VIEW_SHARE_TO_STUDIO,
+          RegisteredCliCommands.EXPERIMENT_VIEW_PUSH,
           { dvcRoot: this.dvcRoot, id: message.payload }
         )
 

extension/src/telemetry/constants.ts

@@ -144,7 +144,6 @@ export interface IEventNamePropertyMapping {
   [EventName.EXPERIMENT_VIEW_BRANCH]: undefined
   [EventName.EXPERIMENT_VIEW_PUSH]: undefined
   [EventName.EXPERIMENT_VIEW_REMOVE]: undefined
-  [EventName.EXPERIMENT_VIEW_SHARE_TO_STUDIO]: undefined
   [EventName.EXPERIMENT_VIEW_SHOW_LOGS]: undefined
   [EventName.EXPERIMENT_VIEW_STOP]: undefined
   [EventName.QUEUE_EXPERIMENT]: undefined

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

@@ -10,7 +10,8 @@ import {
   workspace,
   Uri,
   QuickPickItem,
-  ViewColumn
+  ViewColumn,
+  CancellationToken
 } from 'vscode'
 import { buildExperiments, stubWorkspaceExperimentsGetters } from './util'
 import { Disposable } from '../../../extension'
@@ -34,6 +35,7 @@ import {
 import {
   buildInternalCommands,
   buildMockExperimentsData,
+  bypassProgressCloseDelay,
   closeAllEditors,
   configurationChangeEvent,
   experimentsUpdatedEvent,
@@ -84,6 +86,8 @@ import { DvcReader } from '../../../cli/dvc/reader'
 import { DvcViewer } from '../../../cli/dvc/viewer'
 import { DEFAULT_NB_ITEMS_PER_ROW } from '../../../plots/webview/contract'
 import { GitReader } from '../../../cli/git/reader'
+import { Toast } from '../../../vscode/toast'
+import { Response } from '../../../vscode/response'
 
 const { openFileInEditor } = FileSystem
 
@@ -633,7 +637,7 @@ suite('Experiments Test Suite', () => {
       )
     }).timeout(WEBVIEW_TEST_TIMEOUT)
 
-    it('should handle a message to share an experiment to Studio', async () => {
+    it('should handle a message to push an experiment', async () => {
       const { experiments } = buildExperiments(disposable)
       await experiments.isReady()
 
@@ -643,6 +647,7 @@ suite('Experiments Test Suite', () => {
       const mockMessageReceived = getMessageReceivedEmitter(webview)
 
       const executeCommandSpy = spy(commands, 'executeCommand')
+      const mockExpPush = stub(DvcExecutor.prototype, 'expPush')
 
       const mockGetStudioAccessToken = stub(
         Setup.prototype,
@@ -656,15 +661,24 @@ suite('Experiments Test Suite', () => {
         })
       )
 
+      const mockAskShowOrCloseOrNever = stub(Toast, 'askShowOrCloseOrNever')
+
+      const userPrompted = new Promise(resolve =>
+        mockAskShowOrCloseOrNever.callsFake(() => {
+          resolve(undefined)
+          return Promise.resolve(Response.SHOW)
+        })
+      )
+
       mockMessageReceived.fire({
         payload: mockExpId,
-        type: MessageFromWebviewType.SHARE_EXPERIMENT_TO_STUDIO
+        type: MessageFromWebviewType.PUSH_EXPERIMENT
       })
 
-      await tokenNotFound
+      await Promise.all([tokenNotFound, userPrompted])
 
       expect(executeCommandSpy).to.be.calledWithExactly(
-        RegisteredCommands.SETUP_SHOW
+        RegisteredCommands.SETUP_SHOW_STUDIO_CONNECT
       )
 
       mockGetStudioAccessToken.resetBehavior()
@@ -675,24 +689,41 @@ suite('Experiments Test Suite', () => {
           return 'isat_token'
         })
       )
-      const mockexpPush = stub(DvcExecutor.prototype, 'expPush')
+
+      const mockShowProgress = stub(Toast, 'showProgress')
+      bypassProgressCloseDelay()
+
+      const mockReport = stub()
+
+      mockShowProgress.callsFake((title, callback) => {
+        expect(title).to.equal('exp push')
+
+        const progress = { report: mockReport }
+        return callback(progress, {} as CancellationToken)
+      })
+
       const commandExecuted = new Promise(resolve =>
-        mockexpPush.callsFake(() => {
+        mockExpPush.callsFake(() => {
           resolve(undefined)
           return Promise.resolve(
-            `Pushed experiment ${mockExpId} to Git remote 'origin'`
+            "Experiment major-lamb is up to date on Git remote 'origin'.\nView your experiments at \nhttps://studio.iterative.ai/user/mattseddon/projects/vscode-dvc-demo-ynm6t3jxdx"
           )
         })
       )
 
       mockMessageReceived.fire({
         payload: mockExpId,
-        type: MessageFromWebviewType.SHARE_EXPERIMENT_TO_STUDIO
+        type: MessageFromWebviewType.PUSH_EXPERIMENT
       })
 
       await Promise.all([tokenFound, commandExecuted])
 
-      expect(mockexpPush).to.be.calledWithExactly(dvcDemoPath, mockExpId)
+      expect(mockExpPush).to.be.calledWithExactly(dvcDemoPath, mockExpId)
+      expect(mockReport).to.be.calledWithExactly({
+        increment: 75,
+        message:
+          "Experiment major-lamb is up to date on Git remote 'origin'.\nView your experiments in [Studio](https://studio.iterative.ai/user/mattseddon/projects/vscode-dvc-demo-ynm6t3jxdx)"
+      })
     }).timeout(WEBVIEW_TEST_TIMEOUT)
 
     it("should be able to handle a message to modify an experiment's params and queue an experiment", async () => {

extension/src/vscode/config.ts

@@ -1,12 +1,13 @@
 import { ConfigurationTarget, workspace } from 'vscode'
 
 export enum ConfigKey {
+  DO_NOT_RECOMMEND_ADD_STUDIO_TOKEN = 'dvc.doNotRecommendAddStudioToken',
   DO_NOT_RECOMMEND_RED_HAT = 'dvc.doNotRecommendRedHatExtension',
   DO_NOT_SHOW_CLI_UNAVAILABLE = 'dvc.doNotShowCliUnavailable',
   DO_NOT_SHOW_WALKTHROUGH_AFTER_INSTALL = 'dvc.doNotShowWalkthroughAfterInstall',
+  DVC_PATH = 'dvc.dvcPath',
   EXP_TABLE_HEAD_MAX_HEIGHT = 'dvc.experimentsTableHeadMaxHeight',
   FOCUSED_PROJECTS = 'dvc.focusedProjects',
-  DVC_PATH = 'dvc.dvcPath',
   PYTHON_PATH = 'dvc.pythonPath',
   STUDIO_SHARE_EXPERIMENTS_LIVE = 'dvc.studio.shareExperimentsLive'
 }

extension/src/webview/contract.ts

@@ -25,6 +25,7 @@ export enum MessageFromWebviewType {
   OPEN_PLOTS_WEBVIEW = 'open-plots-webview',
   OPEN_STUDIO = 'open-studio',
   OPEN_STUDIO_PROFILE = 'open-studio-profile',
+  PUSH_EXPERIMENT = 'push-experiment',
   REMOVE_COLUMN_SORT = 'remove-column-sort',
   REMOVE_EXPERIMENT = 'remove-experiment',
   REORDER_COLUMNS = 'reorder-columns',
@@ -36,7 +37,6 @@ export enum MessageFromWebviewType {
   RESIZE_COLUMN = 'resize-column',
   RESIZE_PLOTS = 'resize-plots',
   SAVE_STUDIO_TOKEN = 'save-studio-token',
-  SHARE_EXPERIMENT_TO_STUDIO = 'share-experiment-to-studio',
   SHOW_EXPERIMENT_LOGS = 'show-experiment-logs',
   STOP_EXPERIMENT = 'stop-experiment',
   SORT_COLUMN = 'sort-column',
@@ -154,7 +154,7 @@ export type MessageFromWebview =
     }
   | { type: MessageFromWebviewType.SHOW_EXPERIMENT_LOGS; payload: string }
   | {
-      type: MessageFromWebviewType.SHARE_EXPERIMENT_TO_STUDIO
+      type: MessageFromWebviewType.PUSH_EXPERIMENT
       payload: string
     }
   | {