wip demo

extension/src/setup/index.ts

@@ -4,6 +4,7 @@ import { Disposable, Disposer } from '@hediet/std/disposable'
 import isEmpty from 'lodash.isempty'
 import {
   DvcCliDetails,
+  STUDIO_URL,
   SetupSection,
   SetupData as TSetupData
 } from './webview/contract'
@@ -16,7 +17,7 @@ import { WebviewMessages } from './webview/messages'
 import { validateTokenInput } from './inputBox'
 import { findPythonBinForInstall } from './autoInstall'
 import { run, runWithRecheck, runWorkspace } from './runner'
-import { isStudioAccessToken } from './token'
+import { isStudioAccessToken, pollForStudioToken } from './token'
 import {
   PYTHON_EXTENSION_ACTION,
   pickFocusedProjects,
@@ -69,6 +70,7 @@ import {
   isActivePythonEnvGlobal,
   selectPythonInterpreter
 } from '../extensions/python'
+import { openUrl } from '../vscode/external'
 
 export class Setup
   extends BaseRepository<TSetupData>
@@ -118,7 +120,9 @@ export class Setup
   private dotFolderWatcher?: Disposer
 
   private studioAccessToken: string | undefined = undefined
+  private studioAuthLink: string | undefined = undefined
   private studioIsConnected = false
+  private studioUserCode: string | null = null
   private shareLiveToStudio: boolean | undefined = undefined
 
   private focusedSection: SetupSection | undefined = undefined
@@ -357,8 +361,7 @@ export class Setup
       return
     }
 
-    await this.accessConfig(cwd, Flag.GLOBAL, ConfigKey.STUDIO_TOKEN, token)
-    return this.updateStudioAndSend()
+    return this.saveStudioAccessTokenInConfig(cwd, token)
   }
 
   public getStudioAccessToken() {
@@ -369,6 +372,11 @@ export class Setup
     return this.sendDataToWebview()
   }
 
+  private async saveStudioAccessTokenInConfig(cwd: string, token: string) {
+    await this.accessConfig(cwd, Flag.GLOBAL, ConfigKey.STUDIO_TOKEN, token)
+    return this.updateStudioAndSend()
+  }
+
   private async getDvcCliDetails(): Promise<DvcCliDetails> {
     await this.config.isReady()
     const dvcPath = this.config.getCliPath()
@@ -445,7 +453,8 @@ export class Setup
       pythonBinPath: getBinDisplayText(pythonBinPath),
       remoteList,
       sectionCollapsed: collectSectionCollapsed(this.focusedSection),
-      shareLiveToStudio: !!this.shareLiveToStudio
+      shareLiveToStudio: !!this.shareLiveToStudio,
+      studioUserCode: this.getStudioUserCode()
     })
     this.focusedSection = undefined
   }
@@ -456,8 +465,11 @@ export class Setup
       () => this.initializeGit(),
       (offline: boolean) => this.updateStudioOffline(offline),
       () => this.isPythonExtensionUsed(),
-      () => this.updatePythonEnvironment()
+      () => this.updatePythonEnvironment(),
+      () => this.requestStudioAuth(),
+      () => this.openStudioAuthLink()
     )
+    // add both new actions to above
     this.dispose.track(
       this.onDidReceivedWebviewMessage(message =>
         webviewMessages.handleMessageFromWebview(message)
@@ -805,6 +817,68 @@ export class Setup
     )
   }
 
