Enable saving of Studio access token

.vscode/settings.json

@@ -24,6 +24,7 @@
     "featurize",
     "hardlinks",
     "Interactors",
+    "isat",
     "isdir",
     "isempty",
     "isequal",

README.md

@@ -148,6 +148,7 @@ These are the VS Code [settings] available for the Extension:
 
 | **Option**                             | **Description**                                                                                                                                          |
 | -------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| `dvc.studioAccessToken`                | Token used by DVC to share experiments and plots data to [Iterative Studio](https://studio.iterative.ai).                                                |
 | `dvc.dvcPath`                          | Path or shell command to the DVC binary. Required unless Microsoft's [Python extension] is installed and the `dvc` package found in its environment.     |
 | `dvc.pythonPath`                       | Path to the desired Python interpreter to use with DVC. Should only be utilized when using a virtual environment without Microsoft's [Python extension]. |
 | `dvc.experimentsTableHeadMaxHeight`    | Maximum height of experiment table head rows.                                                                                                            |

extension/package.json

@@ -88,6 +88,11 @@
       }
     ],
     "commands": [
+      {
+        "title": "Connect to Studio",
+        "command": "dvc.studioConnect",
+        "category": "DVC"
+      },
       {
         "title": "%command.addExperimentsTableFilter%",
         "command": "dvc.addExperimentsTableFilter",
@@ -590,6 +595,15 @@
           "description": "%config.pythonPath.description%",
           "type": "string",
           "default": null
+        },
+        "dvc.studioAccessToken": {
+          "title": "%config.studioAccessToken.title%",
+          "description": "%config.studioAccessToken.description%",
+          "type": "string",
+          "pattern": "^isat_",
+          "minLength": 54,
+          "maxLength": 54,
+          "default": null
         }
       }
     },
@@ -786,6 +800,10 @@
           "command": "dvc.stopQueuedExperiments",
           "when": "dvc.commands.available && dvc.project.available"
         },
+        {
+          "command": "dvc.studioConnect",
+          "when": "dvc.commands.available && dvc.project.available && !dvc.studio.connected"
+        },
         {
           "command": "dvc.selectForCompare",
           "when": "false"
@@ -1319,6 +1337,11 @@
           "name": "Views",
           "when": "true"
         },
+        {
+          "id": "dvc.views.studio",
+          "name": "Studio",
+          "when": "!dvc.studio.connected"
+        },
         {
           "id": "dvc.views.experimentsColumnsTree",
           "name": "Columns",
@@ -1356,6 +1379,11 @@
         "contents": "[Show Experiments](command:dvc.showExperiments)\n[Show Plots](command:dvc.showPlots)\n[Show Experiments and Plots](command:dvc.showExperimentsAndPlots)",
         "when": "dvc.commands.available && dvc.project.available && dvc.project.hasData"
       },
+      {
+        "view": "dvc.views.studio",
+        "contents": "[$(plug) Connect](command:dvc.studioConnect)",
+        "when": "!dvc.studio.connected"
+      },
       {
         "view": "dvc.views.webviews",
         "contents": "[Show Experiments](command:dvc.showSetup)\n[Show Plots](command:dvc.showSetup)\n[Show Experiments and Plots](command:dvc.showSetup)",

extension/package.nls.json

@@ -92,5 +92,7 @@
   "config.experimentsTableHeadMaxHeight.title": "Maximum height of Experiment table head rows",
   "config.experimentsTableHeadMaxHeight.description": "Use 0 for infinite height.",
   "config.pythonPath.description": "Path to the desired Python interpreter to use with DVC. Required when using a virtual environment. Overrides any other extension's settings for this extension's purposes.",
-  "config.pythonPath.title": "Python Interpreter"
+  "config.pythonPath.title": "Python Interpreter",
+  "config.studioAccessToken.description": "Token used by DVC to share experiments and plots data to Iterative Studio.",
+  "config.studioAccessToken.title": "Studio Access Token"
 }

extension/src/connect/index.ts

