From 1e5430a54a5583a46a723ded0c9b2274959aa7de Mon Sep 17 00:00:00 2001 From: Szymon Pawlak Date: Wed, 20 Dec 2023 17:10:28 +0100 Subject: [PATCH] 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 --- .../sourcegraph/cody/CodyToolWindowContent.kt | 56 ++++++++++++++----- .../ui/SignInWithEnterpriseInstanceAction.kt | 3 +- .../cody/config/BaseLoginDialog.kt | 4 ++ .../sourcegraph/cody/config/CodyAccount.kt | 2 +- .../cody/config/CodyAccountDetails.kt | 8 ++- .../cody/config/CodyAccountDetailsProvider.kt | 2 +- .../cody/config/CodyAccountListModel.kt | 5 +- .../cody/config/CodyAccountsHost.kt | 8 ++- .../cody/config/CodyPersistentAccountsHost.kt | 5 +- .../cody/config/LogInToSourcegraphAction.kt | 6 +- .../cody/config/SettingsMigration.kt | 46 +++++++++++---- .../graphql/query/getUserDetails.graphql | 1 + 12 files changed, 108 insertions(+), 38 deletions(-) diff --git a/src/main/kotlin/com/sourcegraph/cody/CodyToolWindowContent.kt b/src/main/kotlin/com/sourcegraph/cody/CodyToolWindowContent.kt index d40d35d738..353de3ad7c 100644 --- a/src/main/kotlin/com/sourcegraph/cody/CodyToolWindowContent.kt +++ b/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..." diff --git a/src/main/kotlin/com/sourcegraph/cody/auth/ui/SignInWithEnterpriseInstanceAction.kt b/src/main/kotlin/com/sourcegraph/cody/auth/ui/SignInWithEnterpriseInstanceAction.kt index 021614816f..1d77156f1b 100644 --- a/src/main/kotlin/com/sourcegraph/cody/auth/ui/SignInWithEnterpriseInstanceAction.kt +++ b/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) diff --git a/src/main/kotlin/com/sourcegraph/cody/config/BaseLoginDialog.kt b/src/main/kotlin/com/sourcegraph/cody/config/BaseLoginDialog.kt index dbf9dff73e..206098b935 100644 --- a/src/main/kotlin/com/sourcegraph/cody/config/BaseLoginDialog.kt +++ b/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) } diff --git a/src/main/kotlin/com/sourcegraph/cody/config/CodyAccount.kt b/src/main/kotlin/com/sourcegraph/cody/config/CodyAccount.kt index 8970f2a031..b3f6f290ef 100644 --- a/src/main/kotlin/com/sourcegraph/cody/config/CodyAccount.kt +++ b/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 { diff --git a/src/main/kotlin/com/sourcegraph/cody/config/CodyAccountDetails.kt b/src/main/kotlin/com/sourcegraph/cody/config/CodyAccountDetails.kt index 3af2fa299c..d130331536 100644 --- a/src/main/kotlin/com/sourcegraph/cody/config/CodyAccountDetails.kt +++ b/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 } diff --git a/src/main/kotlin/com/sourcegraph/cody/config/CodyAccountDetailsProvider.kt b/src/main/kotlin/com/sourcegraph/cody/config/CodyAccountDetailsProvider.kt index 1f71de6f36..0a7ff83753 100644 --- a/src/main/kotlin/com/sourcegraph/cody/config/CodyAccountDetailsProvider.kt +++ b/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 = diff --git a/src/main/kotlin/com/sourcegraph/cody/config/CodyAccountListModel.kt b/src/main/kotlin/com/sourcegraph/cody/config/CodyAccountListModel.kt index d6d684b09a..a7d5a37d2d 100644 --- a/src/main/kotlin/com/sourcegraph/cody/config/CodyAccountListModel.kt +++ b/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 } diff --git a/src/main/kotlin/com/sourcegraph/cody/config/CodyAccountsHost.kt b/src/main/kotlin/com/sourcegraph/cody/config/CodyAccountsHost.kt index 21218982db..92fe08b79d 100644 --- a/src/main/kotlin/com/sourcegraph/cody/config/CodyAccountsHost.kt +++ b/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 diff --git a/src/main/kotlin/com/sourcegraph/cody/config/CodyPersistentAccountsHost.kt b/src/main/kotlin/com/sourcegraph/cody/config/CodyPersistentAccountsHost.kt index 2a27edb565..56cf51577a 100644 --- a/src/main/kotlin/com/sourcegraph/cody/config/CodyPersistentAccountsHost.kt +++ b/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) diff --git a/src/main/kotlin/com/sourcegraph/cody/config/LogInToSourcegraphAction.kt b/src/main/kotlin/com/sourcegraph/cody/config/LogInToSourcegraphAction.kt index 5338a47715..44617c75e8 100644 --- a/src/main/kotlin/com/sourcegraph/cody/config/LogInToSourcegraphAction.kt +++ b/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) } } } diff --git a/src/main/kotlin/com/sourcegraph/cody/config/SettingsMigration.kt b/src/main/kotlin/com/sourcegraph/cody/config/SettingsMigration.kt index e65da0b2bc..ac4ed5f084 100644 --- a/src/main/kotlin/com/sourcegraph/cody/config/SettingsMigration.kt +++ b/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) diff --git a/src/main/resources/graphql/query/getUserDetails.graphql b/src/main/resources/graphql/query/getUserDetails.graphql index d348ae4bca..9ac4518607 100644 --- a/src/main/resources/graphql/query/getUserDetails.graphql +++ b/src/main/resources/graphql/query/getUserDetails.graphql @@ -1,5 +1,6 @@ query { currentUser { + id username displayName avatarURL