Provide shareable links for all experiments found on Studio (#4600)

* hack together new integration

* fix webview tests

* fix integration tests

* refactor

* add test for call to Studio

extension/src/data/index.ts

@@ -18,13 +18,27 @@ type LocalExperimentsOutput = {
 
 type RemoteExperimentsOutput = { lsRemoteOutput: string }
 
-export type ExperimentsOutput = LocalExperimentsOutput | RemoteExperimentsOutput
+type StudioExperimentsOutput = {
+  baseUrl: string | undefined
+  live: { baselineSha: string; name: string }[]
+  pushed: string[]
+}
+
+export type ExperimentsOutput =
+  | LocalExperimentsOutput
+  | RemoteExperimentsOutput
+  | StudioExperimentsOutput
 
 export const isRemoteExperimentsOutput = (
   data: ExperimentsOutput
 ): data is RemoteExperimentsOutput =>
   (data as RemoteExperimentsOutput).lsRemoteOutput !== undefined
 
+export const isStudioExperimentsOutput = (
+  data: ExperimentsOutput
+): data is StudioExperimentsOutput =>
+  (data as StudioExperimentsOutput).live !== undefined
+
 export abstract class BaseData<
   T extends
     | { data: PlotsOutputOrError; revs: string[] }

extension/src/experiments/data/index.ts

@@ -1,3 +1,5 @@
+import querystring from 'querystring'
+import fetch from 'node-fetch'
 import { collectBranches, collectFiles } from './collect'
 import {
   EXPERIMENTS_GIT_LOGS_REFS,
@@ -11,7 +13,8 @@ import { ExpShowOutput } from '../../cli/dvc/contract'
 import {
   BaseData,
   ExperimentsOutput,
-  isRemoteExperimentsOutput
+  isRemoteExperimentsOutput,
+  isStudioExperimentsOutput
 } from '../../data'
 import {
   Args,
@@ -22,14 +25,18 @@ import {
 import { COMMITS_SEPARATOR, gitPath } from '../../cli/git/constants'
 import { getGitPath } from '../../fileSystem'
 import { ExperimentsModel } from '../model'
+import { Studio } from '../studio'
+import { STUDIO_URL } from '../../setup/webview/contract'
 
 export class ExperimentsData extends BaseData<ExperimentsOutput> {
   private readonly experiments: ExperimentsModel
+  private readonly studio: Studio
 
   constructor(
     dvcRoot: string,
     internalCommands: InternalCommands,
     experiments: ExperimentsModel,
+    studio: Studio,
     subProjects: string[]
   ) {
     super(
@@ -47,6 +54,7 @@ export class ExperimentsData extends BaseData<ExperimentsOutput> {
     )
 
     this.experiments = experiments
+    this.studio = studio
 
     void this.watchExpGitRefs()
     void this.managedUpdate()
@@ -58,10 +66,13 @@ export class ExperimentsData extends BaseData<ExperimentsOutput> {
   }
 
   public async update(): Promise<void> {
-    await Promise.all([this.updateExpShow(), this.updateRemoteExpRefs()])
+    await Promise.all([
+      this.updateExpShowAndStudio(),
+      this.updateRemoteExpRefs()
+    ])
   }
 
-  private async updateExpShow() {
+  private async updateExpShowAndStudio() {
     await this.updateBranches()
     const [currentBranch, ...branches] = this.experiments.getBranchesToShow()
     const availableNbCommits: { [branch: string]: number } = {}
@@ -77,8 +88,21 @@ export class ExperimentsData extends BaseData<ExperimentsOutput> {
     }
 
     const branchLogs = await Promise.all(promises)
-    const { args, gitLog, rowOrder } = this.collectGitLogAndOrder(branchLogs)
+    const { args, gitLog, rowOrder, shas } =
+      this.collectGitLogAndOrder(branchLogs)
+
+    return Promise.all([
+      this.updateExpShow(args, availableNbCommits, gitLog, rowOrder),
+      this.requestStudioData(shas)
+    ])
+  }
 
+  private async updateExpShow(
+    args: Args,
+    availableNbCommits: { [branch: string]: number },
+    gitLog: string,
+    rowOrder: { branch: string; sha: string }[]
+  ) {
     const expShow = await this.internalCommands.executeCommand<ExpShowOutput>(
       AvailableCommands.EXP_SHOW,
       this.dvcRoot,
@@ -90,6 +114,49 @@ export class ExperimentsData extends BaseData<ExperimentsOutput> {
     this.collectFiles({ expShow })
   }
 
+  private async requestStudioData(shas: string[]) {
+    await this.studio.isReady()
+
+    const defaultData = { baseUrl: undefined, live: [], pushed: [] }
+
+    const studioAccessToken = this.studio.getAccessToken()
+
+    if (!studioAccessToken || shas.length === 0) {
+      this.notifyChanged(defaultData)
+      return
+    }
+
+    const params = querystring.stringify({
+      commits: shas,
+      git_remote_url: this.studio.getGitRemoteUrl()
+    })
+
+    try {
+      const response = await fetch(`${STUDIO_URL}/api/view-links?${params}`, {
+        headers: {
+          Authorization: `token ${studioAccessToken}`
+        },
+        method: 'GET'
+      })
+
+      const { live, pushed, view_url } = (await response.json()) as {
+        live: { baseline_sha: string; name: string }[]
+        pushed: string[]
+        view_url: string
+      }
+      this.notifyChanged({
+        baseUrl: view_url,
+        live: live.map(({ baseline_sha, name }) => ({
+          baselineSha: baseline_sha,
+          name
+        })),
+        pushed
+      })
+    } catch {
+      this.notifyChanged(defaultData)
+    }
+  }
+
   private async collectGitLogByBranch(
     branch: string,
     availableNbCommits: { [branch: string]: number },
@@ -122,6 +189,7 @@ export class ExperimentsData extends BaseData<ExperimentsOutput> {
   ) {
     const rowOrder: { branch: string; sha: string }[] = []
     const args: Args = []
+    const shas = []
     const gitLog: string[] = []
 
     for (const { branch, branchLog } of branchLogs) {
@@ -133,9 +201,10 @@ export class ExperimentsData extends BaseData<ExperimentsOutput> {
           continue
         }
         args.push(ExperimentFlag.REV, sha)
+        shas.push(sha)
       }
     }
-    return { args, gitLog: gitLog.join(COMMITS_SEPARATOR), rowOrder }
+    return { args, gitLog: gitLog.join(COMMITS_SEPARATOR), rowOrder, shas }
   }
 
   private async updateBranches() {
@@ -200,7 +269,10 @@ export class ExperimentsData extends BaseData<ExperimentsOutput> {
   private waitForInitialLocalData() {
     const waitForInitialData = this.dispose.track(
       this.onDidUpdate(data => {
-        if (isRemoteExperimentsOutput(data)) {
+        if (
+          isRemoteExperimentsOutput(data) ||
+          isStudioExperimentsOutput(data)
+        ) {
           return
         }
         this.dispose.untrack(waitForInitialData)

extension/src/experiments/index.ts

@@ -35,7 +35,11 @@ import { DecorationProvider } from './model/decorationProvider'
 import { starredFilter } from './model/filterBy/constants'
 import { ResourceLocator } from '../resourceLocator'
 import { AvailableCommands, InternalCommands } from '../commands/internal'
-import { ExperimentsOutput, isRemoteExperimentsOutput } from '../data'
+import {
+  ExperimentsOutput,
+  isRemoteExperimentsOutput,
+  isStudioExperimentsOutput
+} from '../data'
 import { ViewKey } from '../webview/constants'
 import { BaseRepository } from '../webview/repository'
 import { Title } from '../vscode/title'
@@ -159,6 +163,7 @@ export class Experiments extends BaseRepository<TableData> {
           dvcRoot,
           internalCommands,
           this.experiments,
+          this.studio,
           subProjects
         )
     )
@@ -189,6 +194,13 @@ export class Experiments extends BaseRepository<TableData> {
       return this.webviewMessages.sendWebviewMessage()
     }
 
+    if (isStudioExperimentsOutput(data)) {
+      const { live, pushed, baseUrl } = data
+      this.studio.setBaseUrl(baseUrl)
+      this.experiments.setStudioData(live, pushed)
+      return this.webviewMessages.sendWebviewMessage()
+    }
+
     const { expShow, gitLog, rowOrder, availableNbCommits } = data
 
     const hadCheckpoints = this.hasCheckpoints()
@@ -597,13 +609,15 @@ export class Experiments extends BaseRepository<TableData> {
     return this.data.update()
   }
 
-  public async setStudioBaseUrl(studioToken: string | undefined) {
-    await Promise.all([this.isReady(), this.experiments.isReady()])
-    await this.studio.setBaseUrl(
-      studioToken,
-      this.experiments.getRemoteExpRefs()
-    )
-    return this.webviewMessages.sendWebviewMessage()
+  public setStudioAccessToken(studioAccessToken: string | undefined) {
+    const oldAccessToken = this.studio.getAccessToken()
+    const accessTokenInitialized = this.studio.isAccessTokenSet()
+    this.studio.setAccessToken(studioAccessToken)
+
+    if (!accessTokenInitialized || oldAccessToken === studioAccessToken) {
+      return
+    }
+    return this.data.managedUpdate()
   }
 
   public hasDvcLiveOnlyRunning() {

extension/src/experiments/model/collect.ts

@@ -262,7 +262,7 @@ const collectExpRange = (
   const expState = revs[0]
 
   const { name, rev } = expState
-  const { id: baselineId } = baseline
+  const { id: baselineId, sha: baselineSha } = baseline
 
   const label =
     rev === EXPERIMENT_WORKSPACE_ID
@@ -292,6 +292,8 @@ const collectExpRange = (
     experiment.description = `[${name}]`
   }
 
+  experiment.baselineSha = baselineSha
+
   collectExecutorInfo(experiment, executor)
   collectRunningExperiment(acc, experiment)
 

extension/src/experiments/model/index.test.ts

@@ -57,6 +57,7 @@ describe('ExperimentsModel', () => {
       { main: 6 }
     )
     model.transformAndSetRemote(remoteExpRefsFixture)
+    model.setStudioData([], ['42b8736b08170529903cd203a1f40382a4b4a8cd'])
     expect(model.getRowData()).toStrictEqual(rowsFixture)
   })
 

extension/src/experiments/model/index.ts

@@ -26,7 +26,8 @@ import {
   isRunning,
   GitRemoteStatus,
   RunningExperiment,
-  WORKSPACE_BRANCH
+  WORKSPACE_BRANCH,
+  StudioLinkType
 } from '../webview/contract'
 import { reorderListSubset } from '../../util/array'
 import {
@@ -80,9 +81,13 @@ export class ExperimentsModel extends ModelWithPersistence {
   private filters: Map<string, FilterDefinition> = new Map()
 
   private remoteExpShas?: Set<string>
-  private remoteExpRefs: string[] = []
   private pushing = new Set<string>()
 
+  private studioLiveOnlyExperiments: { baselineSha: string; name: string }[] =
+    []
+
+  private studioPushedExperiments: string[] = []
+
   private currentSorts: SortDefinition[]
   private running: RunningExperiment[] = []
   private startedRunning: Set<string> = new Set()
@@ -176,10 +181,8 @@ export class ExperimentsModel extends ModelWithPersistence {
   }
 
   public transformAndSetRemote(lsRemoteOutput: string) {
-    const { remoteExpShas, remoteExpRefs } =
-      collectRemoteExpDetails(lsRemoteOutput)
+    const { remoteExpShas } = collectRemoteExpDetails(lsRemoteOutput)
     this.remoteExpShas = remoteExpShas
-    this.remoteExpRefs = remoteExpRefs
     this.deferred.resolve()
   }
 
@@ -390,6 +393,16 @@ export class ExperimentsModel extends ModelWithPersistence {
     })
   }
 
+  public getExperimentShas(id: string) {
+    for (const experiment of this.getExperiments()) {
+      if (experiment.id === id) {
+        const { baselineSha, sha } = experiment
+        return { baselineSha, sha }
+      }
+    }
+    return { baselineSha: undefined, sha: undefined }
+  }
+
   public getCommitsAndExperiments() {
     return collectOrderedCommitsAndExperiments(this.commits, commit =>
       this.getExperimentsByCommit(commit)
@@ -543,8 +556,12 @@ export class ExperimentsModel extends ModelWithPersistence {
     return this.availableBranchesToSelect
   }
 
-  public getRemoteExpRefs() {
-    return this.remoteExpRefs
+  public setStudioData(
+    live: { baselineSha: string; name: string }[],
+    pushed: string[]
+  ) {
+    this.studioLiveOnlyExperiments = live
+    this.studioPushedExperiments = pushed
   }
 
   public hasDvcLiveOnlyRunning() {
@@ -596,6 +613,7 @@ export class ExperimentsModel extends ModelWithPersistence {
         const experiment = this.addDetails(originalExperiment)
 
         this.addRemoteStatus(experiment)
+        this.addStudioLinkType(experiment, commit.sha)
 
         return experiment
       })
@@ -623,6 +641,35 @@ export class ExperimentsModel extends ModelWithPersistence {
       : GitRemoteStatus.NOT_ON_REMOTE
   }
 
+  private addStudioLinkType(
+    experiment: Experiment,
+    baselineSha: string | undefined
+  ) {
+    if (
+      this.studioLiveOnlyExperiments.length === 0 &&
+      this.studioPushedExperiments.length === 0
+    ) {
+      return
+    }
+
+    if (
+      experiment.sha &&
+      this.studioPushedExperiments.includes(experiment.sha)
+    ) {
+      experiment.studioLinkType = StudioLinkType.PUSHED
+      return
+    }
+
+    if (
+      this.studioLiveOnlyExperiments.some(
+        ({ baselineSha: expBaselineSha, name }) =>
+          baselineSha === expBaselineSha && experiment.id === name
+      )
+    ) {
+      experiment.studioLinkType = StudioLinkType.LIVE
+    }
+  }
+
   private setColoredStatus(runningExperiments: RunningExperiment[]) {
     this.setRunning(runningExperiments)