Fix EDT violations and chat refresh (#460)

## Test plan Full QA before today's release. --------- Co-authored-by: spawlak <spawlak@virtuslab.com> Co-authored-by: Mikołaj Kondratek <mik.kondratek@gmail.com>
sourcegraph · Feb 1, 2024 · b1f3c6a · b1f3c6a
1 parent fe351ce
commit b1f3c6a
Show file tree

Hide file tree

Showing 7 changed files with 38 additions and 29 deletions.
diff --git a/src/main/kotlin/com/sourcegraph/cody/CodyToolWindowContent.kt b/src/main/kotlin/com/sourcegraph/cody/CodyToolWindowContent.kt
@@ -131,6 +131,7 @@ class CodyToolWindowContent(private val project: Project) {
     allContentLayout.show(allContentPanel, MAIN_PANEL)
   }
 
+  @RequiresEdt
   private fun selectHistory(state: ChatState) {
     val session = AgentChatSessionService.getInstance(project).getOrCreateFromState(state)
     switchToChatSession(session)

diff --git a/src/main/kotlin/com/sourcegraph/cody/chat/AgentChatSessionService.kt b/src/main/kotlin/com/sourcegraph/cody/chat/AgentChatSessionService.kt
@@ -3,6 +3,7 @@ package com.sourcegraph.cody.chat
 import com.intellij.openapi.components.Service
 import com.intellij.openapi.components.service
 import com.intellij.openapi.project.Project
+import com.intellij.util.concurrency.annotations.RequiresEdt
 import com.sourcegraph.cody.agent.CodyAgent
 import com.sourcegraph.cody.history.state.ChatState
 import java.util.concurrent.ConcurrentLinkedQueue
@@ -24,6 +25,7 @@ class AgentChatSessionService(private val project: Project) {
     return chatSessions.remove(session)
   }
 
+  @RequiresEdt
   fun getOrCreateFromState(state: ChatState): AgentChatSession {
     val session = chatSessions.find { it.getInternalId() == state.internalId }
     return session ?: AgentChatSession.createFromState(project, state)

diff --git a/src/main/kotlin/com/sourcegraph/cody/commands/ui/CommandsContextMenu.kt b/src/main/kotlin/com/sourcegraph/cody/commands/ui/CommandsContextMenu.kt
@@ -3,6 +3,7 @@ package com.sourcegraph.cody.commands.ui
 import com.intellij.openapi.actionSystem.ActionManager
 import com.intellij.openapi.actionSystem.AnActionEvent
 import com.intellij.openapi.actionSystem.DefaultActionGroup
+import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.fileEditor.FileEditorManager
 import com.intellij.openapi.project.DumbAwareAction
 import com.intellij.openapi.project.Project
@@ -33,8 +34,10 @@ class CommandsContextMenu {
                   CodyEditorFactoryListener.Util.informAgentAboutEditorChange(
                       it, hasFileChanged = false) {
                         CodyToolWindowContent.executeOnInstanceIfNotDisposed(project) {
-                          switchToChatSession(
-                              AgentChatSession.createFromCommand(project, commandId))
+                          ApplicationManager.getApplication().invokeLater {
+                            switchToChatSession(
+                                AgentChatSession.createFromCommand(project, commandId))
+                          }
                         }
                       }
                 }

diff --git a/src/main/kotlin/com/sourcegraph/cody/commands/ui/CommandsTabPanel.kt b/src/main/kotlin/com/sourcegraph/cody/commands/ui/CommandsTabPanel.kt
@@ -1,6 +1,7 @@
 package com.sourcegraph.cody.commands.ui
 
 import com.intellij.ide.ui.laf.darcula.ui.DarculaButtonUI
+import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.fileEditor.FileEditorManager
 import com.intellij.openapi.project.Project
 import com.intellij.ui.components.JBPanelWithEmptyText
@@ -26,7 +27,7 @@ class CommandsTabPanel(
   private fun executeCommandWithContext(commandId: CommandId) {
     FileEditorManager.getInstance(project).selectedTextEditor?.let {
       CodyEditorFactoryListener.Util.informAgentAboutEditorChange(it, hasFileChanged = false) {
-        executeCommand(commandId)
+        ApplicationManager.getApplication().invokeLater { executeCommand(commandId) }
       }
     }
   }

diff --git a/src/main/kotlin/com/sourcegraph/cody/history/HistoryService.kt b/src/main/kotlin/com/sourcegraph/cody/history/HistoryService.kt
@@ -21,7 +21,9 @@ class HistoryService : SimplePersistentStateComponent<HistoryState>(HistoryState
   fun update(internalId: String, chatMessages: List<ChatMessage>) {
     val found = getChatOrCreate(internalId)
     found.messages = chatMessages.map(::convertToMessageState).toMutableList()
-    found.setUpdatedTimeAt(LocalDateTime.now())
+    if (chatMessages.lastOrNull()?.speaker == Speaker.HUMAN) {
+      found.setUpdatedTimeAt(LocalDateTime.now())
+    }
     synchronized(listeners) { listeners.forEach { it(found) } }
   }
 

diff --git a/src/main/kotlin/com/sourcegraph/cody/history/HistoryTree.kt b/src/main/kotlin/com/sourcegraph/cody/history/HistoryTree.kt
@@ -76,30 +76,35 @@ class HistoryTree(
         addChatToPeriodAndSort(period, chat)
       }
     } else {
-      val period =
+      val currentPeriodText = DurationGroupFormatter.format(chat.getUpdatedTimeAt())
+      val currentPeriod = root.periods().find { it.periodText == currentPeriodText }
+      val leafWithChangedPeriod =
           root
               .periods()
+              .filter { it.periodText != currentPeriodText }
               .flatMap { it.leafs() }
-              .find { it.chat.internalId == chat.internalId }!!
-              .parent as PeriodNode
-      val periodText = DurationGroupFormatter.format(chat.getUpdatedTimeAt())
-      //      Checking if chat needs to be moved to different period
-      if (period.periodText != periodText) {
-        val filter = period.leafs().filter { it.chat.internalId != chat.internalId }
-        if (filter.isEmpty()) {
-          model.removeNodeFromParent(period)
-        } else {
-          period.removeAllChildren()
-          for (child in filter) period.add(child)
+              .find { it.chat.internalId == chat.internalId }
+
+      if (leafWithChangedPeriod != null) {
+        val previousPeriod = leafWithChangedPeriod.parent as? PeriodNode
+        previousPeriod?.let { period ->
+          period.remove(leafWithChangedPeriod)
+          if (period.childCount == 0) period.removeFromParent()
+          model.reload(period)
+        }
+        currentPeriod?.let { period ->
+          addChatToPeriodAndSort(period, chat)
           model.reload(period)
         }
-        val per = root.periods().find { it.periodText == periodText }
-        addChatToPeriodAndSort(per!!, chat)
       } else {
-        val sorted = period.leafs().sortedByDescending { it.chat.getUpdatedTimeAt() }
-        period.removeAllChildren()
-        for (child in sorted) period.add(child)
-        model.reload(period)
+        currentPeriod?.let { period ->
+          val sorted = period.leafs().sortedByDescending { it.chat.getUpdatedTimeAt() }
+          if (period.leafs() != sorted) {
+            period.removeAllChildren()
+            for (child in sorted) period.add(child)
+            model.reload(period)
+          }
+        }
       }
     }
   }

diff --git a/src/main/kotlin/com/sourcegraph/cody/statusbar/CodyAutocompleteStatusService.kt b/src/main/kotlin/com/sourcegraph/cody/statusbar/CodyAutocompleteStatusService.kt
@@ -5,6 +5,7 @@ import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.components.Service
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.project.ProjectManager
+import com.intellij.util.ui.UIUtil
 import com.sourcegraph.cody.agent.CodyAgentService
 import com.sourcegraph.cody.config.CodyAccountManager
 import com.sourcegraph.cody.config.CodyAuthenticationManager
@@ -75,18 +76,12 @@ class CodyAutocompleteStatusService : CodyAutocompleteStatusListener, Disposable
   }
 
   private fun updateCodyStatusBarIcons() {
-    val action = Runnable {
+    UIUtil.invokeLaterIfNeeded {
       val openProjects = ProjectManager.getInstance().openProjects
       openProjects.forEach { project ->
         project.takeIf { !it.isDisposed }?.let { CodyStatusBarWidget.update(it) }
       }
     }
-    val application = ApplicationManager.getApplication()
-    if (application.isDispatchThread) {
-      action.run()
-    } else {
-      application.invokeLater(action)
-    }
   }
 
   private fun getStatus(): CodyAutocompleteStatus {