From ffc667d17ed5b8dce723f11dc02fa0163820c3d2 Mon Sep 17 00:00:00 2001
From: Carl-Robert Linnupuu <carlrobertoh@gmail.com>
Date: Sun, 30 Jun 2024 17:02:55 +0300
Subject: [PATCH] fix: edit code popover loading state

---
 .../editor/EditCodeCompletionListener.kt      |  1 +
 .../editor/EditCodeSubmissionHandler.kt       | 42 +++++++++----------
 .../carlrobert/codegpt/ui/EditCodePopover.kt  |  1 +
 3 files changed, 21 insertions(+), 23 deletions(-)

diff --git a/src/main/kotlin/ee/carlrobert/codegpt/actions/editor/EditCodeCompletionListener.kt b/src/main/kotlin/ee/carlrobert/codegpt/actions/editor/EditCodeCompletionListener.kt
index 5ba4f19d..077a3913 100644
--- a/src/main/kotlin/ee/carlrobert/codegpt/actions/editor/EditCodeCompletionListener.kt
+++ b/src/main/kotlin/ee/carlrobert/codegpt/actions/editor/EditCodeCompletionListener.kt
@@ -36,6 +36,7 @@ class EditCodeCompletionListener(
     }
 
     override fun onError(error: ErrorDetails, ex: Throwable) {
+        observableProperties.loading.set(false)
         OverlayUtil.showNotification(
             "Something went wrong while requesting completion. Please try again.",
             NotificationType.ERROR
diff --git a/src/main/kotlin/ee/carlrobert/codegpt/actions/editor/EditCodeSubmissionHandler.kt b/src/main/kotlin/ee/carlrobert/codegpt/actions/editor/EditCodeSubmissionHandler.kt
index dda10942..a98dcf6e 100644
--- a/src/main/kotlin/ee/carlrobert/codegpt/actions/editor/EditCodeSubmissionHandler.kt
+++ b/src/main/kotlin/ee/carlrobert/codegpt/actions/editor/EditCodeSubmissionHandler.kt
@@ -20,32 +20,28 @@ class EditCodeSubmissionHandler(
     private val previousSourceRef = AtomicReference<String?>(null)
 
     suspend fun handleSubmit(userPrompt: String) {
-        try {
-            observableProperties.loading.set(true)
-            observableProperties.submitted.set(true)
+        observableProperties.loading.set(true)
+        observableProperties.submitted.set(true)
 
-            previousSourceRef.getAndSet(editor.document.text)
-            val (selectionTextRange, selectedText) = readAction {
-                editor.selectionModel.run {
-                    Pair(
-                        TextRange(selectionStart, selectionEnd),
-                        editor.selectionModel.selectedText ?: ""
-                    )
-                }
+        previousSourceRef.getAndSet(editor.document.text)
+        val (selectionTextRange, selectedText) = readAction {
+            editor.selectionModel.run {
+                Pair(
+                    TextRange(selectionStart, selectionEnd),
+                    editor.selectionModel.selectedText ?: ""
+                )
             }
-            runInEdt { editor.selectionModel.removeSelection() }
-
-            // TODO: Support other providers
-            CompletionClientProvider.getCodeGPTClient().getChatCompletionAsync(
-                CompletionRequestProvider.buildEditCodeRequest(
-                    "$userPrompt\n\n$selectedText",
-                    service<CodeGPTServiceSettings>().state.chatCompletionSettings.model
-                ),
-                EditCodeCompletionListener(editor, observableProperties, selectionTextRange)
-            )
-        } finally {
-            observableProperties.loading.set(false)
         }
+        runInEdt { editor.selectionModel.removeSelection() }
+
+        // TODO: Support other providers
+        CompletionClientProvider.getCodeGPTClient().getChatCompletionAsync(
+            CompletionRequestProvider.buildEditCodeRequest(
+                "$userPrompt\n\n$selectedText",
+                service<CodeGPTServiceSettings>().state.chatCompletionSettings.model
+            ),
+            EditCodeCompletionListener(editor, observableProperties, selectionTextRange)
+        )
     }
 
     fun handleAccept() {
diff --git a/src/main/kotlin/ee/carlrobert/codegpt/ui/EditCodePopover.kt b/src/main/kotlin/ee/carlrobert/codegpt/ui/EditCodePopover.kt
index 8191d412..934a72ce 100644
--- a/src/main/kotlin/ee/carlrobert/codegpt/ui/EditCodePopover.kt
+++ b/src/main/kotlin/ee/carlrobert/codegpt/ui/EditCodePopover.kt
@@ -93,6 +93,7 @@ class EditCodePopover(private val editor: Editor) {
                     popup.cancel()
                 }
                     .visibleIf(observableProperties.submitted)
+                    .enabledIf(observableProperties.loading.not())
                 cell(AsyncProcessIcon("edit_code_spinner")).visibleIf(observableProperties.loading)
                 link(CodeGPTBundle.get("shared.discard")) {
                     submissionHandler.handleReject()