Call Studio API will all refs available on the Git remote (#4580)

extension/src/data/index.ts

@@ -16,14 +16,14 @@ type LocalExperimentsOutput = {
   rowOrder: { branch: string; sha: string }[]
 }
 
-type RemoteExperimentsOutput = { remoteExpRefs: string }
+type RemoteExperimentsOutput = { lsRemoteOutput: string }
 
 export type ExperimentsOutput = LocalExperimentsOutput | RemoteExperimentsOutput
 
 export const isRemoteExperimentsOutput = (
   data: ExperimentsOutput
 ): data is RemoteExperimentsOutput =>
-  (data as RemoteExperimentsOutput).remoteExpRefs !== undefined
+  (data as RemoteExperimentsOutput).lsRemoteOutput !== undefined
 
 export abstract class BaseData<
   T extends

extension/src/experiments/data/index.ts

@@ -191,15 +191,15 @@ export class ExperimentsData extends BaseData<ExperimentsOutput> {
   }
 
   private async updateRemoteExpRefs() {
-    const [remoteExpRefs] = await Promise.all([
+    const [lsRemoteOutput] = await Promise.all([
       this.internalCommands.executeCommand(
         AvailableCommands.GIT_GET_REMOTE_EXPERIMENT_REFS,
         this.dvcRoot
       ),
       this.isReady()
     ])
 
-    this.notifyChanged({ remoteExpRefs })
+    this.notifyChanged({ lsRemoteOutput })
   }
 
   private waitForInitialLocalData() {

extension/src/experiments/index.ts

@@ -184,8 +184,8 @@ export class Experiments extends BaseRepository<TableData> {
 
   public async setState(data: ExperimentsOutput) {
     if (isRemoteExperimentsOutput(data)) {
-      const { remoteExpRefs } = data
-      this.experiments.transformAndSetRemote(remoteExpRefs)
+      const { lsRemoteOutput } = data
+      this.experiments.transformAndSetRemote(lsRemoteOutput)
       return this.webviewMessages.sendWebviewMessage()
     }
 
@@ -598,8 +598,11 @@ export class Experiments extends BaseRepository<TableData> {
   }
 
   public async setStudioBaseUrl(studioToken: string | undefined) {
-    await this.isReady()
-    await this.studio.setBaseUrl(studioToken)
+    await Promise.all([this.isReady(), this.experiments.isReady()])
+    await this.studio.setBaseUrl(
+      studioToken,
+      this.experiments.getRemoteExpRefs()
+    )
     return this.webviewMessages.sendWebviewMessage()
   }
 

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

@@ -1,4 +1,4 @@
-import { collectExperiments } from './collect'
+import { collectExperiments, collectRemoteExpDetails } from './collect'
 import { generateTestExpShowOutput } from '../../test/util/experiments'
 import { ExpShowOutput } from '../../cli/dvc/contract'
 
@@ -101,3 +101,36 @@ describe('collectExperiments', () => {
     ])
   })
 })
+
+describe('collectRemoteExpDetails', () => {
+  it('should parse the git ls-remote output', () => {
+    const output = `263e4408e42a0e215b0f70b36b2ab7b65a160d7e        refs/exps/a9/b32d14966b9be1396f2211d9eb743359708a07/vital-taal
+    d4f2a35773ead55b7ce4b596f600e98360e49372        refs/exps/a9/b32d14966b9be1396f2211d9eb743359708a07/whole-bout
+    5af79e8d5e53f4e41221b6a166121d96d50b630a        refs/exps/a9/d8057e088d46842f15c3b6d1bb2e4befd5f677/deism-bots
+    21745a4aa76daf59b49ec81480fe7a89c7ea8fb2        refs/exps/a9/d8057e088d46842f15c3b6d1bb2e4befd5f677/inter-gulf
+    390aef747f45fc49ec8928b24771f8950d057393        refs/exps/a9/d8057e088d46842f15c3b6d1bb2e4befd5f677/known-flus
+    142a803b83ff784ba1106cc4ad0ba03310da6186        refs/exps/a9/d8057e088d46842f15c3b6d1bb2e4befd5f677/tight-lira
+    21ce298cd1743405a0d73f5cb4cf52289ffa3276        refs/exps/bf/6ca8a35911bc6e62fb9bcaa506d4f4e185450c/crumb-orcs`
+    const { remoteExpRefs, remoteExpShas } = collectRemoteExpDetails(output)
+    expect(remoteExpRefs).toStrictEqual([
+      'refs/exps/a9/b32d14966b9be1396f2211d9eb743359708a07/vital-taal',
+      'refs/exps/a9/b32d14966b9be1396f2211d9eb743359708a07/whole-bout',
+      'refs/exps/a9/d8057e088d46842f15c3b6d1bb2e4befd5f677/deism-bots',
+      'refs/exps/a9/d8057e088d46842f15c3b6d1bb2e4befd5f677/inter-gulf',
+      'refs/exps/a9/d8057e088d46842f15c3b6d1bb2e4befd5f677/known-flus',
+      'refs/exps/a9/d8057e088d46842f15c3b6d1bb2e4befd5f677/tight-lira',
+      'refs/exps/bf/6ca8a35911bc6e62fb9bcaa506d4f4e185450c/crumb-orcs'
+    ])
+    expect(remoteExpShas).toStrictEqual(
+      new Set([
+        '263e4408e42a0e215b0f70b36b2ab7b65a160d7e',
+        'd4f2a35773ead55b7ce4b596f600e98360e49372',
+        '5af79e8d5e53f4e41221b6a166121d96d50b630a',
+        '21745a4aa76daf59b49ec81480fe7a89c7ea8fb2',
+        '390aef747f45fc49ec8928b24771f8950d057393',
+        '142a803b83ff784ba1106cc4ad0ba03310da6186',
+        '21ce298cd1743405a0d73f5cb4cf52289ffa3276'
+      ])
+    )
+  })
+})

extension/src/experiments/model/collect.ts

@@ -483,11 +483,15 @@ export const collectRunningInWorkspace = (
   }
 }
 
-export const collectRemoteExpShas = (remoteExpRefs: string): Set<string> => {
+export const collectRemoteExpDetails = (
+  lsRemoteOutput: string
+): { remoteExpShas: Set<string>; remoteExpRefs: string[] } => {
+  const remoteExpRefs: string[] = []
   const remoteExpShas = new Set<string>()
-  for (const ref of trimAndSplit(remoteExpRefs)) {
-    const [sha] = ref.split(/\s/)
+  for (const shaAndRef of trimAndSplit(lsRemoteOutput)) {
+    const [sha, ref] = shaAndRef.split(/\s+/)
     remoteExpShas.add(sha)
+    remoteExpRefs.push(ref.trim())
   }
-  return remoteExpShas
+  return { remoteExpRefs, remoteExpShas }
 }

extension/src/experiments/model/index.ts

@@ -6,7 +6,7 @@ import {
   collectAddRemoveCommitsDetails,
   collectExperiments,
   collectOrderedCommitsAndExperiments,
-  collectRemoteExpShas,
+  collectRemoteExpDetails,
   collectRunningInQueue,
   collectRunningInWorkspace
 } from './collect'
@@ -80,6 +80,7 @@ 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 currentSorts: SortDefinition[]
@@ -174,9 +175,12 @@ export class ExperimentsModel extends ModelWithPersistence {
     this.setColoredStatus(runningExperiments)
   }
 
-  public transformAndSetRemote(remoteExpRefs: string) {
-    const remoteExpShas = collectRemoteExpShas(remoteExpRefs)
+  public transformAndSetRemote(lsRemoteOutput: string) {
+    const { remoteExpShas, remoteExpRefs } =
+      collectRemoteExpDetails(lsRemoteOutput)
     this.remoteExpShas = remoteExpShas
+    this.remoteExpRefs = remoteExpRefs
+    this.deferred.resolve()
   }
 
   public toggleStars(ids: string[]) {
@@ -539,6 +543,10 @@ export class ExperimentsModel extends ModelWithPersistence {
     return this.availableBranchesToSelect
   }
 
+  public getRemoteExpRefs() {
+    return this.remoteExpRefs
+  }
+
   public hasDvcLiveOnlyRunning() {
     return !!this.dvcLiveOnlyExpName
   }

extension/src/experiments/studio.ts

@@ -19,7 +19,10 @@ export class Studio extends Disposable {
     return !!this.baseUrl
   }
 
-  public async setBaseUrl(studioToken: string | undefined) {
+  public async setBaseUrl(
+    studioToken: string | undefined,
+    remoteExpRefs: string[]
+  ) {
     if (!studioToken) {
       this.baseUrl = undefined
       return
@@ -30,7 +33,7 @@ export class Studio extends Disposable {
       this.dvcRoot
     )
 
-    this.baseUrl = await this.getBaseUrl(studioToken, gitRemote)
+    this.baseUrl = await this.getBaseUrl(studioToken, gitRemote, remoteExpRefs)
   }
 
   public getLink(sha: string) {
@@ -40,12 +43,16 @@ export class Studio extends Disposable {
     return `${this.baseUrl}?showOnlySelected=1&experimentReferences=${sha}&activeExperimentReferences=${sha}%3Aprimary`
   }
 
-  private async getBaseUrl(studioToken: string, gitRemoteUrl: string) {
+  private async getBaseUrl(
+    studioToken: string,
+    gitRemoteUrl: string,
+    remoteExpRefs: string[]
+  ) {
     try {
       const response = await fetch(`${STUDIO_URL}/webhook/dvc`, {
         body: JSON.stringify({
           client: 'vscode',
-          refs: [null],
+          refs: { pushed: remoteExpRefs },
           repo_url: gitRemoteUrl
         }),
         headers: {

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

@@ -94,6 +94,7 @@ import { MAX_SELECTED_EXPERIMENTS } from '../../../experiments/model/status'
 import { Pipeline } from '../../../pipeline'
 import { ColumnLike } from '../../../experiments/columns/like'
 import * as Clipboard from '../../../vscode/clipboard'
+import { STUDIO_URL } from '../../../setup/webview/contract'
 
 const { openFileInEditor } = FileSystem
 
@@ -680,20 +681,23 @@ suite('Experiments Test Suite', () => {
     }).timeout(WEBVIEW_TEST_TIMEOUT)
 
     it("should handle a message to copy an experiment's Studio link", async () => {
-      const { mockMessageReceived, experiments } =
+      const { mockMessageReceived, experiments, gitReader } =
         await buildExperimentsWebview({ disposer: disposable })
 
       const viewUrl =
         'https://studio.iterative.ai/user/demo-user/projects/demo-ynm6t3jxdx'
 
-      stub(Fetch, 'default').resolves({
+      const mockGitRemoteUrl = '[email protected]:iterative/vscode-dvc-demo.git'
+      stub(gitReader, 'getRemoteUrl').resolves(mockGitRemoteUrl)
+      const mockStudioToken = 'isat_BBBBBBBBBBBBBBBBBBBBBBBBBBBBBBB'
+      const mockFetch = stub(Fetch, 'default').resolves({
         json: () =>
           Promise.resolve({
             url: viewUrl
           })
       } as unknown as Fetch.Response)
 
-      await experiments.setStudioBaseUrl('isat_BBBBBBBBBBBBBBBBBBBBBBBBBBBBBBB')
+      await experiments.setStudioBaseUrl(mockStudioToken)
       const mockWriteToClipboard = stub(Clipboard, 'writeToClipboard')
       const writeToClipboardCalled = new Promise(resolve =>
         mockWriteToClipboard.callsFake(() => {
@@ -707,6 +711,23 @@ suite('Experiments Test Suite', () => {
         payload: 'exp-e7a67'
       })
 
+      expect(mockFetch).to.be.calledWith(`${STUDIO_URL}/webhook/dvc`, {
+        body: JSON.stringify({
+          client: 'vscode',
+          refs: {
+            pushed: [
+              'refs/exps/a9/b32d14966b9be1396f2211d9eb743359708a07/test-branch'
+            ]
+          },
+          repo_url: mockGitRemoteUrl
+        }),
+        headers: {
+          Authorization: `token ${mockStudioToken}`,
+          'Content-Type': 'application/json'
+        },
+        method: 'POST'
+      })
+
       await writeToClipboardCalled
       const link =
         viewUrl +

extension/src/test/suite/experiments/util.ts

@@ -38,7 +38,7 @@ export const buildExperiments = ({
   dvcRoot = dvcDemoPath,
   expShow = expShowFixture,
   gitLog = gitLogFixture,
-  remoteExpRefs = remoteExpRefsFixture,
+  lsRemoteOutput = remoteExpRefsFixture,
   rowOrder = rowOrderFixture,
   stageList = 'train'
 }: {
@@ -47,7 +47,7 @@ export const buildExperiments = ({
   dvcRoot?: string
   expShow?: ExpShowOutput
   gitLog?: string
-  remoteExpRefs?: string
+  lsRemoteOutput?: string
   rowOrder?: { branch: string; sha: string }[]
   stageList?: string | null
 }) => {
@@ -103,7 +103,7 @@ export const buildExperiments = ({
       gitLog,
       rowOrder
     }),
-    experiments.setState({ remoteExpRefs })
+    experiments.setState({ lsRemoteOutput })
   ])
 
   return {
@@ -144,7 +144,7 @@ export const buildExperimentsWebview = async (inputs: {
   dvcRoot?: string
   expShow?: ExpShowOutput
   gitLog?: string
-  remoteExpRefs?: string
+  lsRemoteOutput?: string
   rowOrder?: { branch: string; sha: string }[]
   stageList?: string | null
 }) => {