From dc2cc3c5a12fb972d4bf6a6a2d3b65c1a6e1d5fa Mon Sep 17 00:00:00 2001
From: Carl-Robert Linnupuu <carlrobertoh@gmail.com>
Date: Thu, 30 Nov 2023 23:28:35 +0200
Subject: [PATCH] fix: UI concurrency issues (run completion events on EDT)

---
 ...WindowCompletionResponseEventListener.java | 87 ++++++++++---------
 .../StandardChatToolWindowTabPanelTest.java   |  2 +-
 2 files changed, 48 insertions(+), 41 deletions(-)

diff --git a/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ToolWindowCompletionResponseEventListener.java b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ToolWindowCompletionResponseEventListener.java
index 33f6ba6a..e51b700d 100644
--- a/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ToolWindowCompletionResponseEventListener.java
+++ b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ToolWindowCompletionResponseEventListener.java
@@ -25,6 +25,7 @@ import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.UUID;
+import javax.swing.SwingUtilities;
 
 abstract class ToolWindowCompletionResponseEventListener implements
     CompletionResponseEventListener {
@@ -80,36 +81,40 @@ abstract class ToolWindowCompletionResponseEventListener implements
 
   @Override
   public void handleError(ErrorDetails error, Throwable ex) {
-    try {
-      if ("insufficient_quota".equals(error.getCode())) {
-        if (SettingsState.getInstance().getSelectedService() == ServiceType.OPENAI) {
-          OpenAISettingsState.getInstance().setOpenAIQuotaExceeded(true);
+    SwingUtilities.invokeLater(() -> {
+      try {
+        if ("insufficient_quota".equals(error.getCode())) {
+          if (SettingsState.getInstance().getSelectedService() == ServiceType.OPENAI) {
+            OpenAISettingsState.getInstance().setOpenAIQuotaExceeded(true);
+          }
+          responseContainer.displayQuotaExceeded();
+        } else {
+          responseContainer.displayError(error.getMessage());
         }
-        responseContainer.displayQuotaExceeded();
-      } else {
-        responseContainer.displayError(error.getMessage());
+      } finally {
+        LOG.error(error.getMessage(), ex);
+        responsePanel.enableActions();
+        stopStreaming(responseContainer);
       }
-    } finally {
-      LOG.error(error.getMessage(), ex);
-      responsePanel.enableActions();
-      stopStreaming(responseContainer);
-    }
+    });
   }
 
   @Override
   public void handleTokensExceeded(Conversation conversation, Message message) {
-    var answer = OverlayUtil.showTokenLimitExceededDialog();
-    if (answer == OK) {
-      TelemetryAction.IDE_ACTION.createActionMessage()
-          .property("action", "DISCARD_TOKEN_LIMIT")
-          .property("model", conversation.getModel())
-          .send();
+    SwingUtilities.invokeLater(() -> {
+      var answer = OverlayUtil.showTokenLimitExceededDialog();
+      if (answer == OK) {
+        TelemetryAction.IDE_ACTION.createActionMessage()
+            .property("action", "DISCARD_TOKEN_LIMIT")
+            .property("model", conversation.getModel())
+            .send();
 
-      conversationService.discardTokenLimits(conversation);
-      handleTokensExceededPolicyAccepted();
-    } else {
-      stopStreaming(responseContainer);
-    }
+        conversationService.discardTokenLimits(conversation);
+        handleTokensExceededPolicyAccepted();
+      } else {
+        stopStreaming(responseContainer);
+      }
+    });
   }
 
   @Override
@@ -118,25 +123,27 @@ abstract class ToolWindowCompletionResponseEventListener implements
       Message message,
       Conversation conversation,
       boolean retry) {
-    try {
-      responsePanel.enableActions();
-      conversationService.saveMessage(fullMessage, message, conversation, retry);
+    conversationService.saveMessage(fullMessage, message, conversation, retry);
 
-      var serpResults = serpResultsMapping.get(message.getId());
-      var containsResults = serpResults != null && !serpResults.isEmpty();
-      if (YouSettingsState.getInstance().isDisplayWebSearchResults() && containsResults) {
-        responseContainer.displaySerpResults(serpResults);
-      }
-
-      if (containsResults) {
-        message.setSerpResults(serpResults);
-      }
-
-      totalTokensPanel.updateUserPromptTokens(userPromptTextArea.getText());
-      totalTokensPanel.updateConversationTokens(conversation);
-    } finally {
-      stopStreaming(responseContainer);
+    var serpResults = serpResultsMapping.get(message.getId());
+    var containsResults = serpResults != null && !serpResults.isEmpty();
+    if (containsResults) {
+      message.setSerpResults(serpResults);
     }
+    var displayResults = YouSettingsState.getInstance().isDisplayWebSearchResults();
+
+    SwingUtilities.invokeLater(() -> {
+      try {
+        responsePanel.enableActions();
+        if (displayResults && containsResults) {
+          responseContainer.displaySerpResults(serpResults);
+        }
+        totalTokensPanel.updateUserPromptTokens(userPromptTextArea.getText());
+        totalTokensPanel.updateConversationTokens(conversation);
+      } finally {
+        stopStreaming(responseContainer);
+      }
+    });
   }
 
   @Override
diff --git a/src/test/java/ee/carlrobert/codegpt/toolwindow/chat/StandardChatToolWindowTabPanelTest.java b/src/test/java/ee/carlrobert/codegpt/toolwindow/chat/StandardChatToolWindowTabPanelTest.java
index f0dd26e1..05d0c465 100644
--- a/src/test/java/ee/carlrobert/codegpt/toolwindow/chat/StandardChatToolWindowTabPanelTest.java
+++ b/src/test/java/ee/carlrobert/codegpt/toolwindow/chat/StandardChatToolWindowTabPanelTest.java
@@ -62,7 +62,7 @@ public class StandardChatToolWindowTabPanelTest extends IntegrationTest {
             "highlightedTokens")
         .containsExactly(
             encodingManager.countTokens(COMPLETION_SYSTEM_PROMPT),
-            encodingManager.countConversationTokens(panel.getConversation()),
+            encodingManager.countTokens(message.getPrompt()),
             0,
             0);
     assertThat(panel.getConversation())