Block further inline edits until previous one completes

src/main/kotlin/com/sourcegraph/cody/agent/CodyAgentService.kt

@@ -56,7 +56,7 @@ class CodyAgentService(project: Project) : Disposable {
 
       agent.client.onEditTaskDidDelete = Consumer { params ->
         FixupService.getInstance(project).getActiveSession()?.let {
-          if (params.id == it.taskId) it.taskDeleted()
+          if (params.id == it.taskId) it.dismiss()
         }
       }
 

src/main/kotlin/com/sourcegraph/cody/commands/CommandId.kt

@@ -2,8 +2,8 @@ package com.sourcegraph.cody.commands
 
 import java.awt.event.KeyEvent
 
-enum class CommandId(val displayName: String, val mnemonic: Int) {
-  Explain("Explain Code", KeyEvent.VK_E),
-  Smell("Smell Code", KeyEvent.VK_S),
-  Test("Generate Test", KeyEvent.VK_T)
+enum class CommandId(val id: String, val displayName: String, val mnemonic: Int) {
+  Explain("cody.command.Explain", "Explain Code", KeyEvent.VK_E),
+  Smell("cody.command.Smell", "Smell Code", KeyEvent.VK_S),
+  Test("cody.command.Test", "Generate Test", KeyEvent.VK_T)
 }

src/main/kotlin/com/sourcegraph/cody/commands/ui/CommandsTabPanel.kt

@@ -13,6 +13,10 @@ import com.intellij.openapi.project.Project
 import com.intellij.ui.components.JBPanelWithEmptyText
 import com.sourcegraph.cody.commands.CommandId
 import com.sourcegraph.cody.config.CodyApplicationSettings
+import com.sourcegraph.cody.edit.FixupService
+import com.sourcegraph.cody.edit.actions.DocumentCodeAction
+import com.sourcegraph.cody.edit.actions.EditCodeAction
+import com.sourcegraph.cody.edit.actions.TestCodeAction
 import com.sourcegraph.cody.ignore.CommandPanelIgnoreBanner
 import com.sourcegraph.cody.ignore.IgnoreOracle
 import com.sourcegraph.cody.ignore.IgnorePolicy
@@ -22,17 +26,21 @@ import java.awt.Dimension
 import java.awt.GridLayout
 import javax.swing.BoxLayout
 import javax.swing.JButton
-import javax.swing.JComponent
 import javax.swing.plaf.ButtonUI
 
 class CommandsTabPanel(
     private val project: Project,
     private val executeCommand: (CommandId) -> Unit
 ) :
     JBPanelWithEmptyText(GridLayout(/* rows = */ 0, /* cols = */ 1)),
-    IgnoreOracle.FocusedFileIgnorePolicyListener {
+    IgnoreOracle.FocusedFileIgnorePolicyListener,
+    FixupService.ActiveFixupSessionStateListener {
   private val ignoreBanner = CommandPanelIgnoreBanner()
-  private val buttons = mutableListOf<JComponent>()
+  private val buttons = mutableMapOf<String, JButton>()
+
+  @Volatile private var ignorePolicy: IgnorePolicy = IgnorePolicy.USE
+
+  @Volatile private var isInlineEditInProgress: Boolean = false
 
   init {
     layout = BoxLayout(this, BoxLayout.Y_AXIS)
@@ -41,25 +49,18 @@ class CommandsTabPanel(
 
     if (ConfigUtil.isFeatureFlagEnabled("cody.feature.inline-edits") ||
         CodyApplicationSettings.instance.isInlineEditionEnabled) {
-      addInlineEditActionButton("cody.editCodeAction")
-      addInlineEditActionButton("cody.documentCodeAction")
-      addInlineEditActionButton("cody.testCodeAction")
+      addInlineEditActionButton(EditCodeAction.ID)
+      addInlineEditActionButton(DocumentCodeAction.ID)
+      addInlineEditActionButton(TestCodeAction.ID)
     }
 
-    addHierarchyListener {
-      if (!project.isDisposed) {
-        if (it.component.isShowing) {
-          IgnoreOracle.getInstance(project).addListener(this)
-        } else {
-          IgnoreOracle.getInstance(project).removeListener(this)
-        }
-      }
-    }
+    IgnoreOracle.getInstance(project).addListener(this)
+    FixupService.getInstance(project).addListener(this)
   }
 
   private fun addInlineEditActionButton(actionId: String) {
     val action = ActionManagerEx.getInstanceEx().getAction(actionId)
-    addButton(action.templatePresentation.text, action.templatePresentation.mnemonic) {
+    addButton(actionId, action.templatePresentation.text, action.templatePresentation.mnemonic) {
       val editor = FileEditorManager.getInstance(project).selectedTextEditor ?: return@addButton
       val dataContext = (editor as? EditorEx)?.dataContext ?: return@addButton
       val managerEx = ActionManagerEx.getInstanceEx()
@@ -70,10 +71,10 @@ class CommandsTabPanel(
   }
 
   private fun addCommandButton(commandId: CommandId) {
-    addButton(commandId.displayName, commandId.mnemonic) { executeCommand(commandId) }
+    addButton(commandId.id, commandId.displayName, commandId.mnemonic) { executeCommand(commandId) }
   }
 
-  private fun addButton(displayName: String, mnemonic: Int, action: () -> Unit) {
+  private fun addButton(actionId: String, displayName: String, mnemonic: Int, action: () -> Unit) {
     val button = JButton(displayName)
     button.mnemonic = mnemonic
     button.displayedMnemonicIndex = displayName.indexOfFirst { it.code == mnemonic }
@@ -84,15 +85,21 @@ class CommandsTabPanel(
     button.addActionListener { action() }
     add(button)
 
-    buttons.add(button)
+    buttons[actionId] = button
   }
 
-  private fun update(policy: IgnorePolicy) {
+  private fun update() {
     // Dis/enable all the buttons.
-    for (button in buttons) {
-      button.isEnabled = policy == IgnorePolicy.USE
+    for ((actionId, button) in buttons) {
+      // We block only DocumentCodeAction and TestCodeAction, for EditCodeAction we will block the
+      // button on the edit dialog
+      val shouldBlockInlineEditAction =
+          isInlineEditInProgress &&
+              (actionId == DocumentCodeAction.ID || actionId == TestCodeAction.ID)
+      button.isEnabled = ignorePolicy == IgnorePolicy.USE && !shouldBlockInlineEditAction
     }
-    when (policy) {
+
+    when (ignorePolicy) {
       IgnorePolicy.USE -> {
         remove(ignoreBanner)
       }
@@ -105,6 +112,12 @@ class CommandsTabPanel(
   }
 
   override fun focusedFileIgnorePolicyChanged(policy: IgnorePolicy) {
-    runInEdt { update(policy) }
+    this.ignorePolicy = policy
+    runInEdt { update() }
+  }
+
+  override fun fixupSessionStateChanged(isInProgress: Boolean) {
+    this.isInlineEditInProgress = isInProgress
+    runInEdt { update() }
   }
 }

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

@@ -80,7 +80,7 @@ class EditCommandPrompt(
     val editor: Editor,
     dialogTitle: String,
     instruction: String? = null
-) : JFrame(), Disposable {
+) : JFrame(), Disposable, FixupService.ActiveFixupSessionStateListener {
   private val logger = Logger.getInstance(EditCommandPrompt::class.java)
 
   private val offset = editor.caretModel.primaryCaret.offset
@@ -321,6 +321,8 @@ class EditCommandPrompt(
     // Close dialog if window loses focus.
     addWindowFocusListener(windowFocusListener)
     addFocusListener(focusListener)
+
+    FixupService.getInstance(controller.project).addListener(this)
   }
 
   override fun setBounds(x: Int, y: Int, width: Int, height: Int) {
@@ -359,7 +361,9 @@ class EditCommandPrompt(
 
   @RequiresEdt
   private fun updateOkButtonState() {
-    okButton.isEnabled = instructionsField.text.isNotBlank()
+    okButton.isEnabled =
+        instructionsField.text.isNotBlank() &&
+            !FixupService.getInstance(controller.project).isEditInProgress()
   }
 
   @RequiresEdt
@@ -687,4 +691,8 @@ class EditCommandPrompt(
       return sb.toString()
     }
   }
+
+  override fun fixupSessionStateChanged(isInProgress: Boolean) {
+    runInEdt { okButton.isEnabled = !isInProgress }
+  }
 }

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

@@ -24,6 +24,9 @@ import java.util.concurrent.atomic.AtomicReference
 class FixupService(val project: Project) : Disposable {
   private val logger = Logger.getInstance(FixupService::class.java)
 
+  private val fixupSessionStateListeners: MutableList<ActiveFixupSessionStateListener> =
+      mutableListOf()
+
   @Volatile private var activeSession: FixupSession? = null
 
   // We only have one editing session at a time in JetBrains, for now.
@@ -76,21 +79,22 @@ class FixupService(val project: Project) : Disposable {
 
   fun getActiveSession(): FixupSession? = activeSession
 
-  fun setActiveSession(session: FixupSession) {
-    activeSession?.let { if (it.isShowingAcceptLens()) it.accept() else it.cancel() }
-    waitUntilActiveSessionIsFinished()
-    activeSession = session
-  }
-
-  fun waitUntilActiveSessionIsFinished() {
-    while (activeSession != null) {
-      Thread.sleep(100)
+  fun startNewSession(newSession: FixupSession) {
+    activeSession?.let { currentSession ->
+      // TODO: This should be done on the agent side, but we would need to add new client capability
+      // (parallel edits, disabled in JetBrains).
+      // We want to enable parallel edits in JetBrains soon, so this would be a wasted effort.
+      if (currentSession.isShowingAcceptLens()) {
+        currentSession.accept()
+        currentSession.dismiss()
+      } else throw IllegalStateException("Cannot start new session when one is already active")
     }
+    activeSession = newSession
   }
 
   fun clearActiveSession() {
-    // N.B. This cannot call back into the activeSession, or it will recurse.
     activeSession = null
+    notifySessionStateChanged()
   }
 
   override fun dispose() {
@@ -110,6 +114,26 @@ class FixupService(val project: Project) : Disposable {
     }
   }
 
+  fun addListener(listener: ActiveFixupSessionStateListener) {
+    synchronized(fixupSessionStateListeners) { fixupSessionStateListeners.add(listener) }
+  }
+
+  fun isEditInProgress(): Boolean {
+    return activeSession?.isShowingWorkingLens() ?: false
+  }
+
+  fun notifySessionStateChanged() {
+    synchronized(fixupSessionStateListeners) {
+      for (listener in fixupSessionStateListeners) {
+        listener.fixupSessionStateChanged(isEditInProgress())
+      }
+    }
+  }
+
+  interface ActiveFixupSessionStateListener {
+    fun fixupSessionStateChanged(isInProgress: Boolean)
+  }
+
   companion object {
     @JvmStatic
     fun getInstance(project: Project): FixupService {

src/main/kotlin/com/sourcegraph/cody/edit/actions/DocumentCodeAction.kt

@@ -1,9 +1,10 @@
 package com.sourcegraph.cody.edit.actions
 
-import com.intellij.openapi.project.DumbAware
-import com.sourcegraph.cody.autocomplete.action.CodyAction
-
 class DocumentCodeAction :
-    EditCommandAction({ editor, fixupService -> fixupService.startDocumentCode(editor) }),
-    CodyAction,
-    DumbAware
+    NonInteractiveEditCommandAction({ editor, fixupService ->
+      fixupService.startDocumentCode(editor)
+    }) {
+  companion object {
+    const val ID: String = "cody.documentCodeAction"
+  }
+}

src/main/kotlin/com/sourcegraph/cody/edit/actions/EditCodeAction.kt

@@ -1,9 +1,8 @@
 package com.sourcegraph.cody.edit.actions
 
-import com.intellij.openapi.project.DumbAware
-import com.sourcegraph.cody.autocomplete.action.CodyAction
-
 class EditCodeAction :
-    EditCommandAction({ editor, fixupService -> fixupService.startCodeEdit(editor) }),
-    CodyAction,
-    DumbAware
+    EditCommandAction({ editor, fixupService -> fixupService.startCodeEdit(editor) }) {
+  companion object {
+    const val ID: String = "cody.editCodeAction"
+  }
+}

src/main/kotlin/com/sourcegraph/cody/edit/actions/NonInteractiveEditCommandAction.kt

@@ -0,0 +1,15 @@
+package com.sourcegraph.cody.edit.actions
+
+import com.intellij.openapi.actionSystem.AnActionEvent
+import com.intellij.openapi.editor.Editor
+import com.sourcegraph.cody.edit.FixupService
+
+open class NonInteractiveEditCommandAction(runAction: (Editor, FixupService) -> Unit) :
+    EditCommandAction(runAction) {
+  override fun update(e: AnActionEvent) {
+    super.update(e)
+
+    val project = e.project ?: return
+    e.presentation.isEnabled = !FixupService.getInstance(project).isEditInProgress()
+  }
+}

src/main/kotlin/com/sourcegraph/cody/edit/actions/TestCodeAction.kt

@@ -1,9 +1,10 @@
 package com.sourcegraph.cody.edit.actions
 
-import com.intellij.openapi.project.DumbAware
-import com.sourcegraph.cody.autocomplete.action.CodyAction
-
 class TestCodeAction :
-    EditCommandAction({ editor, fixupService -> fixupService.startTestCode(editor) }),
-    CodyAction,
-    DumbAware
+    NonInteractiveEditCommandAction({ editor, fixupService ->
+      fixupService.startTestCode(editor)
+    }) {
+  companion object {
+    const val ID: String = "cody.testCodeAction"
+  }
+}