Edit: Add codelens shortcuts (#2757)

## Description

<img width="377" alt="image"
src="https://github.com/sourcegraph/cody/assets/9516420/30f19019-dfb9-41b6-961a-eaffced03dd5">


This PR:
- Adds keyboard shortcuts to common Edit codelens actions (Accept,
Retry, Undo, Cancel)
- These shortcuts are only enabled when an Edit is in an active
actionable state, and it is in a visible document. This will help reduce
the risk of these shortcuts becoming an annoyance if a user doesn't use
Edit much.
- Updates "Retry" text to show "Edit & Retry"


## Test plan

Single edit:

1. Create an edit
2. Try undo/retry/save shortcuts

Multiple edit (actions nearest edit):

1. Create multiple edits
2. Try undo/retry/save shortcuts and confirm that they only trigger
against the nearest edit to the cursor position

Shortcut enablement: 

1. Create an edit
2. Leave the file
3. Try undo/retry/save shortcuts. Confirm nothing happens
4. Return to the file
5. Try undo/retry/save shortcuts. Confirm they work

<!-- Required. See
https://sourcegraph.com/docs/dev/background-information/testing_principles.
-->

---------

Co-authored-by: Beatrix <68532117+abeatrix@users.noreply.github.com>

vscode/CHANGELOG.md

@@ -6,6 +6,7 @@ This is a log of all notable changes to Cody for VS Code. [Unreleased] changes a
 
 ### Added
 
+- Edit: Added keyboard shortcuts for codelens actions such as "Undo" and "Retry" [pull/2757][https://github.com/sourcegraph/cody/pull/2757]
 - Chat: Displays warnings for large @-mentioned files during selection. [pull/3118](https://github.com/sourcegraph/cody/pull/3118)
 
 ### Fixed

vscode/package.json

@@ -528,6 +528,26 @@
                 "command": "cody.autocomplete.manual-trigger",
                 "key": "alt+\\",
                 "when": "editorTextFocus && !editorHasSelection && config.cody.autocomplete.enabled && !inlineSuggestionsVisible"
+            },
+            {
+                "command": "cody.fixup.acceptNearest",
+                "key": "alt+a",
+                "when": "cody.activated && !editorReadonly && cody.hasActionableEdit"
+            },
+            {
+                "command": "cody.fixup.retryNearest",
+                "key": "alt+r",
+                "when": "cody.activated && !editorReadonly && cody.hasActionableEdit"
+            },
+            {
+                "command": "cody.fixup.undoNearest",
+                "key": "alt+x",
+                "when": "cody.activated && !editorReadonly && cody.hasActionableEdit"
+            },
+            {
+                "command": "cody.fixup.cancelNearest",
+                "key": "alt+z",
+                "when": "cody.activated && !editorReadonly && cody.hasActionableEdit"
             }
         ],
         "submenus": [

vscode/src/non-stop/FixupController.ts

@@ -17,7 +17,7 @@ import { countCode } from '../services/utils/code-count'
 import { getEditorInsertSpaces, getEditorTabSize } from '../utils'
 
 import { computeDiff, type Diff } from './diff'
-import { FixupCodeLenses } from './FixupCodeLenses'
+import { FixupCodeLenses } from './codelenses/provider'
 import { ContentProvider } from './FixupContentStore'
 import { FixupDecorator } from './FixupDecorator'
 import { FixupDocumentEditObserver } from './FixupDocumentEditObserver'
@@ -26,9 +26,10 @@ import { FixupFileObserver } from './FixupFileObserver'
 import { FixupScheduler } from './FixupScheduler'
 import { FixupTask, type taskID } from './FixupTask'
 import type { FixupFileCollection, FixupIdleTaskRunner, FixupTextChanged } from './roles'
-import { CodyTaskState } from './utils'
+import { CodyTaskState, getMinimumDistanceToRangeBoundary } from './utils'
 import { getInput } from '../edit/input/get-input'
 import type { AuthProvider } from '../services/AuthProvider'
+import { ACTIONABLE_TASK_STATES, CANCELABLE_TASK_STATES } from './codelenses/constants'
 
 // This class acts as the factory for Fixup Tasks and handles communication between the Tree View and editor
 export class FixupController
@@ -104,6 +105,34 @@ export class FixupController
                 })
                 telemetryRecorder.recordEvent('cody.fixup.codeLens', 'skipFormatting')
                 return this.skipFormatting(id)
