Add self-hosted url info in Studio setup section

extension/src/setup/index.ts

@@ -437,6 +437,7 @@ export class Setup
       pythonBinPath: getBinDisplayText(pythonBinPath),
       remoteList,
       sectionCollapsed: collectSectionCollapsed(this.focusedSection),
+      selfHostedStudioUrl: this.studio.getSelfHostedStudioUrl(),
       shareLiveToStudio: !!this.studio.getShareLiveToStudio()
     })
     this.focusedSection = undefined

extension/src/setup/studio.ts

@@ -58,6 +58,11 @@ export class Studio extends Disposable {
     return this.studioUrl
   }
 
+  public getSelfHostedStudioUrl() {
+    const url = this.getStudioUrl()
+    return url === DEFAULT_STUDIO_URL ? null : url
+  }
+
   public removeStudioAccessToken(dvcRoots: string[]) {
     return this.removeKeyFromConfig(dvcRoots, ConfigKey.STUDIO_TOKEN)
   }

extension/src/setup/webview/contract.ts

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

extension/src/setup/webview/messages.ts

@@ -70,10 +70,14 @@ export class WebviewMessages {
         return commands.executeCommand(
           RegisteredCommands.ADD_STUDIO_ACCESS_TOKEN
         )
+      case MessageFromWebviewType.SAVE_STUDIO_URL:
+        return commands.executeCommand(RegisteredCommands.UPDATE_STUDIO_URL)
       case MessageFromWebviewType.REMOVE_STUDIO_TOKEN:
         return commands.executeCommand(
           RegisteredCommands.REMOVE_STUDIO_ACCESS_TOKEN
         )
+      case MessageFromWebviewType.REMOVE_STUDIO_URL:
+        return commands.executeCommand(RegisteredCommands.REMOVE_STUDIO_URL)
       case MessageFromWebviewType.SET_STUDIO_SHARE_EXPERIMENTS_LIVE:
         return this.updateStudioOffline(message.payload)
       case MessageFromWebviewType.REQUEST_STUDIO_TOKEN:

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

@@ -395,6 +395,7 @@ suite('Setup Test Suite', () => {
         pythonBinPath: undefined,
         remoteList: undefined,
         sectionCollapsed: undefined,
+        selfHostedStudioUrl: null,
         shareLiveToStudio: false
       })
     }).timeout(WEBVIEW_TEST_TIMEOUT)
@@ -441,6 +442,7 @@ suite('Setup Test Suite', () => {
         pythonBinPath: undefined,
         remoteList: undefined,
         sectionCollapsed: undefined,
+        selfHostedStudioUrl: null,
         shareLiveToStudio: true
       })
     }).timeout(WEBVIEW_TEST_TIMEOUT)
@@ -495,6 +497,7 @@ suite('Setup Test Suite', () => {
         pythonBinPath: undefined,
         remoteList: undefined,
         sectionCollapsed: undefined,
+        selfHostedStudioUrl: null,
         shareLiveToStudio: true
       })
     }).timeout(WEBVIEW_TEST_TIMEOUT)
@@ -549,6 +552,7 @@ suite('Setup Test Suite', () => {
         pythonBinPath: undefined,
         remoteList: { [dvcDemoPath]: undefined },
         sectionCollapsed: undefined,
+        selfHostedStudioUrl: null,
         shareLiveToStudio: true
       })
     }).timeout(WEBVIEW_TEST_TIMEOUT)
@@ -1332,6 +1336,68 @@ suite('Setup Test Suite', () => {
       )
     })
 