@@ -0,0 +1,88 @@
+import { commands, workspace } from 'vscode'
+import { isStudioAccessToken, validateTokenInput } from './input'
+import { STUDIO_URL } from './webview/contract'
+import { Resource } from '../resourceLocator'
+import { ViewKey } from '../webview/constants'
+import { MessageFromWebview, MessageFromWebviewType } from '../webview/contract'
+import { BaseRepository } from '../webview/repository'
+import { Logger } from '../common/logger'
+import { getInput, getValidInput } from '../vscode/inputBox'
+import { Title } from '../vscode/title'
+import { ConfigKey, getConfigValue, setUserConfigValue } from '../vscode/config'
+import { openUrl } from '../vscode/external'
+import { ContextKey, setContextValue } from '../vscode/context'
+
+export class Connect extends BaseRepository<undefined> {
+  public readonly viewKey = ViewKey.CONNECT
+
+  constructor(webviewIcon: Resource) {
+    super('', webviewIcon)
+
+    this.dispose.track(
+      this.onDidReceivedWebviewMessage(message =>
+        this.handleMessageFromWebview(message)
+      )
+    )
+
+    void this.setContext()
+
+    this.dispose.track(
+      workspace.onDidChangeConfiguration(e => {
+        if (!e.affectsConfiguration(ConfigKey.STUDIO_ACCESS_TOKEN)) {
+          return
+        }
+        return this.setContext()
+      })
+    )
+  }
+
+  public sendInitialWebviewData(): void {}
+
+  private handleMessageFromWebview(message: MessageFromWebview) {
+    switch (message.type) {
+      case MessageFromWebviewType.OPEN_STUDIO:
+        return openUrl(STUDIO_URL)
+      case MessageFromWebviewType.OPEN_STUDIO_PROFILE:
+        return this.openStudioProfile()
+      case MessageFromWebviewType.SAVE_STUDIO_TOKEN:
+        return this.saveStudioToken()
+      default:
+        Logger.error(`Unexpected message: ${JSON.stringify(message)}`)
+    }
+  }
+
+  private async openStudioProfile() {
+    const username = await getInput(Title.ENTER_STUDIO_USERNAME)
+    if (!username) {
+      return
+    }
+    return openUrl(`${STUDIO_URL}/user/${username}/profile`)
+  }
+
+  private async saveStudioToken() {
+    const token = await getValidInput(
+      Title.ENTER_STUDIO_TOKEN,
+      validateTokenInput
+    )
+    if (!token) {
+      return
+    }
+
+    await setUserConfigValue(ConfigKey.STUDIO_ACCESS_TOKEN, token)
+    return commands.executeCommand(
+      'workbench.action.openSettings',
+      'dvc.studioAccessToken'
+    )
+  }
+
+  private async setContext() {
+    if (
+      isStudioAccessToken(await getConfigValue(ConfigKey.STUDIO_ACCESS_TOKEN))
+    ) {
+      this.webview?.dispose()
+      return setContextValue(ContextKey.STUDIO_CONNECTED, true)
+    }
+
+    return setContextValue(ContextKey.STUDIO_CONNECTED, false)
+  }
+}

extension/src/connect/input.test.ts

@@ -0,0 +1,20 @@
+import { validateTokenInput } from './input'
+
+describe('validateTokenInput', () => {
+  const mockedStudioAccessToken =
+    'isat_1Z4T0zVHvq9Cu03XEe9Zjvx2vkBihfGPdY7FfmEMAagOXfQxU'
+  it('should return the warning if the input is not valid', () => {
+    expect(
+      validateTokenInput(mockedStudioAccessToken.slice(0, -1))
+    ).not.toBeNull()
+    expect(
+      validateTokenInput(
+        mockedStudioAccessToken.slice(1, mockedStudioAccessToken.length)
+      )
+    ).not.toBeNull()
+  })
+
+  it('should return null if the input is valid', () => {
+    expect(validateTokenInput(mockedStudioAccessToken)).toBeNull()
+  })
+})

extension/src/connect/input.ts

@@ -0,0 +1,13 @@
+export const isStudioAccessToken = (text?: string): boolean => {
+  if (!text) {
+    return false
+  }
+  return text.startsWith('isat_') && text.length === 54
+}
+
+export const validateTokenInput = (input: string | undefined) => {
+  if (!isStudioAccessToken(input)) {
+    return 'please enter a valid Studio access token'
+  }
+  return null
+}

extension/src/connect/webview/contract.ts

@@ -0,0 +1 @@
+export const STUDIO_URL = 'https://studio.iterative.ai'

extension/src/extension.ts

@@ -52,12 +52,14 @@ import { stopProcesses } from './processExecution'
 import { Flag } from './cli/dvc/constants'
 import { LanguageClient } from './languageClient'
 import { collectRunningExperimentPids } from './experiments/processExecution/collect'
