Fix internal error on quick project close & open

src/main/java/com/sourcegraph/cody/SubscriptionTabUtils.kt

@@ -7,6 +7,7 @@ import com.sourcegraph.cody.agent.CodyAgentServer
 import com.sourcegraph.cody.agent.protocol.GetFeatureFlag
 import com.sourcegraph.cody.config.CodyAuthenticationManager
 import java.util.concurrent.CompletableFuture
+import java.util.concurrent.TimeUnit
 
 data class SubscriptionTabPanelData(
     val isDotcomAccount: Boolean,
@@ -28,7 +29,10 @@ fun fetchSubscriptionPanelData(
     if (activeAccountType.isDotcomAccount()) {
       ApplicationManager.getApplication().executeOnPooledThread {
         val codyProFeatureFlag =
-            server.evaluateFeatureFlag(GetFeatureFlag.CodyProJetBrains).get() == true
+            server
+                .evaluateFeatureFlag(GetFeatureFlag.CodyProJetBrains)
+                .completeOnTimeout(false, 4, TimeUnit.SECONDS)
+                .get() == true
         if (codyProFeatureFlag) {
           val isCurrentUserPro = getIsCurrentUserPro(server) ?: false
           result.complete(

src/main/java/com/sourcegraph/vcs/RepoUtil.java

@@ -159,7 +159,7 @@ private static String doReplacements(
                 .thenCompose(
                     agent ->
                         agent.getServer().convertGitCloneURLToCodebaseName(new CloneURL(cloneURL)))
-                .completeOnTimeout(/* value= */ null, /* timeout= */ 4, TimeUnit.SECONDS)
+                .completeOnTimeout(null, 4, TimeUnit.SECONDS)
                 .get();
 
         if (codebaseName == null) {

src/main/kotlin/com/sourcegraph/cody/agent/CodyAgentService.kt

@@ -23,12 +23,16 @@ class CodyAgentService(project: Project) : Disposable {
   init {
     onStartup { agent ->
       agent.client.onNewMessage = Consumer { params ->
-        AgentChatSessionService.getInstance(project)
-            .getSession(params.id)
-            ?.receiveMessage(params.message)
+        if (!project.isDisposed) {
+          AgentChatSessionService.getInstance(project)
+              .getSession(params.id)
+              ?.receiveMessage(params.message)
+        }
       }
 
-      AgentChatSessionService.getInstance(project).restoreAllSessions(agent)
+      if (!project.isDisposed) {
+        AgentChatSessionService.getInstance(project).restoreAllSessions(agent)
+      }
     }
   }
 

src/main/kotlin/com/sourcegraph/cody/initialization/EndOfTrialNotificationScheduler.kt

@@ -38,7 +38,10 @@ class EndOfTrialNotificationScheduler private constructor(val project: Project)
 
           CodyAgentService.applyAgentOnBackgroundThread(project) { agent ->
             val currentUserCodySubscription =
-                agent.server.getCurrentUserCodySubscription().get(4, TimeUnit.SECONDS)
+                agent.server
+                    .getCurrentUserCodySubscription()
+                    .completeOnTimeout(null, 4, TimeUnit.SECONDS)
+                    .get()
 
             if (currentUserCodySubscription == null) {
               logger.debug("currentUserCodySubscription is null")
@@ -48,12 +51,15 @@ class EndOfTrialNotificationScheduler private constructor(val project: Project)
             val codyProTrialEnded =
                 agent.server
                     .evaluateFeatureFlag(GetFeatureFlag.CodyProTrialEnded)
-                    .get(4, TimeUnit.SECONDS) == true
+                    .completeOnTimeout(false, 4, TimeUnit.SECONDS)
+                    .get() == true
 
             val useSscForCodySubscription =
                 agent.server
                     .evaluateFeatureFlag(GetFeatureFlag.UseSscForCodySubscription)
-                    .get(4, TimeUnit.SECONDS) == true
+                    .orTimeout(4, TimeUnit.SECONDS)
+                    .completeOnTimeout(false, 4, TimeUnit.SECONDS)
+                    .get() == true
 
             showProperNotificationIfApplicable(
                 currentUserCodySubscription, codyProTrialEnded, useSscForCodySubscription)

src/main/kotlin/com/sourcegraph/cody/statusbar/CodyAutocompleteStatusService.kt

@@ -38,9 +38,11 @@ class CodyAutocompleteStatusService : CodyAutocompleteStatusListener, Disposable
 
   override fun onCodyAutocompleteStatusReset(project: Project) {
     ApplicationManager.getApplication().executeOnPooledThread {
-      val notify = didStatusChanged(project)
-      if (notify) {
-        updateCodyStatusBarIcons()
+      if (!project.isDisposed) {
+        val notify = didStatusChanged(project)
+        if (notify) {
+          updateCodyStatusBarIcons()
+        }
       }
     }
   }

src/main/kotlin/com/sourcegraph/common/UpgradeToCodyProNotification.kt

@@ -5,7 +5,6 @@ import com.intellij.notification.NotificationType
 import com.intellij.notification.impl.NotificationFullContent
 import com.intellij.openapi.actionSystem.AnAction
 import com.intellij.openapi.actionSystem.AnActionEvent
-import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.project.DumbAwareAction
 import com.intellij.openapi.project.Project
 import com.sourcegraph.Icons
@@ -77,9 +76,7 @@ private constructor(title: String, content: String, shouldShowUpgradeOption: Boo
               else -> CodyBundle.getString("UpgradeToCodyProNotification.title.explain")
             }
 
-        ApplicationManager.getApplication().invokeLater {
-          UpgradeToCodyProNotification(title, content, shouldShowUpgradeOption).notify(project)
-        }
+        UpgradeToCodyProNotification(title, content, shouldShowUpgradeOption).notify(project)
       }
     }