+            }),
+            vscode.commands.registerCommand('cody.fixup.cancelNearest', () => {
+                const nearestTask = this.getNearestTask({ filter: { states: CANCELABLE_TASK_STATES } })
+                if (!nearestTask) {
+                    return
+                }
+                return vscode.commands.executeCommand('cody.fixup.codelens.cancel', nearestTask.id)
+            }),
+            vscode.commands.registerCommand('cody.fixup.acceptNearest', () => {
+                const nearestTask = this.getNearestTask({ filter: { states: ACTIONABLE_TASK_STATES } })
+                if (!nearestTask) {
+                    return
+                }
+                return vscode.commands.executeCommand('cody.fixup.codelens.accept', nearestTask.id)
+            }),
+            vscode.commands.registerCommand('cody.fixup.retryNearest', () => {
+                const nearestTask = this.getNearestTask({ filter: { states: ACTIONABLE_TASK_STATES } })
+                if (!nearestTask) {
+                    return
+                }
+                return vscode.commands.executeCommand('cody.fixup.codelens.retry', nearestTask.id)
+            }),
+            vscode.commands.registerCommand('cody.fixup.undoNearest', () => {
+                const nearestTask = this.getNearestTask({ filter: { states: ACTIONABLE_TASK_STATES } })
+                if (!nearestTask) {
+                    return
+                }
+                return vscode.commands.executeCommand('cody.fixup.codelens.undo', nearestTask.id)
             })
         )
         // Observe file renaming and deletion
@@ -946,6 +975,9 @@ export class FixupController
         for (const [file, editors] of editorsByFile.entries()) {
             this.decorator.didChangeVisibleTextEditors(file, editors)
         }
+
+        // Update shortcut enablement for visible files
+        this.codelenses.updateKeyboardShortcutEnablement([...editorsByFile.keys()])
     }
 
     private updateDiffs(): void {
@@ -1130,6 +1162,33 @@ export class FixupController
         }
     }
 
+    private getNearestTask({ filter }: { filter: { states: CodyTaskState[] } }): FixupTask | undefined {
+        const editor = vscode.window.activeTextEditor
+        if (!editor) {
+            return
+        }
+
+        const fixupFile = this.maybeFileForUri(editor.document.uri)
+        if (!fixupFile) {
+            return
+        }
+
+        const position = editor.selection.active
+
+        /**
+         * Get the task closest to the current cursor position from the tasks associated with the current file.
+         */
+        const closestTask = this.tasksForFile(fixupFile)
+            .filter(({ state }) => filter.states.includes(state))
+            .sort(
+                (a, b) =>
+                    getMinimumDistanceToRangeBoundary(position, a.selectionRange) -
+                    getMinimumDistanceToRangeBoundary(position, b.selectionRange)
+            )[0]
+
+        return closestTask
+    }
+
     private reset(): void {
         this.tasks = new Map<taskID, FixupTask>()
     }

vscode/src/non-stop/codelenses/constants.ts

@@ -0,0 +1,19 @@
+import { CodyTaskState } from '../utils'
+
+export const CANCELABLE_TASK_STATES = [
+    CodyTaskState.pending,
+    CodyTaskState.working,
+    CodyTaskState.inserting,
+    CodyTaskState.applying,
+]
+
+export const ACTIONABLE_TASK_STATES = [
+    // User can Accept, Undo, Retry, etc
+    CodyTaskState.applied,
+]
+
+/**
+ * The task states where there is a direct command that the users is likely to action.
+ * This is used to help enable/disable keyboard shortcuts depending on the states in the document
+ */
+export const ALL_ACTIONABLE_TASK_STATES = [...ACTIONABLE_TASK_STATES, ...CANCELABLE_TASK_STATES]

vscode/src/non-stop/codelenses.ts → vscode/src/non-stop/codelenses/items.ts

@@ -2,8 +2,8 @@ import * as vscode from 'vscode'
 
 import { isRateLimitError } from '@sourcegraph/cody-shared'
 