+import { Connect } from './connect'
 
 export class Extension extends Disposable {
   protected readonly internalCommands: InternalCommands
 
   private readonly resourceLocator: ResourceLocator
   private readonly repositories: WorkspaceRepositories
+  private readonly connect: Connect
   private readonly experiments: WorkspaceExperiments
   private readonly plots: WorkspacePlots
   private readonly setup: Setup
@@ -77,6 +79,12 @@ export class Extension extends Disposable {
 
     const stopWatch = new StopWatch()
 
+    this.dispose.track(
+      commands.registerCommand('dvc.studioConnect', () =>
+        this.connect.showWebview()
+      )
+    )
+
     this.dispose.track(getTelemetryReporter())
 
     this.resourceLocator = this.dispose.track(
@@ -185,6 +193,8 @@ export class Extension extends Disposable {
       )
     )
 
+    this.connect = this.dispose.track(new Connect(this.resourceLocator.dvcIcon))
+
     registerExperimentCommands(this.experiments, this.internalCommands)
     registerPlotsCommands(this.plots, this.internalCommands)
     this.internalCommands.registerExternalCommand(

extension/src/telemetry/constants.ts

@@ -28,6 +28,10 @@ export const EventName = Object.assign(
     EXTENSION_EXECUTION_DETAILS_CHANGED: 'extension.executionDetails.changed',
     EXTENSION_LOAD: 'extension.load',
 
+    VIEWS_CONNECT_CLOSED: 'views.connect.closed',
+    VIEWS_CONNECT_CREATED: 'views.connect.created',
+    VIEWS_CONNECT_FOCUS_CHANGED: 'views.connect.focusChanged',
+
     VIEWS_EXPERIMENTS_TABLE_CLOSED: 'views.experimentsTable.closed',
     VIEWS_EXPERIMENTS_TABLE_COLUMNS_REORDERED:
       'views.experimentsTable.columnsReordered',
@@ -209,6 +213,10 @@ export interface IEventNamePropertyMapping {
   [EventName.EXTENSION_SHOW_COMMANDS]: undefined
   [EventName.EXTENSION_SHOW_OUTPUT]: undefined
 
+  [EventName.VIEWS_CONNECT_CLOSED]: undefined
+  [EventName.VIEWS_CONNECT_CREATED]: undefined
+  [EventName.VIEWS_CONNECT_FOCUS_CHANGED]: undefined
+
   [EventName.VIEWS_EXPERIMENTS_TREE_OPENED]: DvcRootCount
   [EventName.VIEWS_EXPERIMENTS_FILTER_BY_TREE_OPENED]: DvcRootCount
   [EventName.VIEWS_EXPERIMENTS_METRICS_AND_PARAMS_TREE_OPENED]: DvcRootCount

extension/src/vscode/config.ts

@@ -5,10 +5,11 @@ export enum ConfigKey {
   DO_NOT_SHOW_CLI_UNAVAILABLE = 'dvc.doNotShowCliUnavailable',
   DO_NOT_SHOW_WALKTHROUGH_AFTER_INSTALL = 'dvc.doNotShowWalkthroughAfterInstall',
   DO_NOT_SHOW_UNABLE_TO_FILTER = 'dvc.doNotShowUnableToFilter',
+  DVC_PATH = 'dvc.dvcPath',
   EXP_TABLE_HEAD_MAX_HEIGHT = 'dvc.experimentsTableHeadMaxHeight',
   FOCUSED_PROJECTS = 'dvc.focusedProjects',
-  DVC_PATH = 'dvc.dvcPath',
-  PYTHON_PATH = 'dvc.pythonPath'
+  PYTHON_PATH = 'dvc.pythonPath',
+  STUDIO_ACCESS_TOKEN = 'dvc.studioAccessToken'
 }
 
 export const getConfigValue = <T = string, D = string>(

extension/src/vscode/context.ts

@@ -1,8 +1,9 @@
 import { commands } from 'vscode'
 
 export enum ContextKey {
-  COMMANDS_AVAILABLE = 'dvc.commands.available',
   CLI_INCOMPATIBLE = 'dvc.cli.incompatible',
+  COMMANDS_AVAILABLE = 'dvc.commands.available',
+  CONNECT_WEBVIEW_ACTIVE = 'dvc.connect.webview.active',
   EXPERIMENT_CHECKPOINTS = 'dvc.experiment.checkpoints',
   EXPERIMENT_FILTERS_SELECTED = 'dvc.experiments.filter.selected',
   EXPERIMENTS_WEBVIEW_ACTIVE = 'dvc.experiments.webview.active',
@@ -15,7 +16,8 @@ export enum ContextKey {
   PROJECT_AVAILABLE = 'dvc.project.available',
   PROJECT_HAS_DATA = 'dvc.project.hasData',
   SCM_RUNNING = 'dvc.scm.command.running',
-  SETUP_WEBVIEW_ACTIVE = 'dvc.setup.webview.active'
+  SETUP_WEBVIEW_ACTIVE = 'dvc.setup.webview.active',
+  STUDIO_CONNECTED = 'dvc.studio.connected'
 }
 
 export const setContextValue = (key: ContextKey, value: unknown) =>

extension/src/vscode/external.ts

@@ -0,0 +1,3 @@
+import { env, Uri } from 'vscode'
+
+export const openUrl = (url: string) => env.openExternal(Uri.parse(url))

extension/src/vscode/title.ts

@@ -7,6 +7,8 @@ export enum Title {
   ENTER_FILTER_VALUE = 'Enter a Filter Value',
   ENTER_RELATIVE_DESTINATION = 'Enter a Destination Relative to the Root',
   ENTER_PATH_OR_CHOOSE_FILE = 'Enter the path to your training script or select it',
+  ENTER_STUDIO_USERNAME = 'Enter your Studio username',
+  ENTER_STUDIO_TOKEN = 'Enter your Studio access token',
   ENTER_STAGE_NAME = 'Enter a name for the main stage of your pipeline',
   GARBAGE_COLLECT_EXPERIMENTS = 'Garbage Collect Experiments',
   SHOW_SETUP = 'Show Setup',

extension/src/webview/constants.ts

@@ -1,8 +1,16 @@
-import { distPath, react, experiments, plots, setup } from 'dvc-vscode-webview'
+import {
+  distPath,
+  connect,
+  react,
+  experiments,
+  plots,
+  setup
+} from 'dvc-vscode-webview'
 import { EventName, IEventNamePropertyMapping } from '../telemetry/constants'
 import { ContextKey } from '../vscode/context'
 
 export enum ViewKey {
+  CONNECT = 'dvc-connect',
   EXPERIMENTS = 'dvc-experiments',
   PLOTS = 'dvc-plots',
   SETUP = 'dvc-setup'
@@ -25,6 +33,18 @@ export const WebviewDetails: {
     viewKey: ViewKey
   }
 } = {
+  [ViewKey.CONNECT]: {
+    contextKey: ContextKey.CONNECT_WEBVIEW_ACTIVE,
+    distPath,
+    eventNames: {
+      closedEvent: EventName.VIEWS_CONNECT_CLOSED,
+      createdEvent: EventName.VIEWS_CONNECT_CREATED,
+      focusChangedEvent: EventName.VIEWS_CONNECT_FOCUS_CHANGED
+    },
+    scripts: [react, connect],
+    title: 'Connect',
+    viewKey: ViewKey.CONNECT
+  },
   [ViewKey.EXPERIMENTS]: {
     contextKey: ContextKey.EXPERIMENTS_WEBVIEW_ACTIVE,
     distPath,

extension/src/webview/contract.ts

@@ -8,7 +8,7 @@ import {
 } from '../plots/webview/contract'
 import { SetupData } from '../setup/webview/contract'
 
-export type WebviewData = TableData | PlotsData | SetupData
+export type WebviewData = TableData | PlotsData | SetupData | {}
 
 export enum MessageFromWebviewType {
   INITIALIZED = 'initialized',
@@ -17,8 +17,10 @@ export enum MessageFromWebviewType {
   CREATE_BRANCH_FROM_EXPERIMENT = 'create-branch-from-experiment',
   FOCUS_FILTERS_TREE = 'focus-filters-tree',
   FOCUS_SORTS_TREE = 'focus-sorts-tree',
-  OPEN_PLOTS_WEBVIEW = 'open-plots-webview',
   OPEN_PARAMS_FILE_TO_THE_SIDE = 'open-params-file-to-the-side',
+  OPEN_PLOTS_WEBVIEW = 'open-plots-webview',
+  OPEN_STUDIO = 'open-studio',
+  OPEN_STUDIO_PROFILE = 'open-studio-in-profile',
   REMOVE_COLUMN_SORT = 'remove-column-sort',
   REMOVE_EXPERIMENT = 'remove-experiment',
   REORDER_COLUMNS = 'reorder-columns',
@@ -30,6 +32,7 @@ export enum MessageFromWebviewType {
   REFRESH_REVISIONS = 'refresh-revisions',
   RESIZE_COLUMN = 'resize-column',
   RESIZE_PLOTS = 'resize-plots',
+  SAVE_STUDIO_TOKEN = 'save-studio-token',
   STOP_EXPERIMENT = 'stop-experiment',
   SORT_COLUMN = 'sort-column',
   TOGGLE_EXPERIMENT = 'toggle-experiment',
@@ -196,6 +199,9 @@ export type MessageFromWebview =
   | { type: MessageFromWebviewType.SHOW_SCM_PANEL }
   | { type: MessageFromWebviewType.INSTALL_DVC }
   | { type: MessageFromWebviewType.SETUP_WORKSPACE }
+  | { type: MessageFromWebviewType.OPEN_STUDIO }
+  | { type: MessageFromWebviewType.OPEN_STUDIO_PROFILE }
+  | { type: MessageFromWebviewType.SAVE_STUDIO_TOKEN }
 
 export type MessageToWebview<T extends WebviewData> = {
   type: MessageToWebviewType.SET_DATA