Several bug fixes for Inline Edits

src/main/java/com/sourcegraph/cody/Icons.java

@@ -49,6 +49,10 @@ interface Chat {
     Icon Download = IconLoader.getIcon("/icons/chat/download.svg", Icons.class);
   }
 
+  interface Edit {
+    Icon Error = IconLoader.getIcon("/icons/edit/error.svg", Icons.class);
+  }
+
   interface LLM {
     Icon Anthropic = IconLoader.getIcon("/icons/chat/llm/anthropic.svg", Icons.class);
     Icon OpenAI = IconLoader.getIcon("/icons/chat/llm/openai.svg", Icons.class);

src/main/kotlin/com/sourcegraph/cody/edit/EditCommandPrompt.kt

@@ -17,6 +17,7 @@ import com.intellij.openapi.fileEditor.FileEditorManagerListener
 import com.intellij.openapi.keymap.KeymapUtil
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.roots.ProjectRootManager
+import com.intellij.openapi.util.Disposer
 import com.intellij.openapi.util.SystemInfo
 import com.intellij.openapi.vfs.VfsUtilCore
 import com.intellij.openapi.vfs.VirtualFile
@@ -185,6 +186,28 @@ class EditCommandPrompt(val controller: FixupService, val editor: Editor, dialog
         foreground = mutedLabelColor()
       }
 
+  private val textFieldListener =
+      object : DocumentListener {
+        override fun insertUpdate(e: DocumentEvent?) {
+          handleDocumentChange()
+        }
+
+        override fun removeUpdate(e: DocumentEvent?) {
+          handleDocumentChange()
+        }
+
+        override fun changedUpdate(e: DocumentEvent?) {
+          handleDocumentChange()
+        }
+
+        private fun handleDocumentChange() {
+          ApplicationManager.getApplication().invokeLater {
+            updateOkButtonState()
+            checkForInterruptions()
+          }
+        }
+      }
+
   private val documentListener =
       object : BulkAwareDocumentListener {
         override fun documentChanged(event: com.intellij.openapi.editor.event.DocumentEvent) {
@@ -231,6 +254,9 @@ class EditCommandPrompt(val controller: FixupService, val editor: Editor, dialog
 
   // Note: Must be created on EDT, although we can't annotate it as such.
   init {
+    // Register with FixupService as a failsafe if the project closes. Normally we're disposed
+    // sooner, when the dialog is closed or focus is lost.
+    Disposer.register(controller, this)
     connection = ApplicationManager.getApplication().messageBus.connect(this)
     registerListeners()
     // Don't reset the session, just any previous instructions dialog.
@@ -298,6 +324,7 @@ class EditCommandPrompt(val controller: FixupService, val editor: Editor, dialog
   private fun unregisterListeners() {
     try {
       editor.document.removeDocumentListener(documentListener)
+      instructionsField.document.removeDocumentListener(textFieldListener)
 
       removeWindowFocusListener(windowFocusListener)
       removeFocusListener(focusListener)
@@ -318,27 +345,7 @@ class EditCommandPrompt(val controller: FixupService, val editor: Editor, dialog
 
   @RequiresEdt
   private fun setupTextField() {
-    instructionsField.document.addDocumentListener(
-        object : DocumentListener {
-          override fun insertUpdate(e: DocumentEvent?) {
-            handleDocumentChange()
-          }
-
-          override fun removeUpdate(e: DocumentEvent?) {
-            handleDocumentChange()
-          }
-
-          override fun changedUpdate(e: DocumentEvent?) {
-            handleDocumentChange()
-          }
-
-          private fun handleDocumentChange() {
-            ApplicationManager.getApplication().invokeLater {
-              updateOkButtonState()
-              checkForInterruptions()
-            }
-          }
-        })
+    instructionsField.document.addDocumentListener(textFieldListener)
   }
 
   @RequiresEdt
@@ -393,9 +400,10 @@ class EditCommandPrompt(val controller: FixupService, val editor: Editor, dialog
       instructionsField.text?.let { lastPrompt = it } // Save last thing they typed.
       connection?.disconnect()
       connection = null
-      dispose()
     } catch (x: Exception) {
       logger.warn("Error cancelling edit command prompt", x)
+    } finally {
+      dispose()
     }
   }
 

src/main/kotlin/com/sourcegraph/cody/edit/FixupService.kt

@@ -89,6 +89,13 @@ class FixupService(val project: Project) : Disposable {
         logger.warn("Error disposing session", x)
       }
     }
+    currentEditPrompt.get()?.let {
+      try {
+        Disposer.dispose(it)
+      } catch (x: Exception) {
+        logger.warn("Error disposing prompt", x)
+      }
+    }
   }
 
   companion object {

src/main/kotlin/com/sourcegraph/cody/edit/sessions/FixupSession.kt

@@ -71,25 +71,25 @@ abstract class FixupSession(
   private val performedActions: MutableList<FixupUndoableAction> = mutableListOf()
 
   init {
-    triggerDocumentCodeAsync()
-    // Kotlin doesn't like leaking 'this' before constructors are finished.
+    triggerFixupAsync()
     ApplicationManager.getApplication().invokeLater { Disposer.register(controller, this) }
   }
 
   private val document
     get() = editor.document
 
   @RequiresEdt
-  private fun triggerDocumentCodeAsync() {
+  private fun triggerFixupAsync() {
     // Those lookups require us to be on the EDT.
     val file = FileDocumentManager.getInstance().getFile(document)
     val textFile = file?.let { ProtocolTextDocument.fromVirtualFile(editor, it) } ?: return
 
     CodyAgentService.withAgent(project) { agent ->
       workAroundUninitializedCodebase()
 
-      // Force a round-trip to get folding ranges before showing lenses.
+      // Spend a turn to get folding ranges before showing lenses.
       ensureSelectionRange(agent, textFile)
+
       showWorkingGroup()
 
       // All this because we can get the workspace/edit before the request returns!
@@ -108,7 +108,7 @@ abstract class FixupSession(
           }
           .exceptionally { error: Throwable? ->
             if (!(error is CancellationException || error is CompletionException)) {
-              logger.warn("Error while generating doc string: $error")
+              showErrorGroup("Error while generating code: ${error?.localizedMessage}")
             }
             finish()
             null
@@ -208,10 +208,12 @@ abstract class FixupSession(
     } catch (x: Exception) {
       logger.debug("Session cleanup error", x)
     }
-    try {
-      Disposer.dispose(this)
-    } catch (x: Exception) {
-      logger.warn("Error disposing fixup session $this", x)
+    ApplicationManager.getApplication().invokeLater {
+      try { // Disposing inlay requires EDT.
+        Disposer.dispose(this)
+      } catch (x: Exception) {
+        logger.warn("Error disposing fixup session $this", x)
+      }
     }
   }
 

src/main/kotlin/com/sourcegraph/cody/edit/widget/LensGroupFactory.kt

@@ -4,6 +4,7 @@ import com.intellij.openapi.diagnostic.Logger
 import com.sourcegraph.cody.Icons
 import com.sourcegraph.cody.edit.EditCommandPrompt
 import com.sourcegraph.cody.edit.sessions.FixupSession
+import javax.swing.Icon
 
 /** Handles assembling standard groups of lenses. */
 class LensGroupFactory(val session: FixupSession) {
@@ -66,8 +67,7 @@ class LensGroupFactory(val session: FixupSession) {
   }
 
   private fun addLogo(group: LensWidgetGroup) {
-    group.addWidget(LensIcon(group, Icons.StatusBar.CodyAvailable))
-    addSpacer(group)
+    addIcon(group, Icons.StatusBar.CodyAvailable)
   }
 
   private fun addSpacer(group: LensWidgetGroup) {
@@ -84,7 +84,11 @@ class LensGroupFactory(val session: FixupSession) {
   }
 
   private fun addErrorIcon(group: LensWidgetGroup) {
-    addLabel(group, " ! ") // TODO: Change to LensIcon when we get SVG
+    addIcon(group, Icons.Edit.Error)
+  }
+
+  private fun addIcon(group: LensWidgetGroup, icon: Icon) {
+    group.addWidget(LensIcon(group, icon))
     addSpacer(group)
   }
 

src/main/kotlin/com/sourcegraph/cody/edit/widget/LensIcon.kt

@@ -11,8 +11,8 @@ class LensIcon(group: LensWidgetGroup, val icon: Icon) : LensWidget(group) {
 
   private var scaledImage: Image? = null
 
-  // Squish the icon down to make it better fit the text. Eyeballed based on logo image.
-  private val scaleFactor = 0.6
+  // Squish the icon down to make it better fit the text.
+  private val scaleFactor = 0.9
 
   private fun getScaleFactor(fontMetrics: FontMetrics) = (fontMetrics.height * scaleFactor).toInt()
 

src/main/kotlin/com/sourcegraph/cody/ui/FrameMover.kt

@@ -133,6 +133,9 @@ class FrameMover(private val frame: JFrame, private val titleBar: JComponent) :
   }
 
   private fun resizeDialog(e: MouseEvent) {
+    val currentTime = System.currentTimeMillis()
+    if (currentTime - lastUpdateTime <= 16) return
+
     val newX = e.xOnScreen
     val newY = e.yOnScreen
     val deltaX = newX - lastMouseX
@@ -173,9 +176,10 @@ class FrameMover(private val frame: JFrame, private val titleBar: JComponent) :
       else -> {}
     }
 
-    frame.setSize(newWidth, newHeight)
+    SwingUtilities.invokeLater { frame.setSize(newWidth, newHeight) }
     lastMouseX = newX
     lastMouseY = newY
+    lastUpdateTime = currentTime
   }
 
   private fun moveDialog(e: MouseEvent) {