Fix missing progress indicator on autocompletion loading

sourcegraph · Jan 18, 2024 · de69354 · de69354
1 parent 1a8123a
commit de69354
Showing 1 changed file with 27 additions and 23 deletions.
diff --git a/src/main/kotlin/com/sourcegraph/cody/autocomplete/CodyAutocompleteManager.kt b/src/main/kotlin/com/sourcegraph/cody/autocomplete/CodyAutocompleteManager.kt
@@ -46,6 +46,7 @@ import difflib.Patch
 import java.util.concurrent.CancellationException
 import java.util.concurrent.CompletableFuture
 import java.util.concurrent.CompletionException
+import java.util.concurrent.TimeUnit
 import java.util.concurrent.atomic.AtomicReference
 import java.util.stream.Collectors
 import org.eclipse.lsp4j.jsonrpc.ResponseErrorException
@@ -218,7 +219,7 @@ class CodyAutocompleteManager {
                             position)))
     notifyApplication(CodyAutocompleteStatus.AutocompleteInProgress)
 
-    val result = CompletableFuture<Void?>()
+    val resultOuter = CompletableFuture<Void?>()
     CodyAgentService.applyAgentOnBackgroundThread(project) { agent ->
       val completions = agent.server.autocompleteExecute(params)
 
@@ -227,33 +228,36 @@ class CodyAutocompleteManager {
       // correctly propagate the cancellation to the agent.
       cancellationToken.onCancellationRequested { completions.cancel(true) }
 
-      completions
-          .handle { result, error ->
-            if (error != null) {
-              if (triggerKind == InlineCompletionTriggerKind.INVOKE ||
-                  !UpgradeToCodyProNotification.isFirstRLEOnAutomaticAutocompletionsShown) {
-                handleError(project, error)
+      ApplicationManager.getApplication().executeOnPooledThread {
+        completions
+            .handle { result, error ->
+              if (error != null) {
+                if (triggerKind == InlineCompletionTriggerKind.INVOKE ||
+                    !UpgradeToCodyProNotification.isFirstRLEOnAutomaticAutocompletionsShown) {
+                  handleError(project, error)
+                }
+              } else if (result != null && result.items.isNotEmpty()) {
+                UpgradeToCodyProNotification.isFirstRLEOnAutomaticAutocompletionsShown = false
+                UpgradeToCodyProNotification.autocompleteRateLimitError.set(null)
+                processAutocompleteResult(editor, offset, triggerKind, result, cancellationToken)
               }
-            } else if (result != null && result.items.isNotEmpty()) {
-              UpgradeToCodyProNotification.isFirstRLEOnAutomaticAutocompletionsShown = false
-              UpgradeToCodyProNotification.autocompleteRateLimitError.set(null)
-              processAutocompleteResult(editor, offset, triggerKind, result, cancellationToken)
+              null
             }
-            null
-          }
-          .exceptionally { error: Throwable? ->
-            if (!(error is CancellationException || error is CompletionException)) {
-              logger.warn("failed autocomplete request $params", error)
+            .exceptionally { error: Throwable? ->
+              if (!(error is CancellationException || error is CompletionException)) {
+                logger.warn("failed autocomplete request $params", error)
+              }
+              null
             }
-            null
-          }
-          .thenAccept {
-            resetApplication(project)
-            result.complete(null)
-          }
+            .completeOnTimeout(null, 3, TimeUnit.SECONDS)
+            .get()
+
+        resetApplication(project)
+        resultOuter.complete(null)
+      }
     }
 
-    return result
+    return resultOuter
   }
 
   private fun handleError(project: Project, error: Throwable?) {