Patch Studio API update timinig issue

extension/src/experiments/index.ts

@@ -10,6 +10,7 @@ import omit from 'lodash.omit'
 import { ColumnLike, addStarredToColumns } from './columns/like'
 import { setContextForEditorTitleIcons } from './context'
 import { ExperimentsModel } from './model'
+import { collectRemoteExpShas } from './model/collect'
 import {
   pickExperiment,
   pickExperiments,
@@ -229,9 +230,10 @@ export class Experiments extends BaseRepository<TableData> {
   }
 
   public async unsetPushing(ids: string[]) {
-    await this.update()
+    await Promise.all([this.update(), this.patchStudioApiTimingIssue(ids)])
+
     this.experiments.unsetPushing(ids)
-    return this.notifyChanged()
+    return this.webviewMessages.sendWebviewMessage()
   }
 
   public hasCheckpoints() {
@@ -759,4 +761,56 @@ export class Experiments extends BaseRepository<TableData> {
     const columnLikes = addStarredToColumns(columns)
     return pickColumnToFilter(columnLikes)
   }
+
+  private async patchStudioApiTimingIssue(ids: string[]) {
+    if (!this.studio.isConnected()) {
+      return
+    }
+
+    const [, { lsRemoteOutput }] = await Promise.all([
+      this.waitForStudioUpdate(),
+      this.waitForRemoteUpdate()
+    ])
+
+    const remoteExpShas = collectRemoteExpShas(lsRemoteOutput)
+
+    const shas = []
+    for (const { id, sha } of this.experiments.getExperiments()) {
+      if (ids.includes(id) && sha && remoteExpShas.has(sha)) {
+        shas.push(sha)
+      }
+    }
+
+    this.experiments.assumePushed(shas)
+  }
+
+  private waitForStudioUpdate() {
+    return new Promise(resolve => {
+      const listener = this.dispose.track(
+        this.data.onDidUpdate(data => {
+          if (!isStudioExperimentsOutput(data)) {
+            return
+          }
+          this.dispose.untrack(listener)
+          listener.dispose()
+          resolve(undefined)
+        })
+      )
+    })
+  }
+
+  private waitForRemoteUpdate(): Promise<{ lsRemoteOutput: string }> {
+    return new Promise(resolve => {
+      const listener = this.dispose.track(
+        this.data.onDidUpdate(data => {
+          if (!isRemoteExperimentsOutput(data)) {
+            return
+          }
+          this.dispose.untrack(listener)
+          listener.dispose()
+          resolve(data)
+        })
+      )
+    })
+  }
 }

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

@@ -1,4 +1,4 @@
-import { collectExperiments, collectRemoteExpRefs } from './collect'
+import { collectExperiments, collectRemoteExpShas } from './collect'
 import { generateTestExpShowOutput } from '../../test/util/experiments'
 import { ExpShowOutput } from '../../cli/dvc/contract'
 
@@ -102,7 +102,7 @@ describe('collectExperiments', () => {
   })
 })
 
