diff --git a/src/main/kotlin/com/sourcegraph/cody/autocomplete/CodyAutocompleteManager.kt b/src/main/kotlin/com/sourcegraph/cody/autocomplete/CodyAutocompleteManager.kt index ea517bc62c..fd91bc5297 100644 --- 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() + val resultOuter = CompletableFuture() 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?) {