Fix disappearing subscription tab when switching accounts (#208)

It fixes #186 
## Test plan
* Add user with enterprise account
* Add another user with dotcom account
* Switch between enterprise and dotcom accounts
* Subscription tab should appear for dotcom account and disappear for
enterprise account


https://github.com/sourcegraph/jetbrains/assets/9321940/309974c6-fc1a-45a6-912f-ce3cacfab9d7


<!-- All pull requests REQUIRE a test plan:
https://sourcegraph.com/docs/dev/background-information/testing_principles

Why does it matter?

These test plans are there to demonstrate that are following industry
standards which are important or critical for our customers.
They might be read by customers or an auditor. There are meant be simple
and easy to read. Simply explain what you did to ensure
your changes are correct!

Here are a non exhaustive list of test plan examples to help you:

- Making changes on a given feature or component:
- "Covered by existing tests" or "CI" for the shortest possible plan if
there is zero ambiguity
  - "Added new tests"
- "Manually tested" (if non trivial, share some output, logs, or
screenshot)
- Updating docs:
  - "previewed locally"
  - share a screenshot if you want to be thorough
- Updating deps, that would typically fail immediately in CI if
incorrect
  - "CI"
  - "locally tested"
-->

src/main/kotlin/com/sourcegraph/cody/CodyToolWindowContent.kt

@@ -131,26 +131,52 @@ class CodyToolWindowContent(private val project: Project) : UpdatableChat {
 
   private fun addNewSubscriptionTab(server: CodyAgentServer) {
     val activeAccountType = CodyAuthenticationManager.instance.getActiveAccount(project)
-    if (activeAccountType != null && activeAccountType.isDotcomAccount()) {
-      val codyProFeatureFlag = server.evaluateFeatureFlag(GetFeatureFlag("CodyProJetBrains"))
-      if (codyProFeatureFlag.get() != null && codyProFeatureFlag.get()!!) {
-        val isCurrentUserPro =
-            server
-                .isCurrentUserPro()
-                .exceptionally { e ->
-                  logger.warn("Error getting user pro status", e)
-                  null
-                }
-                .get()
-        if (isCurrentUserPro != null) {
-          val subscriptionPanel = createSubscriptionTab(isCurrentUserPro)
-          tabbedPane.insertTab(
-              "Subscription", null, subscriptionPanel, null, SUBSCRIPTION_TAB_INDEX)
+    if (activeAccountType != null) {
+      val jetbrainsUserId = activeAccountType.id
+      var agentUserId = getUserId(server)
+      var retryCount = 3
+      while (jetbrainsUserId != agentUserId && retryCount > 0) {
+        Thread.sleep(200)
+        retryCount--
+        logger.warn("Retrying call for userId from agent")
+        agentUserId = getUserId(server)
+      }
+      if (jetbrainsUserId != agentUserId) {
+        logger.error("User id in JetBrains is different from agent")
+        return
+      }
+
+      if (activeAccountType.isDotcomAccount()) {
+        val codyProFeatureFlag = server.evaluateFeatureFlag(GetFeatureFlag("CodyProJetBrains"))
+        if (codyProFeatureFlag.get() != null && codyProFeatureFlag.get()!!) {
+          val isCurrentUserPro =
+              server
+                  .isCurrentUserPro()
+                  .exceptionally { e ->
+                    logger.warn("Error getting user pro status", e)
+                    null
+                  }
+                  .get()
+          if (isCurrentUserPro != null) {
+            val subscriptionPanel = createSubscriptionTab(isCurrentUserPro)
+            tabbedPane.insertTab(
+                "Subscription", null, subscriptionPanel, null, SUBSCRIPTION_TAB_INDEX)
+          }
         }
       }
     }
   }
 
+  private fun getUserId(server: CodyAgentServer): String? {
+    return server
+        .currentUserId()
+        .exceptionally {
+          logger.warn("Unable to fetch user id from agent")
+          null
+        }
+        .get()
+  }
+
   private fun refreshRecipes() {
     recipesPanel.removeAll()
     recipesPanel.emptyText.text = "Loading commands..."

src/main/kotlin/com/sourcegraph/cody/auth/ui/SignInWithEnterpriseInstanceAction.kt

@@ -23,7 +23,8 @@ class SignInWithEnterpriseInstanceAction(
 
     dialog.setServer(defaultServer)
     if (dialog.showAndGet()) {
-      accountsHost.addAccount(dialog.server, dialog.login, dialog.displayName, dialog.token)
+      accountsHost.addAccount(
+          dialog.server, dialog.login, dialog.displayName, dialog.token, dialog.id)
       if (project != null && ConfigUtil.isCodyEnabled()) {
         // Open Cody sidebar
         val toolWindowManager = ToolWindowManager.getInstance(project)

src/main/kotlin/com/sourcegraph/cody/config/BaseLoginDialog.kt

@@ -25,6 +25,9 @@ abstract class BaseLoginDialog(
 
   protected val loginPanel = CodyLoginPanel(executorFactory, isAccountUnique)
 
+  var id: String = ""
+    private set
+
   var login: String = ""
     private set
 
@@ -71,6 +74,7 @@ abstract class BaseLoginDialog(
           login = details.username
           displayName = details.displayName
           token = newToken
+          id = details.id
 
           close(OK_EXIT_CODE, true)
         }

src/main/kotlin/com/sourcegraph/cody/config/CodyAccount.kt

@@ -21,7 +21,7 @@ data class CodyAccount(
     @Property(style = Property.Style.ATTRIBUTE, surroundWithTag = false)
     override val server: SourcegraphServerPath =
         SourcegraphServerPath.from(ConfigUtil.DOTCOM_URL, ""),
-    @Attribute("id") override val id: String = generateId(),
+    @Attribute("id") override var id: String = generateId(),
 ) : ServerAccount() {
 
   fun isCodyApp(): Boolean {

src/main/kotlin/com/sourcegraph/cody/config/CodyAccountDetails.kt

@@ -2,8 +2,12 @@ package com.sourcegraph.cody.config
 
 import com.sourcegraph.cody.auth.AccountDetails
 
-class CodyAccountDetails(val username: String, val displayName: String?, val avatarURL: String?) :
-    AccountDetails {
+class CodyAccountDetails(
+    val id: String,
+    val username: String,
+    val displayName: String?,
+    val avatarURL: String?
+) : AccountDetails {
   override val name: String
     get() = displayName ?: username
 }

src/main/kotlin/com/sourcegraph/cody/config/CodyAccountDetailsProvider.kt

@@ -29,7 +29,7 @@ class CodyAccountDetailsProvider(
     return ProgressManager.getInstance()
         .submitIOTask(indicator) { indicator ->
           if (account.isCodyApp()) {
-            val details = CodyAccountDetails(account.name, account.name, null)
+            val details = CodyAccountDetails(account.id, account.name, account.name, null)
             DetailsLoadingResult(details, IconUtil.toBufferedImage(defaultIcon), null, false)
           } else {
             val accountDetails =

src/main/kotlin/com/sourcegraph/cody/config/CodyAccountListModel.kt

@@ -69,9 +69,10 @@ class CodyAccountListModel(private val project: Project) :
       server: SourcegraphServerPath,
       login: String,
       displayName: String?,
-      token: String
+      token: String,
+      id: String
   ) {
-    val account = CodyAccount.create(login, displayName, server)
+    val account = CodyAccount.create(login, displayName, server, id)
     if (accountsListModel.isEmpty) {
       activeAccount = account
     }

src/main/kotlin/com/sourcegraph/cody/config/CodyAccountsHost.kt

@@ -4,7 +4,13 @@ import com.intellij.openapi.actionSystem.DataKey
 import com.intellij.openapi.util.Key
 
 interface CodyAccountsHost {
-  fun addAccount(server: SourcegraphServerPath, login: String, displayName: String?, token: String)
+  fun addAccount(
+      server: SourcegraphServerPath,
+      login: String,
+      displayName: String?,
+      token: String,
+      id: String
+  )
 
   fun isAccountUnique(login: String, server: SourcegraphServerPath): Boolean
 

src/main/kotlin/com/sourcegraph/cody/config/CodyPersistentAccountsHost.kt

@@ -12,9 +12,10 @@ class CodyPersistentAccountsHost(private val project: Project?) : CodyAccountsHo
       server: SourcegraphServerPath,
       login: String,
       displayName: String?,
-      token: String
+      token: String,
+      id: String
   ) {
-    val codyAccount = CodyAccount.create(login, displayName, server)
+    val codyAccount = CodyAccount.create(login, displayName, server, id)
     CodyAuthenticationManager.instance.updateAccountToken(codyAccount, token)
     if (project != null) {
       CodyAuthenticationManager.instance.setActiveAccount(project, codyAccount)

src/main/kotlin/com/sourcegraph/cody/config/LogInToSourcegraphAction.kt

@@ -35,7 +35,8 @@ class LogInToSourcegraphAction : BaseAddAccountWithTokenAction() {
             authMethod)
     dialog.setServer(defaultServer)
     if (dialog.showAndGet()) {
-      accountsHost.addAccount(dialog.server, dialog.login, dialog.displayName, dialog.token)
+      accountsHost.addAccount(
+          dialog.server, dialog.login, dialog.displayName, dialog.token, dialog.id)
     }
   }
 }
@@ -54,7 +55,8 @@ class AddCodyEnterpriseAccountAction : BaseAddAccountWithTokenAction() {
 
     dialog.setServer(defaultServer)
     if (dialog.showAndGet()) {
-      accountsHost.addAccount(dialog.server, dialog.login, dialog.displayName, dialog.token)
+      accountsHost.addAccount(
+          dialog.server, dialog.login, dialog.displayName, dialog.token, dialog.id)
     }
   }
 }

src/main/kotlin/com/sourcegraph/cody/config/SettingsMigration.kt

@@ -14,30 +14,57 @@ import com.intellij.openapi.wm.ToolWindowManager
 import com.sourcegraph.cody.CodyToolWindowFactory
 import com.sourcegraph.cody.api.SourcegraphApiRequestExecutor
 import com.sourcegraph.cody.api.SourcegraphApiRequests
-import com.sourcegraph.cody.auth.Account
 import com.sourcegraph.cody.initialization.Activity
 import com.sourcegraph.config.AccessTokenStorage
 import com.sourcegraph.config.CodyApplicationService
 import com.sourcegraph.config.CodyProjectService
 import com.sourcegraph.config.ConfigUtil
 import com.sourcegraph.config.UserLevelConfig
-import java.util.UUID
 import java.util.concurrent.CompletableFuture
 
+private const val RUN_ONCE_CODY_ACCOUNTS_IDS_REFRESH = "CodyAccountsIdsRefresh"
+private const val RUN_ONCE_TOGGLE_CODY_TOOL_WINDOW_AFTER_MIGRATION =
+    "ToggleCodyToolWindowAfterMigration"
+private const val RUN_ONCE_CODY_APPLICATION_SETTINGS_MIGRATION = "CodyApplicationSettingsMigration"
+private const val RUN_ONCE_CODY_PROJECT_SETTINGS_MIGRATION = "CodyProjectSettingsMigration"
+
 class SettingsMigration : Activity {
 
   private val codyAuthenticationManager = CodyAuthenticationManager.instance
 
   override fun runActivity(project: Project) {
-    RunOnceUtil.runOnceForProject(project, "CodyProjectSettingsMigration") {
+    RunOnceUtil.runOnceForProject(project, RUN_ONCE_CODY_PROJECT_SETTINGS_MIGRATION) {
       val customRequestHeaders = extractCustomRequestHeaders(project)
       migrateProjectSettings(project)
       migrateAccounts(project, customRequestHeaders)
     }
-    RunOnceUtil.runOnceForApp("CodyApplicationSettingsMigration") { migrateApplicationSettings() }
-    RunOnceUtil.runOnceForApp("ToggleCodyToolWindowAfterMigration") {
+    RunOnceUtil.runOnceForApp(RUN_ONCE_CODY_APPLICATION_SETTINGS_MIGRATION) {
+      migrateApplicationSettings()
+    }
+    RunOnceUtil.runOnceForApp(RUN_ONCE_TOGGLE_CODY_TOOL_WINDOW_AFTER_MIGRATION) {
       ApplicationManager.getApplication().invokeLater { toggleCodyToolbarWindow(project) }
     }
+    RunOnceUtil.runOnceForApp(RUN_ONCE_CODY_ACCOUNTS_IDS_REFRESH) {
+      val customRequestHeaders = extractCustomRequestHeaders(project)
+      refreshAccountsIds(customRequestHeaders)
+    }
+  }
+
+  private fun refreshAccountsIds(customRequestHeaders: String) {
+    codyAuthenticationManager.getAccounts().forEach { codyAccount ->
+      val server = SourcegraphServerPath.from(codyAccount.server.url, customRequestHeaders)
+      val token = codyAuthenticationManager.getTokenForAccount(codyAccount)
+      if (token != null) {
+        loadUserDetails(
+            SourcegraphApiRequestExecutor.Factory.instance,
+            token,
+            EmptyProgressIndicator(ModalityState.NON_MODAL),
+            server) {
+              codyAccount.id = it.id
+              codyAuthenticationManager.updateAccountToken(codyAccount, token)
+            }
+      }
+    }
   }
 
   private fun toggleCodyToolbarWindow(project: Project) {
@@ -70,8 +97,7 @@ class SettingsMigration : Activity {
             dotcomAccessToken,
             server,
             requestExecutorFactory,
-            EmptyProgressIndicator(ModalityState.NON_MODAL),
-            customRequestHeaders)
+            EmptyProgressIndicator(ModalityState.NON_MODAL))
       } else {
         addAccountIfUnique(
             dotcomAccessToken,
@@ -118,10 +144,9 @@ class SettingsMigration : Activity {
       server: SourcegraphServerPath,
       requestExecutorFactory: SourcegraphApiRequestExecutor.Factory,
       progressIndicator: EmptyProgressIndicator,
-      id: String = UUID.randomUUID().toString(),
   ) {
     loadUserDetails(requestExecutorFactory, accessToken, progressIndicator, server) {
-      addAccount(CodyAccount.create(it.name, it.displayName, server, id), accessToken)
+      addAccount(CodyAccount.create(it.name, it.displayName, server, it.id), accessToken)
     }
   }
 
@@ -131,10 +156,9 @@ class SettingsMigration : Activity {
       server: SourcegraphServerPath,
       requestExecutorFactory: SourcegraphApiRequestExecutor.Factory,
       progressIndicator: EmptyProgressIndicator,
-      id: String = Account.generateId(),
   ) {
     loadUserDetails(requestExecutorFactory, accessToken, progressIndicator, server) {
-      val codyAccount = CodyAccount.create(it.name, it.displayName, server, id)
+      val codyAccount = CodyAccount.create(it.name, it.displayName, server, it.id)
       addAccount(codyAccount, accessToken)
       if (CodyAuthenticationManager.instance.getActiveAccount(project) == null)
           CodyAuthenticationManager.instance.setActiveAccount(project, codyAccount)

src/main/resources/graphql/query/getUserDetails.graphql

@@ -1,5 +1,6 @@
 query {
     currentUser {
+        id
         username
         displayName
         avatarURL