Got Accept, Undo and Cancel working with telemetry

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

@@ -1,5 +1,6 @@
 package com.sourcegraph.cody.edit
 
+import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.diagnostic.Logger
 import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.fileEditor.FileDocumentManager
@@ -8,6 +9,7 @@ import com.intellij.util.concurrency.annotations.RequiresEdt
 import com.sourcegraph.cody.agent.CodyAgent
 import com.sourcegraph.cody.agent.CodyAgentCodebase
 import com.sourcegraph.cody.agent.CodyAgentService.Companion.withAgent
+import com.sourcegraph.cody.agent.CommandExecuteParams
 import com.sourcegraph.cody.edit.FixupService.Companion.backgroundThread
 import com.sourcegraph.cody.edit.widget.LensWidgetGroup
 import java.util.concurrent.CancellationException
@@ -23,11 +25,11 @@ class DocumentCodeSession(editor: Editor) : FixupSession(editor) {
 
   private val lensActionCallbacks =
       mapOf(
-          "cody.fixup.codelens.accept" to { accept() },
-          "cody.fixup.codelens.cancel" to { cancel() },
-          "cody.fixup.codelens.retry" to { retry() },
-          "cody.fixup.codelens.diff" to { showDiff() },
-          "cody.fixup.codelens.undo" to { undo() },
+          ACTION_ACCEPT to { accept() },
+          ACTION_CANCEL to { cancel() },
+          ACTION_RETRY to { retry() },
+          ACTION_DIFF to { diff() },
+          ACTION_UNDO to { undo() },
       )
 
   init {
@@ -132,13 +134,16 @@ class DocumentCodeSession(editor: Editor) : FixupSession(editor) {
   }
 
   override fun accept() {
-    // TODO: Telemetry -- see get-inputs.ts
-    // telemetryRecorder.recordEvent('cody.fixup.input.model', 'selected')
+    withAgent(project) { agent ->
+      agent.server.commandExecute(CommandExecuteParams(ACTION_ACCEPT, listOf(taskId!!)))
+    }
     finish()
   }
 
   override fun cancel() {
-    // TODO: Telemetry
+    withAgent(project) { agent ->
+      agent.server.commandExecute(CommandExecuteParams(ACTION_CANCEL, listOf(taskId!!)))
+    }
     if (performedEdits) {
       undo()
     } else {
@@ -147,18 +152,33 @@ class DocumentCodeSession(editor: Editor) : FixupSession(editor) {
   }
 
   override fun retry() {
-    // TODO: Telemetry
-    FixupService.instance.documentCode(editor) // Disposes this session.
+    // 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(editor, "Edit instructions and Retry").displayPromptUI()
   }
 
   // Brings up a diff view showing the changes the AI made.
-  private fun showDiff() {
+  private fun diff() {
+    // The FixupController issues a vscode.diff command to show the smart diff in the
+    // handler for cody.fixup.codelens.diff. TODO: Register a handler in the Agent
+    // and send a new RPC to the client to display the diff, maybe just a notification.
     logger.warn("Code Lenses: Show Diff")
   }
 
   fun undo() {
-    // TODO: Telemetry
+    withAgent(project) { agent ->
+      agent.server.commandExecute(CommandExecuteParams(ACTION_UNDO, listOf(taskId!!)))
+    }
     undoEdits()
     finish()
   }
+
+  companion object {
+    const val ACTION_ACCEPT = "cody.fixup.codelens.accept"
+    const val ACTION_CANCEL = "cody.fixup.codelens.cancel"
+    const val ACTION_RETRY = "cody.fixup.codelens.retry"
+    const val ACTION_DIFF = "cody.fixup.codelens.diff"
+    const val ACTION_UNDO = "cody.fixup.codelens.undo"
+  }
 }

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

@@ -20,7 +20,7 @@ import javax.swing.event.DocumentEvent
 import javax.swing.event.DocumentListener
 
 /** Pop up a user interface for giving Cody instructions to fix up code at the cursor. */
-class EditCommandPrompt(val editor: Editor) {
+class EditCommandPrompt(val editor: Editor, val dialogTitle: String) {
   private val logger = Logger.getInstance(EditCommandPrompt::class.java)
   private val offset = editor.caretModel.primaryCaret.offset
   private val controller = FixupService.instance
@@ -143,7 +143,7 @@ class EditCommandPrompt(val editor: Editor) {
       DialogWrapper(editor.project, false, IdeModalityType.MODELESS) {
     init {
       init()
-      title = "Edit Code with Cody"
+      title = dialogTitle
 
       instructionsField.text = controller.getLastPrompt()
       updateOkButtonState()

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

@@ -22,7 +22,7 @@ class FixupService : Disposable {
   // Prompt user for instructions for editing selected code.
   fun startCodeEdit(editor: Editor) {
     if (!isEligibleForInlineEdit(editor)) return
-    EditCommandPrompt(editor).displayPromptUI()
+    EditCommandPrompt(editor, "Edit Code with Cody").displayPromptUI()
   }
 
   // Generate and insert a doc string for the current code.

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

@@ -62,6 +62,7 @@ abstract class FixupSession(val editor: Editor) : Disposable {
     WriteCommandAction.runWriteCommandAction(editor.project ?: return) {
       val doc: Document = editor.document
       val project = editor.project ?: return@runWriteCommandAction
+      // TODO: For all 3 of these, we should use a marked range to track it over edits.
       for (edit in edits) {
         // TODO: handle options if present (currently just undo bounds)
         when (edit.type) {
@@ -70,17 +71,16 @@ abstract class FixupSession(val editor: Editor) : Disposable {
           "delete" -> performDelete(doc, edit)
           else -> getLogger().warn("Unknown edit type: ${edit.type}")
         }
-        // TODO: Would be nice to group all the edits into a single undo action.
+        // TODO: Group all the edits into a single UndoableAction.
         UndoManager.getInstance(project)
-                .undoableActionPerformed(FixupUndoableAction.from(editor, edit))
+            .undoableActionPerformed(FixupUndoableAction.from(editor, edit))
       }
     }
   }
 
   private fun performReplace(doc: Document, edit: TextEdit) {
     val (start, end) = edit.range?.toOffsets(doc) ?: return
     doc.replaceString(start, end, edit.value ?: return)
-
   }
 
   private fun performInsert(doc: Document, edit: TextEdit): TextEdit? {

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

@@ -3,10 +3,12 @@ package com.sourcegraph.cody.edit
 import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.command.undo.DocumentReference
 import com.intellij.openapi.command.undo.DocumentReferenceManager
+import com.intellij.openapi.command.undo.UndoManager
 import com.intellij.openapi.command.undo.UndoableAction
 import com.intellij.openapi.diagnostic.Logger
 import com.intellij.openapi.editor.Document
 import com.intellij.openapi.editor.Editor
+import com.intellij.openapi.editor.RangeMarker
 import com.sourcegraph.cody.agent.protocol.TextEdit
 
 abstract class FixupUndoableAction(val editor: Editor, val edit: TextEdit) : UndoableAction {
@@ -23,16 +25,26 @@ abstract class FixupUndoableAction(val editor: Editor, val edit: TextEdit) : Und
     return true
   }
 
+  // TODO: Use this.
+  fun createRangeMarker(start: Int, end: Int): RangeMarker {
+    val rangeMarker = document.createRangeMarker(start, end)
+    rangeMarker.isGreedyToLeft = true
+    rangeMarker.isGreedyToRight = true
+    return rangeMarker
+  }
+
   class InsertUndoableAction(editor: Editor, edit: TextEdit) : FixupUndoableAction(editor, edit) {
 
     override fun undo() {
+      if (UndoManager.getInstance(editor.project ?: return).isUndoInProgress) return
       val offsets = (edit.range ?: return).toOffsets(document)
       ApplicationManager.getApplication().runWriteAction {
         document.deleteString(offsets.first, offsets.second)
       }
     }
 
     override fun redo() {
+      if (UndoManager.getInstance(editor.project ?: return).isUndoInProgress) return
       val offset = edit.position?.toOffset(document) ?: return
       ApplicationManager.getApplication().runWriteAction {
         document.insertString(offset, edit.value ?: "")
@@ -43,18 +55,24 @@ abstract class FixupUndoableAction(val editor: Editor, val edit: TextEdit) : Und
   class ReplaceUndoableAction(editor: Editor, edit: TextEdit) : FixupUndoableAction(editor, edit) {
 
     override fun undo() {
+      if (UndoManager.getInstance(editor.project ?: return).isUndoInProgress) return
+
       TODO("fix")
     }
 
     override fun redo() {
-      TODO("fix")
+      if (UndoManager.getInstance(editor.project ?: return).isUndoInProgress) return
+
+      TODO("finish")
     }
   }
 
   class DeleteUndoableAction(editor: Editor, edit: TextEdit) : FixupUndoableAction(editor, edit) {
     var oldText = ""
 
     override fun undo() {
+      if (UndoManager.getInstance(editor.project ?: return).isUndoInProgress) return
+
       val offsets = (edit.range ?: return).toOffsets(document)
       val deleted = document.text.substring(offsets.first, offsets.second)
       ApplicationManager.getApplication().runWriteAction {
@@ -64,6 +82,8 @@ abstract class FixupUndoableAction(val editor: Editor, val edit: TextEdit) : Und
     }
 
     override fun redo() {
+      if (UndoManager.getInstance(editor.project ?: return).isUndoInProgress) return
+
       val offsets = (edit.range ?: return).toOffsets(document)
       ApplicationManager.getApplication().runWriteAction {
         document.deleteString(offsets.first, offsets.second)

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

@@ -1,6 +1,7 @@
 package com.sourcegraph.cody.edit.widget
 
 import com.intellij.openapi.Disposable
+import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.diagnostic.Logger
 import com.intellij.openapi.util.Disposer
 import java.awt.FontMetrics
@@ -10,34 +11,26 @@ abstract class LensWidget(val parentGroup: LensWidgetGroup) : Disposable {
   protected val logger = Logger.getInstance(LensWidget::class.java)
   protected var mouseInBounds = false
 
-  init {
-    Disposer.register(parentGroup, this)
-  }
-
   abstract fun calcWidthInPixels(fontMetrics: FontMetrics): Int
 
   abstract fun calcHeightInPixels(fontMetrics: FontMetrics): Int
 
   abstract fun paint(g: Graphics2D, x: Float, y: Float)
 
-  /**
-   * Optional method for updating the widget state, useful for animations.
-   */
+  /** Optional method for updating the widget state, useful for animations. */
   open fun update() {
     // Default implementation does nothing
   }
 
-  /**
-   * Called only when widget is clicked.
-   * Coordinates are relative to the widget.
-   */
+  /** Called only when widget is clicked. Coordinates are relative to the widget. */
   open fun onClick(x: Int, y: Int): Boolean {
     return false
   }
 
   open fun onMouseEnter() {
     mouseInBounds = true
   }
+
   open fun onMouseExit() {
     mouseInBounds = false
   }

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

@@ -1,6 +1,7 @@
 package com.sourcegraph.cody.edit.widget
 
 import com.intellij.openapi.Disposable
+import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.diagnostic.Logger
 import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.editor.EditorCustomElementRenderer
@@ -16,10 +17,7 @@ import com.sourcegraph.cody.Icons
 import com.sourcegraph.cody.agent.protocol.DisplayCodeLensParams
 import com.sourcegraph.cody.agent.protocol.ProtocolCodeLens
 import com.sourcegraph.cody.edit.FixupSession
-import java.awt.Font
-import java.awt.FontMetrics
-import java.awt.Graphics2D
-import java.awt.Point
+import java.awt.*
 import java.awt.geom.Rectangle2D
 
 operator fun Point.component1() = this.x
@@ -82,6 +80,8 @@ class LensWidgetGroup(val session: FixupSession, parentComponent: Editor) :
 
   var inlay: Inlay<EditorCustomElementRenderer>? = null
 
+  private var prevCursor: Cursor? = null
+
   init {
     Disposer.register(session, this)
     editor.addEditorMouseListener(mouseClickListener)
@@ -218,6 +218,7 @@ class LensWidgetGroup(val session: FixupSession, parentComponent: Editor) :
         if (text.isNotEmpty()) separator = true
       }
     }
+    widgets.forEach { Disposer.register(this, it) }
   }
 
   // Dispatch mouse click events to the appropriate widget.
@@ -232,6 +233,15 @@ class LensWidgetGroup(val session: FixupSession, parentComponent: Editor) :
     val (x, y) = e.mouseEvent.point
     val widget = findWidgetAt(x, y)
     val lastWidget = lastHoveredWidget
+
+    if (widget != null) {
+      prevCursor = e.editor.contentComponent.cursor
+      e.editor.contentComponent.cursor = Cursor.getPredefinedCursor(Cursor.HAND_CURSOR)
+    } else if (prevCursor != null) {
+      e.editor.contentComponent.cursor = prevCursor!!
+      prevCursor = null
+    }
+
     // Check if the mouse has moved from one widget to another or from/to outside
     if (widget != lastWidget) {
       lastWidget?.onMouseExit()