Redirect user to add new Studio access token on 401 response (#3311)

Author: mattseddon

extension/src/patch.ts

@@ -1,5 +1,6 @@
+import { commands } from 'vscode'
 import omit from 'lodash.omit'
-import fetch, { Response } from 'node-fetch'
+import fetch, { Response as FetchResponse } from 'node-fetch'
 import {
   EXPERIMENT_WORKSPACE_ID,
   ExperimentFields,
@@ -9,7 +10,10 @@ import {
 } from './cli/dvc/contract'
 import { AvailableCommands, InternalCommands } from './commands/internal'
 import { Args, ExperimentFlag } from './cli/dvc/constants'
+import { Response as UserResponse } from './vscode/response'
 import { Toast } from './vscode/toast'
+import { Modal } from './vscode/modal'
+import { RegisteredCommands } from './commands/external'
 
 export const STUDIO_ENDPOINT = 'https://studio.iterative.ai/api/live'
 
@@ -93,7 +97,7 @@ const collectExperimentDetails = (
 const sendPostRequest = (
   studioAccessToken: string,
   body: StartRequestBody | DataRequestBody | DoneRequestBody
-): Promise<Response> =>
+): Promise<FetchResponse> =>
   fetch(STUDIO_ENDPOINT, {
     body: JSON.stringify(body),
     headers: {
@@ -103,6 +107,16 @@ const sendPostRequest = (
     method: 'POST'
   })
 
+const showUnauthorized = async () => {
+  const response = await Modal.errorWithOptions(
+    'The current Studio access token is invalid. Please add a new token.',
+    UserResponse.SHOW
+  )
+  if (response === UserResponse.SHOW) {
+    return commands.executeCommand(RegisteredCommands.CONNECT_SHOW)
+  }
+}
+
 const shareWithProgress = (
   experimentDetails: ExperimentDetails,
   repoUrl: string,
@@ -121,7 +135,16 @@ const shareWithProgress = (
       increment: 0,
       message: 'Initializing experiment...'
     })
-    await sendPostRequest(studioAccessToken, { ...base, type: 'start' })
+    const response = await sendPostRequest(studioAccessToken, {
+      ...base,
+      type: 'start'
+    })
+
+    if (response.status === 401) {
+      progress.report({ increment: 100, message: 'Access unauthorized' })
+      void showUnauthorized()
+      return Toast.delayProgressClosing()
+    }
 
     progress.report({ increment: 33, message: 'Sending data...' })
     await sendPostRequest(studioAccessToken, {

extension/src/test/suite/patch.test.ts

@@ -1,8 +1,9 @@
 import { join } from 'path'
 import { afterEach, beforeEach, describe, it, suite } from 'mocha'
-import { restore, stub } from 'sinon'
+import { restore, spy, stub } from 'sinon'
 import { expect } from 'chai'
 import * as Fetch from 'node-fetch'
+import { commands } from 'vscode'
 import { buildInternalCommands, closeAllEditors } from './util'
 import { PROGRESS_TEST_TIMEOUT } from './timeouts'
 import { Disposable } from '../../extension'
@@ -12,6 +13,9 @@ import expShowFixture from '../fixtures/expShow/base/output'
 import { dvcDemoPath } from '../util'
 import { ExperimentFields } from '../../cli/dvc/contract'
 import { Toast } from '../../vscode/toast'
+import { Modal } from '../../vscode/modal'
+import { Response } from '../../vscode/response'
+import { RegisteredCommands } from '../../commands/external'
 
 suite('Patch Test Suite', () => {
   const disposable = Disposable.fn()
@@ -27,7 +31,7 @@ suite('Patch Test Suite', () => {
 
   describe('exp push patch', () => {
     it('should share an experiment to Studio', async () => {
-      const mockFetch = stub(Fetch, 'default').resolves(undefined)
+      const mockFetch = stub(Fetch, 'default').resolves({} as Fetch.Response)
       const mockStudioAccessToken = 'isat_12123123123123123'
       const mockRepoUrl = 'https://github.com/iterative/vscode-dvc-demo'
 
@@ -98,6 +102,66 @@ suite('Patch Test Suite', () => {
       })
     }).timeout(PROGRESS_TEST_TIMEOUT)
 
+    it('should show an error modal if Studio returns a 401 response', async () => {
+      const mockFetch = stub(Fetch, 'default').resolves({
+        status: 401
+      } as Fetch.Response)
+      const mockStudioAccessToken = 'isat_12123123123123123'
+      const mockRepoUrl = 'https://github.com/iterative/vscode-dvc-demo'
+
+      const executeCommandSpy = spy(commands, 'executeCommand')
+
+      const { internalCommands, gitReader, dvcReader } =
+        buildInternalCommands(disposable)
+
+      const mockGetRemoteUrl = stub(gitReader, 'getRemoteUrl').resolves(
+        mockRepoUrl
+      )
+      const mockExpShow = stub(dvcReader, 'expShow').resolves(expShowFixture)
+
+      const mockErrorWithOptions = stub(Modal, 'errorWithOptions').resolves(
+        Response.SHOW
+      )
+
+      registerPatchCommand(internalCommands)
+
+      await internalCommands.executeCommand(
+        AvailableCommands.EXP_PUSH,
+        mockStudioAccessToken,
+        dvcDemoPath,
+        'exp-e7a67'
+      )
+
+      expect(mockGetRemoteUrl).to.be.calledOnce
+      expect(mockExpShow).to.be.calledOnce
+      expect(mockFetch).to.be.calledOnce
+      expect(mockErrorWithOptions).to.be.calledOnce
+      expect(executeCommandSpy).to.be.calledWithExactly(
+        RegisteredCommands.CONNECT_SHOW
+      )
+
+      const { name } = expShowFixture[
+        '53c3851f46955fa3e2b8f6e1c52999acc8c9ea77'
+      ]['4fb124aebddb2adf1545030907687fa9a4c80e70'].data as ExperimentFields
+
+      const baseBody = {
+        baseline_sha: '53c3851f46955fa3e2b8f6e1c52999acc8c9ea77',
+        client: 'vscode',
+        name,
+        repo_url: mockRepoUrl
+      }
+      const headers = {
+        Authorization: `token ${mockStudioAccessToken}`,
+        'Content-type': 'application/json'
+      }
+
+      expect(mockFetch).to.be.calledWithExactly(STUDIO_ENDPOINT, {
+        body: JSON.stringify({ ...baseBody, type: 'start' }),
+        headers,
+        method: 'POST'
+      })
+    }).timeout(PROGRESS_TEST_TIMEOUT)
+
     it('should show an error message if the experiment cannot be found', async () => {
       const mockShowError = stub(Toast, 'showError').resolves(undefined)
 

extension/src/vscode/modal.ts

@@ -2,6 +2,7 @@ import { window } from 'vscode'
 import { Response } from './response'
 
 enum Level {
+  ERROR = 'Error',
   INFORMATION = 'Information',
   WARNING = 'Warning'
 }
@@ -15,6 +16,10 @@ export class Modal {
     return Modal.show(Level.WARNING, text, ...items)
   }
 
+  public static errorWithOptions(text: string, ...items: Response[]) {
+    return Modal.show(Level.ERROR, text, ...items)
+  }
+
   private static show(
     level: Level,
     message: string,