Minor refactoring of inline edit commands (#1338)

## Changes

Minor refactoring of inline edit commands to remove some duplication -
in preparation of adding `Test` command.

## Test plan

N/A, No real code changes

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

@@ -12,6 +12,7 @@ import com.intellij.ui.components.fields.ExpandableTextField
 import com.intellij.util.concurrency.annotations.RequiresEdt
 import com.sourcegraph.cody.agent.protocol.ModelUsage
 import com.sourcegraph.cody.chat.ui.LlmDropdown
+import com.sourcegraph.cody.edit.sessions.EditCodeSession
 import java.awt.BorderLayout
 import java.awt.Dimension
 import java.awt.Graphics
@@ -191,7 +192,7 @@ class EditCommandPrompt(val controller: FixupService, val editor: Editor, val di
           return
         }
         controller.addSession(
-            EditSession(controller, editor, project, editor.document, text, llmDropdown.item))
+            EditCodeSession(controller, editor, project, editor.document, text, llmDropdown.item))
       }
     }
 

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

@@ -7,6 +7,8 @@ import com.intellij.openapi.diagnostic.Logger
 import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.project.Project
 import com.sourcegraph.cody.agent.protocol.EditTask
+import com.sourcegraph.cody.edit.sessions.DocumentCodeSession
+import com.sourcegraph.cody.edit.sessions.FixupSession
 import com.sourcegraph.config.ConfigUtil.isCodyEnabled
 import com.sourcegraph.utils.CodyEditorUtil
 

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

@@ -10,6 +10,7 @@ import com.intellij.openapi.editor.Document
 import com.intellij.openapi.editor.RangeMarker
 import com.intellij.openapi.util.TextRange
 import com.sourcegraph.cody.agent.protocol.TextEdit
