Improve Studio Token Authentication URI Handler (#5049)

Author: julieg18

extension/src/setup/studio.ts

@@ -1,4 +1,4 @@
-import { Event, EventEmitter } from 'vscode'
+import { Event, EventEmitter, Disposable as VSCodeDisposable } from 'vscode'
 import fetch from 'node-fetch'
 import { STUDIO_URL } from './webview/contract'
 import { AvailableCommands, InternalCommands } from '../commands/internal'
@@ -27,6 +27,7 @@ export class Studio extends Disposable {
   private studioAccessToken: string | undefined = undefined
   private studioIsConnected = false
   private shareLiveToStudio: boolean | undefined = undefined
+  private studioAccessTokenUriHandler: VSCodeDisposable | undefined = undefined
 
   constructor(
     internalCommands: InternalCommands,
@@ -108,6 +109,8 @@ export class Studio extends Disposable {
   }
 
   public async requestStudioTokenAuthentication() {
+    this.resetAccessTokenUriHandler()
+
     const response = await this.fetchFromStudio(
       `${STUDIO_URL}/api/device-login`,
       {
@@ -134,9 +137,17 @@ export class Studio extends Disposable {
     verificationUrlWithParams.searchParams.append('code', userCode)
 
     await openUrl(verificationUrlWithParams.toString())
-    void waitForUriResponse('/studio-complete-auth', () => {
-      void this.requestStudioToken(deviceCode, tokenUri)
-    })
+    this.studioAccessTokenUriHandler = waitForUriResponse(
+      '/studio-complete-auth',
+      () => {
+        void this.requestStudioToken(deviceCode, tokenUri)
+      }
+    )
+  }
+
+  private resetAccessTokenUriHandler() {
+    this.studioAccessTokenUriHandler?.dispose()
+    this.studioAccessTokenUriHandler = undefined
   }
 
   private fetchFromStudio(reqUri: string, body: Record<string, unknown>) {
@@ -171,6 +182,7 @@ export class Studio extends Disposable {
   }
 
   private requestStudioToken(deviceCode: string, tokenUri: string) {
+    this.resetAccessTokenUriHandler()
     return Toast.showProgress('Connecting to Studio', async progress => {
       progress.report({ increment: 0 })
       progress.report({ increment: 25, message: 'Fetching token...' })

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

@@ -873,6 +873,7 @@ suite('Setup Test Suite', () => {
       const mockGetCallbackUrl = stub(ExternalUtil, 'getCallBackUrl')
       const mockOpenUrl = stub(ExternalUtil, 'openUrl')
       const mockWaitForUriRes = stub(ExternalUtil, 'waitForUriResponse')
+      const mockUriHandlerDispose = stub()
       const mockStudioRes = {
         device_code: 'Yi-NPd9ggvNUDBcam5bP8iivbtLhnqVgM_lSSbilqNw',
         token_uri: 'https://studio.iterative.ai/api/device-login/token',
@@ -890,6 +891,7 @@ suite('Setup Test Suite', () => {
         resolve =>
           mockWaitForUriRes.onFirstCall().callsFake((_, onResponse) => {
             resolve(onResponse)
+            return { dispose: mockUriHandlerDispose }
           })
       )
 
@@ -942,6 +944,7 @@ suite('Setup Test Suite', () => {
 
       await failedTokenEvent
 
+      expect(mockUriHandlerDispose).to.be.calledOnce
       expect(mockFetch).to.be.calledTwice
       expect(mockFetch).to.be.calledWithExactly(mockStudioRes.token_uri, {
         body: JSON.stringify({

extension/src/vscode/external.test.ts

@@ -56,7 +56,7 @@ describe('waitForUriResponse', () => {
 
     const uriWithoutPath = Uri.from({
       authority: 'iterative.dvc',
-      path: '/not-correct-path',
+      path: '/path-but-not-the-right-one',
       scheme: mockedUriScheme
     })
     mockHandleUriResponse(uriWithoutPath)

extension/src/vscode/external.ts

@@ -10,12 +10,11 @@ export const getCallBackUrl = async (path: string) => {
   return uri.toString()
 }
 
-export const waitForUriResponse = (path: string, onResponse: () => unknown) => {
+export const waitForUriResponse = (path: string, onResponse: () => unknown) =>
   window.registerUriHandler({
     handleUri(uri: Uri): ProviderResult<void> {
-      if (uri.path.startsWith(path)) {
+      if (uri.path === path) {
         onResponse()
       }
     }
   })
-}