From 3f63f9579b1ea6e8dfb2453ad4d9a76feff31967 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Miko=C5=82aj=20Kondratek?= <mik.kondratek@gmail.com>
Date: Wed, 10 Jan 2024 15:38:49 +0100
Subject: [PATCH] Fix slow op on refreshSubscriptionTab + restart the agent on
 jetbrains and agent's user id mismatch (#284)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

## Test plan
- put a breakpoint in `addNewSubscriptionTab` and set `agentUserId` a
dummy value so it mismatch the `jetbrainsUserId`(that will test the
error handling)
  - internal ide error should not appear
- agent should restart and a new proper matching id should be assigned

- in `addNewSubscriptionTab` increase the timeout & runIde
  - UI should not freeeze 🥶
---
 .../sourcegraph/cody/CodyToolWindowContent.kt | 34 +++++++++++--------
 .../AccountSettingChangeListener.kt           |  5 ++-
 2 files changed, 24 insertions(+), 15 deletions(-)

diff --git a/src/main/kotlin/com/sourcegraph/cody/CodyToolWindowContent.kt b/src/main/kotlin/com/sourcegraph/cody/CodyToolWindowContent.kt
index da2032d034..39309ea4cd 100644
--- a/src/main/kotlin/com/sourcegraph/cody/CodyToolWindowContent.kt
+++ b/src/main/kotlin/com/sourcegraph/cody/CodyToolWindowContent.kt
@@ -4,23 +4,24 @@ import com.intellij.icons.AllIcons
 import com.intellij.ide.ui.laf.darcula.ui.DarculaButtonUI
 import com.intellij.openapi.actionSystem.*
 import com.intellij.openapi.application.ApplicationManager
+import com.intellij.openapi.application.invokeLater
 import com.intellij.openapi.components.Service
 import com.intellij.openapi.diagnostic.Logger
 import com.intellij.openapi.fileEditor.FileEditorManager
 import com.intellij.openapi.project.DumbAwareAction
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.ui.VerticalFlowLayout
-import com.intellij.openapi.util.Disposer
 import com.intellij.ui.SimpleTextAttributes
 import com.intellij.ui.components.JBPanelWithEmptyText
 import com.intellij.ui.components.JBTabbedPane
 import com.intellij.util.IconUtil
+import com.intellij.util.concurrency.annotations.RequiresBackgroundThread
 import com.intellij.util.concurrency.annotations.RequiresEdt
 import com.intellij.util.ui.JBUI
-import com.intellij.vcs.log.runInEdtAsync
 import com.intellij.xml.util.XmlStringUtil
 import com.sourcegraph.cody.agent.CodyAgent.Companion.getInitializedServer
 import com.sourcegraph.cody.agent.CodyAgent.Companion.isConnected
+import com.sourcegraph.cody.agent.CodyAgentManager
 import com.sourcegraph.cody.agent.CodyAgentManager.tryRestartingAgentIfNotRunning
 import com.sourcegraph.cody.agent.CodyAgentServer
 import com.sourcegraph.cody.agent.protocol.*
@@ -111,20 +112,21 @@ class CodyToolWindowContent(private val project: Project) : UpdatableChat {
     refreshPanelsVisibility()
 
     addWelcomeMessage()
-    refreshSubscriptionTab()
+    ApplicationManager.getApplication().executeOnPooledThread { refreshSubscriptionTab() }
     loadNewChatId()
   }
 
+  @RequiresBackgroundThread
   fun refreshSubscriptionTab() {
-    runInEdtAsync(Disposer.newCheckedDisposable()) {
-      tryRestartingAgentIfNotRunning(project)
-      getInitializedServer(project).thenAccept { server ->
-        if (tabbedPane.tabCount < SUBSCRIPTION_TAB_INDEX + 1) {
-          addNewSubscriptionTab(server)
-        } else {
+    tryRestartingAgentIfNotRunning(project)
+    getInitializedServer(project).thenAccept { server ->
+      if (tabbedPane.tabCount < SUBSCRIPTION_TAB_INDEX + 1) {
+        addNewSubscriptionTab(server)
+      } else {
+        ApplicationManager.getApplication().invokeLater {
           tabbedPane.remove(SUBSCRIPTION_TAB_INDEX)
-          addNewSubscriptionTab(server)
         }
+        addNewSubscriptionTab(server)
       }
     }
   }
@@ -163,7 +165,9 @@ class CodyToolWindowContent(private val project: Project) : UpdatableChat {
       }
       if (jetbrainsUserId != agentUserId) {
         if (agentUserId != null) {
-          logger.error("User id in JetBrains is different from agent")
+          logger.warn("User id in JetBrains is different from agent: restarting agent...")
+          CodyAgentManager.restartAgent(project)
+          refreshSubscriptionTab()
           return
         }
         return
@@ -181,9 +185,11 @@ class CodyToolWindowContent(private val project: Project) : UpdatableChat {
                   }
                   .get()
           if (isCurrentUserPro != null) {
-            val subscriptionPanel = createSubscriptionTab(isCurrentUserPro)
-            tabbedPane.insertTab(
-                "Subscription", null, subscriptionPanel, null, SUBSCRIPTION_TAB_INDEX)
+            ApplicationManager.getApplication().invokeLater {
+              val subscriptionPanel = createSubscriptionTab(isCurrentUserPro)
+              tabbedPane.insertTab(
+                  "Subscription", null, subscriptionPanel, null, SUBSCRIPTION_TAB_INDEX)
+            }
           }
         }
       }
diff --git a/src/main/kotlin/com/sourcegraph/cody/config/notification/AccountSettingChangeListener.kt b/src/main/kotlin/com/sourcegraph/cody/config/notification/AccountSettingChangeListener.kt
index fea7ee10d1..4921642434 100644
--- a/src/main/kotlin/com/sourcegraph/cody/config/notification/AccountSettingChangeListener.kt
+++ b/src/main/kotlin/com/sourcegraph/cody/config/notification/AccountSettingChangeListener.kt
@@ -1,5 +1,6 @@
 package com.sourcegraph.cody.config.notification
 
+import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.components.Service
 import com.intellij.openapi.project.Project
 import com.sourcegraph.cody.CodyToolWindowContent
@@ -53,7 +54,9 @@ class AccountSettingChangeListener(project: Project) : ChangeListener(project) {
             UpgradeToCodyProNotification.autocompleteRateLimitError.set(null)
             UpgradeToCodyProNotification.chatRateLimitError.set(null)
             CodyAutocompleteStatusService.resetApplication(project)
-            codyToolWindowContent.refreshSubscriptionTab()
+            ApplicationManager.getApplication().executeOnPooledThread {
+              codyToolWindowContent.refreshSubscriptionTab()
+            }
 
             // Log install events
             if (context.serverUrlChanged) {