+import com.sourcegraph.cody.edit.sessions.FixupSession
 
 abstract class FixupUndoableAction(
     val session: FixupSession,

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

@@ -0,0 +1,9 @@
+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

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

@@ -0,0 +1,9 @@
+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

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

@@ -0,0 +1,21 @@
+package com.sourcegraph.cody.edit.actions
+
+import com.intellij.openapi.actionSystem.AnActionEvent
+import com.intellij.openapi.editor.Editor
+import com.intellij.openapi.editor.actionSystem.EditorAction
+import com.intellij.openapi.project.DumbAware
+import com.sourcegraph.cody.autocomplete.action.CodyAction
+import com.sourcegraph.cody.config.CodyApplicationSettings
+import com.sourcegraph.cody.edit.FixupService
+import com.sourcegraph.config.ConfigUtil
+
+open class EditCommandAction(runAction: (Editor, FixupService) -> Unit) :
+    EditorAction(EditCommandActionHandler(runAction)), CodyAction, DumbAware {
+  override fun update(e: AnActionEvent) {
+    super.update(e)
+
+    e.presentation.isVisible =
+        ConfigUtil.isFeatureFlagEnabled("cody.feature.inline-edits") ||
+            CodyApplicationSettings.instance.isInlineEditionEnabled
+  }
+}

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

@@ -0,0 +1,27 @@
+package com.sourcegraph.cody.edit.actions
+
+import com.intellij.openapi.actionSystem.DataContext
+import com.intellij.openapi.diagnostic.Logger
+import com.intellij.openapi.editor.Caret
+import com.intellij.openapi.editor.Editor
+import com.intellij.openapi.editor.actionSystem.EditorActionHandler
+import com.sourcegraph.cody.edit.FixupService
+import com.sourcegraph.utils.CodyEditorUtil
+
+open class EditCommandActionHandler(val runAction: (Editor, FixupService) -> Unit) :
+    EditorActionHandler() {
+  private val logger = Logger.getInstance(EditCommandActionHandler::class.java)
+
+  override fun isEnabledForCaret(editor: Editor, caret: Caret, dataContext: DataContext?): Boolean {
+    return CodyEditorUtil.isEditorValidForAutocomplete(editor)
+  }
+
+  override fun doExecute(editor: Editor, where: Caret?, dataContext: DataContext?) {
+    val project = editor.project
+    if (project == null) {
+      logger.warn("No project found, cannot run FixupService action for ${this.javaClass.name}")
+    } else {
+      runAction(editor, FixupService.getInstance(project))
+    }
+  }
+}

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

@@ -0,0 +1,20 @@
+package com.sourcegraph.cody.edit.sessions
+
+import com.intellij.openapi.editor.Document
+import com.intellij.openapi.editor.Editor
+import com.intellij.openapi.project.Project
+import com.sourcegraph.cody.agent.CodyAgent
+import com.sourcegraph.cody.agent.protocol.EditTask
+import com.sourcegraph.cody.edit.FixupService
+import java.util.concurrent.CompletableFuture
+
+class DocumentCodeSession(
+    controller: FixupService,
+    editor: Editor,
+    project: Project,
+    document: Document
+) : FixupSession(controller, editor, project, document) {
+  override fun makeEditingRequest(agent: CodyAgent): CompletableFuture<EditTask> {
+    return agent.server.commandsDocument()
+  }
+}

src/main/kotlin/com/sourcegraph/cody/edit/EditSession.kt → src/main/kotlin/com/sourcegraph/cody/edit/sessions/EditCodeSession.kt

@@ -1,43 +1,30 @@
-package com.sourcegraph.cody.edit
+package com.sourcegraph.cody.edit.sessions
 
-import com.intellij.openapi.diagnostic.Logger
 import com.intellij.openapi.editor.Document
 import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.project.Project
 import com.sourcegraph.cody.agent.CodyAgent
 import com.sourcegraph.cody.agent.protocol.ChatModelsResponse
 import com.sourcegraph.cody.agent.protocol.EditTask
 import com.sourcegraph.cody.agent.protocol.InlineEditParams
+import com.sourcegraph.cody.edit.FixupService
 import java.util.concurrent.CompletableFuture
 
 /**
  * Manages the state machine for inline-edit requests.
  *
  * @param instructions The user's instructions for fixing up the code.
  */
-class EditSession(
+class EditCodeSession(
     controller: FixupService,
     editor: Editor,
     project: Project,
     document: Document,
     val instructions: String,
     private val chatModelProvider: ChatModelsResponse.ChatModelProvider,
 ) : FixupSession(controller, editor, project, document) {
-  private val logger = Logger.getInstance(EditSession::class.java)
-
   override fun makeEditingRequest(agent: CodyAgent): CompletableFuture<EditTask> {
     val params = InlineEditParams(instructions, chatModelProvider.model)
     return agent.server.commandsEdit(params)
   }
-
-  override fun dispose() {}
-
-  override fun diff() {}
-
-  override fun retry() {
-    // TODO: The actual prompt is displayed as ghost text in the text input field.
-    // E.g. "Write a brief documentation comment for the selected code <etc.>"
-    // We need to send the prompt along with the lenses, so that the client can display it.
-    EditCommandPrompt(controller, editor, "Edit instructions and Retry").displayPromptUI()
-  }
 }

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

@@ -1,4 +1,4 @@
-package com.sourcegraph.cody.edit
+package com.sourcegraph.cody.edit.sessions
 
 import com.intellij.openapi.Disposable
 import com.intellij.openapi.application.ApplicationManager
@@ -30,6 +30,10 @@ import com.sourcegraph.cody.agent.protocol.Range
 import com.sourcegraph.cody.agent.protocol.TaskIdParam
 import com.sourcegraph.cody.agent.protocol.TextEdit
 import com.sourcegraph.cody.agent.protocol.WorkspaceEditParams
+import com.sourcegraph.cody.edit.EditCommandPrompt
+import com.sourcegraph.cody.edit.FixupService
+import com.sourcegraph.cody.edit.InsertUndoableAction
+import com.sourcegraph.cody.edit.ReplaceUndoableAction
 import com.sourcegraph.cody.edit.widget.LensGroupFactory
 import com.sourcegraph.cody.edit.widget.LensWidgetGroup
 import java.util.concurrent.CancellationException
@@ -218,9 +222,19 @@ abstract class FixupSession(
     }
   }
 
-  abstract fun retry()
+  fun retry() {
+    // TODO: The actual prompt we sent is displayed as ghost text in the text input field, in VS
+    // Code.
+    // E.g. "Write a brief documentation comment for the selected code <etc.>"
+    // We need to send the prompt along with the lenses, so that the client can display it.
+    EditCommandPrompt(controller, editor, "Edit instructions and Retry").displayPromptUI()
+  }
 
-  abstract fun diff()
+  fun diff() {
+    // TODO: Use DiffManager and bring up a diff of the changed region.
+    // You can see it in action now by clicking the green gutter to the left of Cody changes.
+    logger.warn("Code Lenses: Show Diff")
+  }
 
   fun undo() {
     CodyAgentService.withAgent(project) { agent ->
@@ -334,6 +348,8 @@ abstract class FixupSession(
     return FileEditorManager.getInstance(project).getEditors(file).firstOrNull()
   }
 
+  override fun dispose() {}
+
   companion object {
     // Lens actions the user can take; we notify the Agent when they are taken.
     const val COMMAND_ACCEPT = "cody.fixup.codelens.accept"

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

@@ -2,7 +2,7 @@ package com.sourcegraph.cody.edit.widget
 
 import com.intellij.openapi.editor.event.EditorMouseEvent
 import com.intellij.ui.JBColor
-import com.sourcegraph.cody.edit.FixupSession
+import com.sourcegraph.cody.edit.sessions.FixupSession
 import java.awt.Font
 import java.awt.FontMetrics
 import java.awt.Graphics2D

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

@@ -2,7 +2,7 @@ package com.sourcegraph.cody.edit.widget
 
 import com.intellij.openapi.diagnostic.Logger
 import com.sourcegraph.cody.Icons
-import com.sourcegraph.cody.edit.FixupSession
+import com.sourcegraph.cody.edit.sessions.FixupSession
 
 /** Handles assembling standard groups of lenses. */
 class LensGroupFactory(val session: FixupSession) {

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

@@ -16,8 +16,12 @@ import com.intellij.openapi.editor.markup.TextAttributes
 import com.intellij.openapi.util.Disposer
 import com.intellij.ui.Gray
 import com.sourcegraph.cody.agent.protocol.Range
-import com.sourcegraph.cody.edit.FixupSession
-import java.awt.*
+import com.sourcegraph.cody.edit.sessions.FixupSession
+import java.awt.Cursor
+import java.awt.Font
+import java.awt.FontMetrics
+import java.awt.Graphics2D
+import java.awt.Point
 import java.awt.geom.Rectangle2D
 import java.util.concurrent.atomic.AtomicBoolean