From 782ac24bc9570543eb666b3f250682a775c5d5c6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Miko=C5=82aj=20Kondratek?= <mik.kondratek@gmail.com>
Date: Thu, 18 Jan 2024 17:57:03 +0100
Subject: [PATCH] Fix missing progress indicator on autocompletion loading
 (#314)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

![image](https://github.com/sourcegraph/jetbrains/assets/19799111/7caa5ca0-f6bd-4543-bfdb-2cee04b06266)

Note it's a π PR ❗

---

This PR fixes the missing progress indicator on the autocompletion
loading (fixes https://github.com/sourcegraph/jetbrains/issues/312).

## Test plan
- play with autocomplete
---
 .../autocomplete/CodyAutocompleteManager.kt   | 56 ++++++++++---------
 1 file changed, 30 insertions(+), 26 deletions(-)

diff --git a/src/main/kotlin/com/sourcegraph/cody/autocomplete/CodyAutocompleteManager.kt b/src/main/kotlin/com/sourcegraph/cody/autocomplete/CodyAutocompleteManager.kt
index 6d53e0b03d..9b41b7859c 100644
--- a/src/main/kotlin/com/sourcegraph/cody/autocomplete/CodyAutocompleteManager.kt
+++ b/src/main/kotlin/com/sourcegraph/cody/autocomplete/CodyAutocompleteManager.kt
@@ -47,6 +47,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
@@ -219,7 +220,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)
 
@@ -228,36 +229,39 @@ 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)
+                ApplicationManager.getApplication().executeOnPooledThread {
+                  CodyToolWindowContent.getInstance(project).refreshSubscriptionTab()
+                }
+                processAutocompleteResult(editor, offset, triggerKind, result, cancellationToken)
               }
-            } else if (result != null && result.items.isNotEmpty()) {
-              UpgradeToCodyProNotification.isFirstRLEOnAutomaticAutocompletionsShown = false
-              UpgradeToCodyProNotification.autocompleteRateLimitError.set(null)
-              ApplicationManager.getApplication().executeOnPooledThread {
-                CodyToolWindowContent.getInstance(project).refreshSubscriptionTab()
-              }
-              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?) {