+  private getStudioUserCode() {
+    return this.studioUserCode
+  }
+
+  private async requestStudioAuth() {
+    const response = await fetch(`${STUDIO_URL}/api/device-login`, {
+      body: JSON.stringify({
+        client_name: 'vscode'
+      }),
+      headers: {
+        'Content-Type': 'application/json'
+      },
+      method: 'POST'
+    })
+
+    const {
+      token_uri: tokenUri,
+      verification_uri: verificationUri,
+      user_code: userCode,
+      device_code: deviceCode
+    } = (await response.json()) as {
+      token_uri: string
+      verification_uri: string
+      user_code: string
+      device_code: string
+    }
+    this.studioAuthLink = verificationUri
+    this.studioUserCode = userCode
+
+    await this.sendDataToWebview()
+    void this.requestStudioToken(deviceCode, tokenUri)
+  }
+
+  private async requestStudioToken(
+    studioDeviceCode: string,
+    studioTokenRequestUri: string
+  ) {
+    const token = await pollForStudioToken(
+      studioTokenRequestUri,
+      studioDeviceCode
+    )
+
+    this.studioAccessToken = token
+    this.studioUserCode = null
+    this.studioAuthLink = undefined
+
+    const cwd = this.getCwd()
+
+    if (!cwd) {
+      return
+    }
+
+    return this.saveStudioAccessTokenInConfig(cwd, token)
+  }
+
+  private openStudioAuthLink() {
+    if (!this.studioAuthLink) {
+      return
+    }
+    void openUrl(this.studioAuthLink)
+  }
+
   private getCwd() {
     if (!this.getCliCompatible()) {
       return

extension/src/setup/token.ts

@@ -4,3 +4,35 @@ export const isStudioAccessToken = (text?: string): boolean => {
   }
   return text.startsWith('isat_') && text.length >= 53
 }
+
+// chose the simplest way to do this, in reality we need a way to
+// offer a straightforward way to stop the polling either because the
+// user cancels or possibly because were getting errors
+export const pollForStudioToken = async (
+  tokenUri: string,
+  deviceCode: string
+): Promise<string> => {
+  const response = await fetch(tokenUri, {
+    body: JSON.stringify({
+      code: deviceCode
+    }),
+    headers: {
+      'Content-Type': 'application/json'
+    },
+    method: 'POST'
+  })
+
+  if (response.status === 400) {
+    await new Promise(resolve => setTimeout(resolve, 2000))
+    return pollForStudioToken(tokenUri, deviceCode)
+  }
+  if (response.status !== 200) {
+    await new Promise(resolve => setTimeout(resolve, 2000))
+    return pollForStudioToken(tokenUri, deviceCode)
+  }
+
+  const { access_token: accessToken } = (await response.json()) as {
+    access_token: string
+  }
+  return accessToken
+}

extension/src/setup/webview/contract.ts

@@ -23,6 +23,7 @@ export type SetupData = {
   remoteList: RemoteList
   sectionCollapsed: typeof DEFAULT_SECTION_COLLAPSED | undefined
   shareLiveToStudio: boolean
+  studioUserCode: string | null
   isAboveLatestTestedVersion: boolean | undefined
 }
 

extension/src/setup/webview/messages.ts

@@ -21,19 +21,25 @@ export class WebviewMessages {
   private readonly updateStudioOffline: (offline: boolean) => Promise<void>
   private readonly isPythonExtensionUsed: () => Promise<boolean>
   private readonly updatePythonEnv: () => Promise<void>
+  private readonly requestStudioAuth: () => Promise<void>
+  private readonly openStudioAuthLink: () => void
 
   constructor(
     getWebview: () => BaseWebview<TSetupData> | undefined,
     initializeGit: () => void,
     updateStudioOffline: (shareLive: boolean) => Promise<void>,
     isPythonExtensionUsed: () => Promise<boolean>,
-    updatePythonEnv: () => Promise<void>
+    updatePythonEnv: () => Promise<void>,
+    requestStudioAuth: () => Promise<void>,
+    openStudioAuthLink: () => void
   ) {
     this.getWebview = getWebview
     this.initializeGit = initializeGit
     this.updateStudioOffline = updateStudioOffline
     this.isPythonExtensionUsed = isPythonExtensionUsed
     this.updatePythonEnv = updatePythonEnv
+    this.requestStudioAuth = requestStudioAuth
+    this.openStudioAuthLink = openStudioAuthLink
   }
 
   public sendWebviewMessage(data: SetupData) {
@@ -68,6 +74,8 @@ export class WebviewMessages {
         return this.openStudio()
       case MessageFromWebviewType.OPEN_STUDIO_PROFILE:
         return this.openStudioProfile()
+      case MessageFromWebviewType.OPEN_STUDIO_AUTH_LINK:
+        return this.openStudioAuthLink()
       case MessageFromWebviewType.SAVE_STUDIO_TOKEN:
         return commands.executeCommand(
           RegisteredCommands.ADD_STUDIO_ACCESS_TOKEN
@@ -78,6 +86,8 @@ export class WebviewMessages {
         )
       case MessageFromWebviewType.SET_STUDIO_SHARE_EXPERIMENTS_LIVE:
         return this.updateStudioOffline(message.payload)
+      case MessageFromWebviewType.REQUEST_STUDIO_AUTH:
+        return this.requestStudioAuth()
       case MessageFromWebviewType.OPEN_EXPERIMENTS_WEBVIEW:
         return commands.executeCommand(RegisteredCommands.EXPERIMENT_SHOW)
       case MessageFromWebviewType.REMOTE_ADD:

extension/src/webview/contract.ts

@@ -32,6 +32,7 @@ export enum MessageFromWebviewType {
   OPEN_PLOTS_WEBVIEW = 'open-plots-webview',
   OPEN_STUDIO = 'open-studio',
   OPEN_STUDIO_PROFILE = 'open-studio-profile',
+  OPEN_STUDIO_AUTH_LINK = 'open-studio-auth-link',
   PUSH_EXPERIMENT = 'push-experiment',
   REMOVE_COLUMN_FILTERS = 'remove-column-filters',
   REMOVE_COLUMN_SORT = 'remove-column-sort',
@@ -47,6 +48,7 @@ export enum MessageFromWebviewType {
   RESET_COMMITS = 'reset-commits',
   RESIZE_COLUMN = 'resize-column',
   RESIZE_PLOTS = 'resize-plots',
+  REQUEST_STUDIO_AUTH = 'request-studio-authentication',
   SAVE_STUDIO_TOKEN = 'save-studio-token',
   SET_COMPARISON_MULTI_PLOT_VALUE = 'update-comparison-multi-plot-value',
   SET_SMOOTH_PLOT_VALUE = 'update-smooth-plot-value',
@@ -292,7 +294,9 @@ export type MessageFromWebview =
   | { type: MessageFromWebviewType.SETUP_WORKSPACE }
   | { type: MessageFromWebviewType.OPEN_STUDIO }
   | { type: MessageFromWebviewType.OPEN_STUDIO_PROFILE }
+  | { type: MessageFromWebviewType.OPEN_STUDIO_AUTH_LINK }
   | { type: MessageFromWebviewType.SAVE_STUDIO_TOKEN }
+  | { type: MessageFromWebviewType.REQUEST_STUDIO_AUTH }
   | { type: MessageFromWebviewType.ADD_CONFIGURATION }
   | { type: MessageFromWebviewType.ZOOM_PLOT; payload?: string }
   | { type: MessageFromWebviewType.OPEN_EXPERIMENTS_WEBVIEW }

webview/src/setup/components/App.test.tsx

@@ -41,7 +41,8 @@ const DEFAULT_DATA = {
   pythonBinPath: undefined,
   remoteList: undefined,
   sectionCollapsed: undefined,
-  shareLiveToStudio: false
+  shareLiveToStudio: false,
+  studioUserCode: null
 }
 
 const renderApp = (overrideData: Partial<SetupData> = {}) => {

webview/src/setup/components/App.tsx

@@ -33,7 +33,8 @@ import {
 import { updateRemoteList } from '../state/remoteSlice'
 import {
   updateIsStudioConnected,
-  updateShareLiveToStudio
+  updateShareLiveToStudio,
+  updateStudioUserCode
 } from '../state/studioSlice'
 import { setStudioShareExperimentsLive } from '../util/messages'
 
@@ -119,6 +120,9 @@ export const feedStore = (
       case 'shareLiveToStudio':
         dispatch(updateShareLiveToStudio(data.data.shareLiveToStudio))
         continue
+      case 'studioUserCode':
+        dispatch(updateStudioUserCode(data.data.studioUserCode))
+        continue
       case 'remoteList':
         dispatch(updateRemoteList(data.data.remoteList))
         continue

webview/src/setup/components/studio/Connect.tsx

@@ -1,46 +1,49 @@
 import React from 'react'
+import { useSelector } from 'react-redux'
 import { STUDIO_URL } from 'dvc/src/setup/webview/contract'
-import {
-  openStudio,
-  openStudioProfile,
-  saveStudioToken
-} from '../../util/messages'
+import { requestStudioAuth, openStudioAuthLink } from '../../util/messages'
 import { EmptyState } from '../../../shared/components/emptyState/EmptyState'
 import { Button } from '../../../shared/components/button/Button'
+import { SetupState } from '../../store'
 
 export const Connect: React.FC = () => {
+  const { studioUserCode } = useSelector((state: SetupState) => state.studio)
   return (
     <EmptyState isFullScreen={false}>
       <div data-testid="setup-studio-content">
         <h1>
           Connect to <a href={STUDIO_URL}>Studio</a>
         </h1>
-        <p>
-          Share experiments and plots with collaborators directly from your IDE.
-          Start sending data with an{' '}
-          <a href="https://dvc.org/doc/studio/user-guide/projects-and-experiments/live-metrics-and-plots#set-up-an-access-token">
-            access token
-          </a>{' '}
-          generated from your Studio profile page.
-        </p>
-        <Button
-          appearance="primary"
-          isNested={false}
-          text="Sign In"
-          onClick={openStudio}
-        />
-        <Button
-          appearance="secondary"
-          isNested={true}
-          text="Get Token"
-          onClick={openStudioProfile}
-        />
-        <Button
-          appearance="secondary"
-          isNested={true}
-          text="Save Token"
-          onClick={saveStudioToken}
-        />
+        {studioUserCode ? (
+          <>
+            <p>
+              We sent a token request to Studio. Enter the code shown below to
+              verify your identity.
+            </p>
+            <p>{studioUserCode}</p>
+            <Button
+              appearance="secondary"
+              text="Verify Identity"
+              onClick={openStudioAuthLink}
+            />
+          </>
+        ) : (
+          <>
+            <p>
+              Share experiments and plots with collaborators directly from your
+              IDE. Start sending data with an{' '}
+              <a href="https://dvc.org/doc/studio/user-guide/projects-and-experiments/live-metrics-and-plots#set-up-an-access-token">
+                access token
+              </a>{' '}
+              generated from your Studio profile page.
+            </p>
+            <Button
+              appearance="secondary"
+              text="Get Token"
+              onClick={requestStudioAuth}
+            />
+          </>
+        )}
         <p>
           Don&apos;t Have an account? <a href={STUDIO_URL}>Get Started</a>
         </p>

webview/src/setup/state/dvcSlice.ts

@@ -8,6 +8,7 @@ export type DvcState = Omit<
   | 'needsGitCommit'
   | 'sectionCollapsed'
   | 'shareLiveToStudio'
+  | 'studioUserCode'
 >
 
 export const dvcInitialState: DvcState = {