Edit: introduce CodyTaskState.Reverting and FixupTask.retryID

agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/protocol_generated/CodyTaskState.kt

@@ -1,5 +1,5 @@
 @file:Suppress("FunctionName", "ClassName", "unused", "EnumEntryName", "UnusedImport")
 package com.sourcegraph.cody.protocol_generated
 
-typealias CodyTaskState = String // One of: Idle, Working, Inserting, Applying, Formatting, Applied, Finished, Error, Pending
+typealias CodyTaskState = String // One of: Idle, Working, Inserting, Applying, Formatting, Applied, Reverting, Finished, Error, Pending
 

agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/protocol_generated/Constants.kt

@@ -71,6 +71,7 @@ object Constants {
   const val Applying = "Applying"
   const val Formatting = "Formatting"
   const val Applied = "Applied"
+  const val Reverting = "Reverting"
   const val Finished = "Finished"
   const val Error = "Error"
   const val Pending = "Pending"

agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/protocol_generated/EditTask.kt

@@ -3,7 +3,7 @@ package com.sourcegraph.cody.protocol_generated
 
 data class EditTask(
   val id: String,
-  val state: CodyTaskState, // Oneof: Idle, Working, Inserting, Applying, Formatting, Applied, Finished, Error, Pending
+  val state: CodyTaskState, // Oneof: Idle, Working, Inserting, Applying, Formatting, Applied, Reverting, Finished, Error, Pending
   val error: CodyError? = null,
   val selectionRange: Range,
   val instruction: String? = null,

vscode/src/non-stop/FixupController.ts

@@ -132,14 +132,29 @@ export class FixupController
     }
 
     /**
-     * Reverts an applied fixup task by replacing the edited code range with the original code.
-     *
-     * TODO: It is possible the original code is out of date if the user edited it whilst the fixup was running.
-     * Handle this case better. Possibly take a copy of the previous code just before the fixup is applied.
+     * Reverts an applied fixup task by replacing the edited code range with the
+     * original code.
      */
     public async undo(task: FixupTask): Promise<void> {
+        const terminalState = await this.doUndo(task)
+        if (terminalState) {
+            this.setTaskState(task, terminalState)
+        }
+    }
+
+    /**
+     * Reverts an applied fixup task by replacing the edited code range with the
+     * original code.
+     *
+     * @returns the terminal state to set for the task.
+     *
+     * TODO: It is possible the original code is out of date if the user edited
+     * it whilst the fixup was running. Handle this case better. Possibly take a
+     * copy of the previous code just before the fixup is applied.
+     */
+    private async doUndo(task: FixupTask): Promise<CodyTaskState | undefined> {
         if (task.state !== CodyTaskState.Applied) {
-            return
+            return undefined
         }
 
         this.undoCommandEvent.fire(task.id)
@@ -153,9 +168,12 @@ export class FixupController
 
         const replacementText = task.replacement
         if (!replacementText) {
-            return
+            return undefined
         }
 
+        // From here on we will start reverting the task and will end up in a terminal state.
+        this.setTaskState(task, CodyTaskState.Reverting)
+
         editor.revealRange(task.selectionRange)
 
         // TODO: We should figure out why `editor.edit` often returns `null` instead of a boolean.
@@ -173,8 +191,6 @@ export class FixupController
             ...task.telemetryMetadata,
         }
 
-        this.setTaskState(task, editOk ? CodyTaskState.Finished : CodyTaskState.Error)
-
         const { metadata, privateMetadata } = splitSafeMetadata(legacyMetadata)
         if (!editOk) {
             telemetryService.log('CodyVSCodeExtension:fixup:revert:failed', legacyMetadata, {
@@ -199,6 +215,7 @@ export class FixupController
                 },
             })
         }
+        return editOk ? CodyTaskState.Finished : CodyTaskState.Error
     }
 
     // Undo the specified task, then prompt for a new set of instructions near
@@ -228,9 +245,9 @@ export class FixupController
         const updatedRange = input.range.isEqual(task.selectionRange) ? task.originalRange : input.range
 
         // Revert and remove the previous task
-        await this.undo(task)
+        const terminalState = await this.doUndo(task)
 
-        return executeEdit({
+        const newTask = await executeEdit({
             configuration: {
                 range: updatedRange,
                 instruction: input.instruction,
@@ -242,6 +259,13 @@ export class FixupController
             },
             source,
         })
+
+        task.retryID = newTask?.id
+        if (terminalState) {
+            this.setTaskState(task, terminalState)
+        }
+
+        return newTask
     }
 
     // FixupFileCollection
@@ -1066,7 +1090,7 @@ export class FixupController
         // User has changed an applied task, so we assume the user has accepted the change and wants to take control.
         // This helps ensure that the codelens doesn't stay around unnecessarily and become an annoyance.
         // Note: This will also apply if the user attempts to undo the applied change.
-        if (task.state === CodyTaskState.Applied) {
+        if (task.state === CodyTaskState.Applied || task.state === CodyTaskState.Reverting) {
             this.accept(task)
             return
         }

vscode/src/non-stop/FixupTask.ts

@@ -57,6 +57,10 @@ export class FixupTask {
      * they may run into an error that we can't anticipate
      */
     public formattingResolver: ((value: boolean) => void) | null = null
+    /**
+     * The ID of the task that is created when a user retries a task.
+     */
+    public retryID: FixupTaskID | undefined
 
     constructor(
         /**

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

@@ -34,6 +34,7 @@ export function getLensesForTask(task: FixupTask): vscode.CodeLens[] {
             const skip = getFormattingSkipLens(codeLensRange, task.id)
             return [title, skip]
         }
+        case CodyTaskState.Reverting:
         case CodyTaskState.Applied: {
             const accept = getAcceptLens(codeLensRange, task.id)
             const retry = getRetryLens(codeLensRange, task.id)

vscode/src/non-stop/utils.ts

@@ -32,6 +32,11 @@ export enum CodyTaskState {
      * E.g. Undo the change, Retry the change, View the diff.
      */
     Applied = 'Applied',
+    /**
+     * The response has been applied to the document, and we are satisfied enough to present it to the user.
+     * This is a transitional state from Applied to a terminal state which is triggered by a user undoing or retrying a change.
+     */
+    Reverting = 'Reverting',
     /**
      * Terminal state. The response has been "accepted" by the user. This is either by:
      * - Clicking "Accept" via the CodeLens