-import type { FixupTask } from './FixupTask'
-import { CodyTaskState } from './utils'
+import type { FixupTask } from '../FixupTask'
+import { CodyTaskState } from '../utils'
 
 // Create Lenses based on state
 export function getLensesForTask(task: FixupTask): vscode.CodeLens[] {
@@ -17,30 +17,30 @@ export function getLensesForTask(task: FixupTask): vscode.CodeLens[] {
             return [title, cancel]
         }
         case CodyTaskState.inserting: {
-            let title = getInsertingLens(codeLensRange)
             if (isTest) {
-                title = getUnitTestLens(codeLensRange)
+                return [getUnitTestLens(codeLensRange)]
             }
-            return [title]
+            return [getInsertingLens(codeLensRange), getCancelLens(codeLensRange, task.id)]
         }
         case CodyTaskState.applying: {
             const title = getApplyingLens(codeLensRange)
-            return [title]
+            const cancel = getCancelLens(codeLensRange, task.id)
+            return [title, cancel]
         }
         case CodyTaskState.formatting: {
             const title = getFormattingLens(codeLensRange)
             const skip = getFormattingSkipLens(codeLensRange, task.id)
             return [title, skip]
         }
         case CodyTaskState.applied: {
-            const title = getAppliedLens(codeLensRange, task.id)
             const accept = getAcceptLens(codeLensRange, task.id)
             const retry = getRetryLens(codeLensRange, task.id)
             const undo = getUndoLens(codeLensRange, task.id)
+            const showDiff = getDiffLens(codeLensRange, task.id)
             if (isTest) {
                 return [accept, undo]
             }
-            return [title, accept, retry, undo]
+            return [accept, retry, undo, showDiff]
         }
         case CodyTaskState.error: {
             const title = getErrorLens(codeLensRange, task)
@@ -135,8 +135,9 @@ function getFormattingSkipLens(codeLensRange: vscode.Range, id: string): vscode.
 
 function getCancelLens(codeLensRange: vscode.Range, id: string): vscode.CodeLens {
     const lens = new vscode.CodeLens(codeLensRange)
+    const shortcut = process.platform === 'darwin' ? '⌥Z' : 'Alt+Z'
     lens.command = {
-        title: 'Cancel',
+        title: `Cancel [${shortcut}]`,
         command: 'cody.fixup.codelens.cancel',
         arguments: [id],
     }
@@ -153,10 +154,10 @@ function getDiscardLens(codeLensRange: vscode.Range, id: string): vscode.CodeLen
     return lens
 }
 
-function getAppliedLens(codeLensRange: vscode.Range, id: string): vscode.CodeLens {
+function getDiffLens(codeLensRange: vscode.Range, id: string): vscode.CodeLens {
     const lens = new vscode.CodeLens(codeLensRange)
     lens.command = {
-        title: '$(cody-logo) Show Diff',
+        title: 'Show Diff',
         command: 'cody.fixup.codelens.diff',
         arguments: [id],
     }
@@ -165,8 +166,9 @@ function getAppliedLens(codeLensRange: vscode.Range, id: string): vscode.CodeLen
 
 function getRetryLens(codeLensRange: vscode.Range, id: string): vscode.CodeLens {
     const lens = new vscode.CodeLens(codeLensRange)
+    const shortcut = process.platform === 'darwin' ? '⌥R' : 'Alt+R'
     lens.command = {
-        title: 'Retry',
+        title: `Edit & Retry (${shortcut})`,
         command: 'cody.fixup.codelens.retry',
         arguments: [id],
     }
@@ -175,8 +177,9 @@ function getRetryLens(codeLensRange: vscode.Range, id: string): vscode.CodeLens
 
 function getUndoLens(codeLensRange: vscode.Range, id: string): vscode.CodeLens {
     const lens = new vscode.CodeLens(codeLensRange)
+    const shortcut = process.platform === 'darwin' ? '⌥X' : 'Alt+X'
     lens.command = {
-        title: 'Undo',
+        title: `Undo (${shortcut})`,
         command: 'cody.fixup.codelens.undo',
         arguments: [id],
     }
@@ -185,8 +188,9 @@ function getUndoLens(codeLensRange: vscode.Range, id: string): vscode.CodeLens {
 
 function getAcceptLens(codeLensRange: vscode.Range, id: string): vscode.CodeLens {
     const lens = new vscode.CodeLens(codeLensRange)
+    const shortcut = process.platform === 'darwin' ? '⌥A' : 'Alt+A'
     lens.command = {
-        title: 'Accept',
+        title: `$(cody-logo) Accept (${shortcut})`,
         command: 'cody.fixup.codelens.accept',
         arguments: [id],
     }

vscode/src/non-stop/FixupCodeLenses.ts → vscode/src/non-stop/codelenses/provider.ts

@@ -1,9 +1,11 @@
 import * as vscode from 'vscode'
 
-import { getLensesForTask } from './codelenses'
-import type { FixupTask } from './FixupTask'
-import type { FixupFileCollection } from './roles'
-import { CodyTaskState } from './utils'
+import { getLensesForTask } from './items'
+import type { FixupTask } from '../FixupTask'
+import type { FixupFileCollection } from '../roles'
+import { CodyTaskState } from '../utils'
+import type { FixupFile } from '../FixupFile'
+import { ALL_ACTIONABLE_TASK_STATES } from './constants'
 
 export class FixupCodeLenses implements vscode.CodeLensProvider {
     private taskLenses = new Map<FixupTask, vscode.CodeLens[]>()
@@ -39,6 +41,7 @@ export class FixupCodeLenses implements vscode.CodeLensProvider {
     }
 
     public didUpdateTask(task: FixupTask): void {
+        this.updateKeyboardShortcutEnablement([task.fixupFile])
         if (task.state === CodyTaskState.finished) {
             this.removeLensesFor(task)
             return
@@ -48,6 +51,7 @@ export class FixupCodeLenses implements vscode.CodeLensProvider {
     }
 
     public didDeleteTask(task: FixupTask): void {
+        this.updateKeyboardShortcutEnablement([task.fixupFile])
         this.removeLensesFor(task)
     }
 
@@ -58,6 +62,21 @@ export class FixupCodeLenses implements vscode.CodeLensProvider {
         }
     }
 
+    /**
+     * For a set of active files, check to see if any tasks within these files are currently actionable.
+     * If they are, enable the code lens keyboard shortcuts in the editor.
+     */
+    public updateKeyboardShortcutEnablement(activeFiles: FixupFile[]): void {
+        const allTasks = activeFiles
+            .filter(file =>
+                vscode.window.visibleTextEditors.some(editor => editor.document.uri === file.uri)
+            )
+            .flatMap(file => this.files.tasksForFile(file))
+
+        const hasActionableEdit = allTasks.some(task => ALL_ACTIONABLE_TASK_STATES.includes(task.state))
+        void vscode.commands.executeCommand('setContext', 'cody.hasActionableEdit', hasActionableEdit)
+    }
+
     private notifyCodeLensesChanged(): void {
         this._onDidChangeCodeLenses.fire()
     }

vscode/src/non-stop/utils.test.ts

@@ -0,0 +1,27 @@
+import { describe, expect, it } from 'vitest'
+import * as vscode from 'vscode'
+
+import { getMinimumDistanceToRangeBoundary } from './utils'
+
+describe('getMinimumDistanceToRangeBoundary', () => {
+    it('returns start distance when position is before range', () => {
+        const position = new vscode.Position(5, 0)
+        const range = new vscode.Range(10, 0, 20, 0)
+        const minDistance = getMinimumDistanceToRangeBoundary(position, range)
+        expect(minDistance).toBe(5)
+    })
+
+    it('returns end distance when position is after range', () => {
+        const position = new vscode.Position(25, 0)
+        const range = new vscode.Range(10, 0, 20, 0)
+        const minDistance = getMinimumDistanceToRangeBoundary(position, range)
+        expect(minDistance).toBe(5)
+    })
+
+    it('returns smaller of start and end distances when position is in range', () => {
+        const position = new vscode.Position(18, 0)
+        const range = new vscode.Range(10, 0, 20, 0)
+        const minDistance = getMinimumDistanceToRangeBoundary(position, range)
+        expect(minDistance).toBe(2)
+    })
+})

vscode/src/non-stop/utils.ts

@@ -1,3 +1,5 @@
+import type * as vscode from 'vscode'
+
 export enum CodyTaskState {
     idle = 1,
     working = 2,
@@ -19,3 +21,15 @@ export function isTerminalCodyTaskState(state: CodyTaskState): boolean {
             return false
     }
 }
+
+/**
+ * Calculates the minimum distance from the given position to the start or end of the provided range.
+ */
+export function getMinimumDistanceToRangeBoundary(
+    position: vscode.Position,
+    range: vscode.Range
+): number {
+    const startDistance = Math.abs(position.line - range.start.line)
+    const endDistance = Math.abs(position.line - range.end.line)
+    return Math.min(startDistance, endDistance)
+}

vscode/test/e2e/command-edit.test.ts

@@ -34,9 +34,9 @@ test('code lenses for edit (fixup) task', async ({ page, sidebar }) => {
     const inputBox = page.getByPlaceholder(/^Enter edit instructions \(type @ to include code/)
     const instruction = 'replace hello with goodbye'
     const inputTitle = /^Edit index.html:(\d+).* with Cody$/
-    const showDiffLens = page.getByRole('button', { name: 'A Show Diff' })
+    const showDiffLens = page.getByRole('button', { name: 'Show Diff' })
     const acceptLens = page.getByRole('button', { name: 'Accept' })
-    const retryLens = page.getByRole('button', { name: 'Retry' })
+    const retryLens = page.getByRole('button', { name: 'Edit & Retry' })
     const undoLens = page.getByRole('button', { name: 'Undo' })
 
     // Wait for the input box to appear with the document name in title