-describe('collectRemoteExpRefs', () => {
+describe('collectRemoteExpShas', () => {
   it('should parse the git ls-remote output', () => {
     const output = `263e4408e42a0e215b0f70b36b2ab7b65a160d7e        refs/exps/a9/b32d14966b9be1396f2211d9eb743359708a07/vital-taal
     d4f2a35773ead55b7ce4b596f600e98360e49372        refs/exps/a9/b32d14966b9be1396f2211d9eb743359708a07/whole-bout
@@ -111,7 +111,7 @@ describe('collectRemoteExpRefs', () => {
     390aef747f45fc49ec8928b24771f8950d057393        refs/exps/a9/d8057e088d46842f15c3b6d1bb2e4befd5f677/known-flus
     142a803b83ff784ba1106cc4ad0ba03310da6186        refs/exps/a9/d8057e088d46842f15c3b6d1bb2e4befd5f677/tight-lira
     21ce298cd1743405a0d73f5cb4cf52289ffa3276        refs/exps/bf/6ca8a35911bc6e62fb9bcaa506d4f4e185450c/crumb-orcs`
-    const remoteExpShas = collectRemoteExpRefs(output)
+    const remoteExpShas = collectRemoteExpShas(output)
     expect(remoteExpShas).toStrictEqual(
       new Set([
         '263e4408e42a0e215b0f70b36b2ab7b65a160d7e',

extension/src/experiments/model/collect.ts

@@ -485,7 +485,7 @@ export const collectRunningInWorkspace = (
   }
 }
 
-export const collectRemoteExpRefs = (lsRemoteOutput: string): Set<string> => {
+export const collectRemoteExpShas = (lsRemoteOutput: string): Set<string> => {
   const acc = new Set<string>()
   for (const shaAndRef of trimAndSplit(lsRemoteOutput)) {
     const [sha] = shaAndRef.split(/\s+/)

extension/src/experiments/model/index.ts

@@ -6,7 +6,7 @@ import {
   collectAddRemoveCommitsDetails,
   collectExperiments,
   collectOrderedCommitsAndExperiments,
-  collectRemoteExpRefs,
+  collectRemoteExpShas,
   collectRunningInQueue,
   collectRunningInWorkspace
 } from './collect'
@@ -181,7 +181,7 @@ export class ExperimentsModel extends ModelWithPersistence {
   }
 
   public transformAndSetRemote(lsRemoteOutput: string) {
-    const remoteExpShas = collectRemoteExpRefs(lsRemoteOutput)
+    const remoteExpShas = collectRemoteExpShas(lsRemoteOutput)
     this.remoteExpShas = remoteExpShas
   }
 
@@ -563,6 +563,15 @@ export class ExperimentsModel extends ModelWithPersistence {
     this.studioPushedExperiments = pushed
   }
 
+  public assumePushed(shas: string[]) {
+    for (const sha of shas) {
+      if (this.studioPushedExperiments.includes(sha)) {
+        continue
+      }
+      this.studioPushedExperiments.push(sha)
+    }
+  }
+
   public hasDvcLiveOnlyRunning() {
     return !!this.dvcLiveOnlyExpName
   }

extension/src/experiments/studio.ts

@@ -36,6 +36,10 @@ export class Studio extends DeferredDisposable {
     return this.accessTokenSet
   }
 
+  public isConnected() {
+    return !!this.baseUrl
+  }
+
   public getAccessToken() {
     return this.studioAccessToken
   }

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

@@ -13,7 +13,8 @@ import {
   CancellationToken,
   WorkspaceConfiguration,
   MessageItem,
-  ConfigurationTarget
+  ConfigurationTarget,
+  EventEmitter
 } from 'vscode'
 import {
   DEFAULT_EXPERIMENTS_OUTPUT,
@@ -94,6 +95,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 { ExperimentsOutput } from '../../../data'
 
 const { openFileInEditor } = FileSystem
 
@@ -593,8 +595,18 @@ suite('Experiments Test Suite', () => {
     }).timeout(WEBVIEW_TEST_TIMEOUT)
 
     it('should handle a message to push an experiment', async () => {
-      const { experiments, mockMessageReceived } =
-        await stubWorkspaceGettersWebview(disposable)
+      const {
+        experiments,
+        experimentsModel,
+        messageSpy,
+        mockMessageReceived,
+        webview
+      } = await stubWorkspaceGettersWebview(disposable)
+
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      const studio = (experiments as any).studio
+
+      const mockIsConnected = stub(studio, 'isConnected').returns(false)
 
       const mockExpId = 'exp-e7a67'
 
@@ -634,6 +646,15 @@ suite('Experiments Test Suite', () => {
         RegisteredCommands.SETUP_SHOW_STUDIO_CONNECT
       )
 
+      const experimentWithoutLink = experimentsModel
+        .getRowData()[1]
+        .subRows?.find(({ id }) => id === mockExpId)
+
+      expect(experimentWithoutLink?.studioLinkType).not.to.equal(
+        StudioLinkType.PUSHED
+      )
+
+      mockIsConnected.restore()
       mockGetStudioAccessToken.resetBehavior()
 
       const tokenFound = new Promise(resolve =>
@@ -664,6 +685,25 @@ suite('Experiments Test Suite', () => {
         })
       )
 
+      const dataUpdated = disposable.track(
+        new EventEmitter<ExperimentsOutput>()
+      )
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      ;(experiments as any).data.onDidUpdate = dataUpdated.event
+
+      let calls = 0
+
+      const remoteUpdated = new Promise(resolve =>
+        disposable.track(
+          dataUpdated.event(() => {
+            calls = calls + 1
+            if (calls === 2) {
+              resolve(undefined)
+            }
+          })
+        )
+      )
+
       mockMessageReceived.fire({
         payload: [mockExpId],
         type: MessageFromWebviewType.PUSH_EXPERIMENT
@@ -677,6 +717,28 @@ suite('Experiments Test Suite', () => {
         message:
           "Experiment major-lamb is up to date on Git remote 'origin'.\nView your experiments in [Studio](https://studio.iterative.ai/user/mattseddon/projects/vscode-dvc-demo-ynm6t3jxdx)"
       })
+
+      messageSpy.restore()
+      const mockShow = stub(webview, 'show')
+
+      const messageSent = new Promise(resolve =>
+        mockShow.callsFake(() => {
+          resolve(undefined)
+          return Promise.resolve(true)
+        })
+      )
+      dataUpdated.fire({ live: [], pushed: [], baseUrl: '' })
+      dataUpdated.fire({
+        lsRemoteOutput: `42b8736b08170529903cd203a1f40382a4b4a8cd        refs/exps/a9/b32d14966b9be1396f2211d9eb743359708a07/test-branch
+        4fb124aebddb2adf1545030907687fa9a4c80e70        refs/exps/a9/53c3851f46955fa3e2b8f6e1c52999acc8c9ea77/${mockExpId}`
+      })
+      await Promise.all([remoteUpdated, messageSent])
+
+      const experimentWithLink = experimentsModel
+        .getRowData()[1]
+        .subRows?.find(({ id }) => id === mockExpId)
+
+      expect(experimentWithLink?.studioLinkType).to.equal(StudioLinkType.PUSHED)
     }).timeout(WEBVIEW_TEST_TIMEOUT)
 
     it("should handle a message to copy an experiment's Studio link", async () => {

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

@@ -384,7 +384,8 @@ export const stubWorkspaceGettersWebview = async (
     experiments,
     experimentsModel,
     messageSpy,
-    mockMessageReceived
+    mockMessageReceived,
+    webview
   } = await buildExperimentsWebview({ disposer })
 
   return {
@@ -395,6 +396,7 @@ export const stubWorkspaceGettersWebview = async (
     experimentsModel,
     messageSpy,
     ...stubWorkspaceExperiments(dvcRoot, experiments),
-    mockMessageReceived
+    mockMessageReceived,
+    webview
   }
 }