+    it("should handle a message from the webview to update the user's self hosted url", async () => {
+      const { setup, mockExecuteCommand } = buildSetup({
+        disposer: disposable
+      })
+
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      stub(Setup.prototype as any, 'getCliCompatible').returns(true)
+
+      const webview = await setup.showWebview()
+      await webview.isReady()
+
+      const mockMessageReceived = getMessageReceivedEmitter(webview)
+
+      const commandCalled = new Promise(resolve =>
+        mockExecuteCommand.callsFake(() => {
+          resolve(undefined)
+          return Promise.resolve(undefined)
+        })
+      )
+
+      mockMessageReceived.fire({
+        type: MessageFromWebviewType.SAVE_STUDIO_URL
+      })
+
+      await commandCalled
+
+      expect(mockExecuteCommand).to.be.calledWithExactly(
+        RegisteredCommands.UPDATE_STUDIO_URL
+      )
+    }).timeout(WEBVIEW_TEST_TIMEOUT)
+
+    it("should handle a message from the webview to remove the user's self hosted url", async () => {
+      const { setup, mockExecuteCommand } = buildSetup({
+        disposer: disposable
+      })
+
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      stub(Setup.prototype as any, 'getCliCompatible').returns(true)
+
+      const webview = await setup.showWebview()
+      await webview.isReady()
+
+      const commandCalled = new Promise(resolve =>
+        mockExecuteCommand.callsFake(() => {
+          resolve(undefined)
+          return Promise.resolve(undefined)
+        })
+      )
+
+      const mockMessageReceived = getMessageReceivedEmitter(webview)
+
+      mockMessageReceived.fire({
+        type: MessageFromWebviewType.SAVE_STUDIO_URL
+      })
+
+      await commandCalled
+
+      expect(mockExecuteCommand).to.be.calledWithExactly(
+        RegisteredCommands.UPDATE_STUDIO_URL
+      )
+    }).timeout(WEBVIEW_TEST_TIMEOUT)
+
     it('should check if experiments and dvc are setup', async () => {
       const { setup } = buildSetup({
         disposer: disposable,

extension/src/webview/contract.ts

@@ -47,6 +47,7 @@ export enum MessageFromWebviewType {
   RESIZE_PLOTS = 'resize-plots',
   REQUEST_STUDIO_TOKEN = 'request-studio-token',
   SAVE_STUDIO_TOKEN = 'save-studio-token',
+  SAVE_STUDIO_URL = 'save-studio-url',
   SET_COMPARISON_MULTI_PLOT_VALUE = 'update-comparison-multi-plot-value',
   SET_SMOOTH_PLOT_VALUE = 'update-smooth-plot-value',
   SHOW_EXPERIMENT_LOGS = 'show-experiment-logs',
@@ -71,6 +72,7 @@ export enum MessageFromWebviewType {
   REMOTE_REMOVE = 'remote-remove',
   REMOVE_CUSTOM_PLOTS = 'remove-custom-plots',
   REMOVE_STUDIO_TOKEN = 'remove-studio-token',
+  REMOVE_STUDIO_URL = 'remove-studio-url',
   MODIFY_WORKSPACE_PARAMS_AND_QUEUE = 'modify-workspace-params-and-queue',
   MODIFY_WORKSPACE_PARAMS_AND_RUN = 'modify-workspace-params-and-run',
   MODIFY_WORKSPACE_PARAMS_RESET_AND_RUN = 'modify-workspace-params-reset-and-run',
@@ -224,6 +226,7 @@ export type MessageFromWebview =
       type: MessageFromWebviewType.REMOVE_CUSTOM_PLOTS
     }
   | { type: MessageFromWebviewType.REMOVE_STUDIO_TOKEN }
+  | { type: MessageFromWebviewType.REMOVE_STUDIO_URL }
   | {
       type: MessageFromWebviewType.REORDER_PLOTS_COMPARISON
       payload: string[]
@@ -291,6 +294,7 @@ export type MessageFromWebview =
   | { type: MessageFromWebviewType.UPGRADE_DVC }
   | { type: MessageFromWebviewType.SETUP_WORKSPACE }
   | { type: MessageFromWebviewType.SAVE_STUDIO_TOKEN }
+  | { type: MessageFromWebviewType.SAVE_STUDIO_URL }
   | { type: MessageFromWebviewType.REQUEST_STUDIO_TOKEN }
   | { type: MessageFromWebviewType.ADD_CONFIGURATION }
   | { type: MessageFromWebviewType.ZOOM_PLOT; payload?: string }

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

@@ -41,6 +41,7 @@ const DEFAULT_DATA = {
   pythonBinPath: undefined,
   remoteList: undefined,
   sectionCollapsed: undefined,
+  selfHostedStudioUrl: null,
   shareLiveToStudio: false
 }
 
@@ -675,6 +676,65 @@ describe('App', () => {
         within(iconWrapper).getByTestId(TooltipIconType.WARNING)
       ).toBeInTheDocument()
     })
+
+    it('should show the self hosted url with actions to change it if the user has set one', () => {
+      const selfHostedUrl = 'https://studio.example.com'
+      renderApp({ selfHostedStudioUrl: selfHostedUrl })
+
+      const urlDetails = screen.getByTestId('studio-url-details')
+
+      expect(
+        within(urlDetails).getByText('Self-Hosted Url:')
+      ).toBeInTheDocument()
+      expect(within(urlDetails).getByText(selfHostedUrl)).toBeInTheDocument()
+
+      const updateUrlBtn = within(urlDetails).getByText('Update')
+      const removeUrlBtn = within(urlDetails).getByText('Remove')
+
+      expect(updateUrlBtn).toBeInTheDocument()
+      expect(removeUrlBtn).toBeInTheDocument()
+
+      mockPostMessage.mockClear()
+      fireEvent.click(updateUrlBtn)
+
+      expect(mockPostMessage).toHaveBeenCalledTimes(1)
+      expect(mockPostMessage).toHaveBeenCalledWith({
+        type: MessageFromWebviewType.SAVE_STUDIO_URL
+      })
+
+      mockPostMessage.mockClear()
+      fireEvent.click(removeUrlBtn)
+
+      expect(mockPostMessage).toHaveBeenCalledTimes(1)
+      expect(mockPostMessage).toHaveBeenCalledWith({
+        type: MessageFromWebviewType.REMOVE_STUDIO_URL
+      })
+    })
+
+    it('should show the self hosted url with "Not found" with an action to add one if the user has not set one', () => {
+      renderApp()
+
+      const urlDetails = screen.getByTestId('studio-url-details')
+
+      expect(
+        within(urlDetails).getByText('Self-Hosted Url:')
+      ).toBeInTheDocument()
+      expect(within(urlDetails).getByText('Not found')).toBeInTheDocument()
+
+      const addUrlBtn = within(urlDetails).getByText('Add Url')
+
+      expect(addUrlBtn).toBeInTheDocument()
+
+      mockPostMessage.mockClear()
+      fireEvent.click(addUrlBtn)
+
+      expect(mockPostMessage).toHaveBeenCalledTimes(1)
+      expect(mockPostMessage).toHaveBeenCalledWith({
+        type: MessageFromWebviewType.SAVE_STUDIO_URL
+      })
+
+      mockPostMessage.mockClear()
+    })
   })
 
   describe('Studio connected', () => {
@@ -718,6 +778,24 @@ describe('App', () => {
         within(iconWrapper).getByTestId(TooltipIconType.PASSED)
       ).toBeInTheDocument()
     })
+
+    it('should show the self hosted url with actions to change it if the user has set one', () => {
+      const selfHostedUrl = 'https://studio.example.com'
+      renderApp({ isStudioConnected: true, selfHostedStudioUrl: selfHostedUrl })
+
+      const urlDetails = screen.getByTestId('studio-url-details')
+
+      expect(
+        within(urlDetails).getByText('Self-Hosted Url:')
+      ).toBeInTheDocument()
+      expect(within(urlDetails).getByText(selfHostedUrl)).toBeInTheDocument()
+
+      const updateUrlBtn = within(urlDetails).getByText('Update')
+      const removeUrlBtn = within(urlDetails).getByText('Remove')
+
+      expect(updateUrlBtn).toBeInTheDocument()
+      expect(removeUrlBtn).toBeInTheDocument()
+    })
   })
 
   describe('focused section', () => {

webview/src/setup/components/App.tsx

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