Edit: Use a tracked originalRange when retrying a task (#1926)

## Description


https://github.com/sourcegraph/cody/assets/9516420/0f6ac95a-621c-42f2-b42c-b3dd395eaaf3

Follow up on:
#1892 (comment)

We were retrying fixup tasks with the `selectionRange` - this is not
suitable as this range may shrink or expand depending on the code that
the LLM produces.

Instead we should use the `originalRange` - the range that is provided
when the fixup task is created. We still need to track the positions of
the range, as it could shift due to code being added above.

## Test plan

1. Create fixups
2. Press "retry"
3. Check retry range is correct (code is not unnecessarily deleted,
unless caused by the LLM)

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

vscode/CHANGELOG.md

@@ -24,6 +24,8 @@ Starting from `0.2.0`, Cody is using `major.EVEN_NUMBER.patch` for release versi
 - Chat: Add delays before sending webview ready events to prevent premature sending. This fixes issue where chat panel fails to load when multiple chat panels are opened simultaneously. [pull/1836](https://github.com/sourcegraph/cody/pull/1836)
 - Autocomplete: Fixes a bug that caused autocomplete to be triggered at the end of a block or function invocation. [pull/1864](https://github.com/sourcegraph/cody/pull/1864)
 - Edit: Incoming edits that are afixed to the selected code and now handled properly (e.g. docstrings). [pull/1724](https://github.com/sourcegraph/cody/pull/1724)
+- Chat: Allowed backspace and delete keys to remove characters in chat messages input box.
+- Edit: Retrying an edit will now correctly use the original intended range. [pull/1926](https://github.com/sourcegraph/cody/pull/1926)
 - Chat: Allowed backspace and delete keys to remove characters in chat messages input box. [pull/1906](https://github.com/sourcegraph/cody/pull/1906)
 - Chat: The commands display box in the chat input box now uses the same styles as the @ command results box. [pull/1962](https://github.com/sourcegraph/cody/pull/1962)
 - Chat: Fix chat command selection to only filter on '/' prefix. [pull/1980](https://github.com/sourcegraph/cody/pull/1980)

vscode/src/non-stop/FixupController.ts

@@ -860,7 +860,7 @@ export class FixupController
         if (!task) {
             return
         }
-        const previousRange = task.selectionRange
+        const previousRange = task.originalRange
         const previousInstruction = task.instruction
         const document = await vscode.workspace.openTextDocument(task.fixupFile.uri)
 
@@ -872,7 +872,7 @@ export class FixupController
 
         void vscode.commands.executeCommand(
             'cody.command.edit-code',
-            { range: previousRange, instruction, document, intent: task.intent },
+            { range: previousRange, instruction, document, intent: task.intent, insertMode: task.insertMode },
             'code-lens'
         )
     }

vscode/src/non-stop/FixupDocumentEditObserver.ts

@@ -2,7 +2,7 @@ import * as vscode from 'vscode'
 
 import { Edit, Position, Range } from './diff'
 import { FixupFileCollection, FixupTextChanged } from './roles'
-import { TextChange, updateRangeMultipleChanges } from './tracked-range'
+import { TextChange, updateFixedRange, updateRangeMultipleChanges } from './tracked-range'
 
 // This does some thunking to manage the two range types: diff ranges, and
 // text change ranges.
@@ -61,9 +61,11 @@ export class FixupDocumentEditObserver {
         const tasks = this.provider_.tasksForFile(file)
         // Notify which tasks have changed text or the range edits apply to
         for (const task of tasks) {
-            const targetRange = task.selectionRange
             for (const edit of event.contentChanges) {
-                if (edit.range.end.isBefore(targetRange.start) || edit.range.start.isAfter(targetRange.end)) {
+                if (
+                    edit.range.end.isBefore(task.selectionRange.start) ||
+                    edit.range.start.isAfter(task.selectionRange.end)
+                ) {
                     continue
                 }
                 this.provider_.textDidChange(task)
@@ -83,6 +85,13 @@ export class FixupDocumentEditObserver {
                 task.selectionRange = updatedRange
                 this.provider_.rangeDidChange(task)
             }
+
+            // We keep track of where the original range should be, so we can re-use it for retries.
+            // Note: This range doesn't expand or shrink, it needs to match the original range as applied to `task.original`
+            const updatedFixedRange = updateRangeMultipleChanges(task.originalRange, changes, {}, updateFixedRange)
+            if (!updatedFixedRange.isEqual(task.originalRange)) {
+                task.originalRange = updatedFixedRange
+            }
         }
     }
 }

vscode/src/non-stop/FixupTask.ts

@@ -16,6 +16,11 @@ export class FixupTask {
      * The original text that we're working on updating. Set when we start an LLM spin.
      */
     public original = ''
+    /**
+     * The original range that we're working on updating.
+     * Used to perform an accurate retry. We cannot use `selectionRange` as that range may expand with the replacement code.
+     */
+    public originalRange: vscode.Range
     /** The text of the streaming turn of the LLM, if any */
     public inProgressReplacement: string | undefined
     /** The text of the last completed turn of the LLM, if any */
@@ -29,6 +34,11 @@ export class FixupTask {
     public diff: Diff | undefined
     /** The number of times we've submitted this to the LLM. */
     public spinCount = 0
+    /**
+     * A callback to skip formatting.
+     * We use the users' default editor formatter so it is possible that
+     * they may run into an error that we can't anticipate
+     */
     public formattingResolver: ((value: boolean) => void) | null = null
 
     constructor(
@@ -44,6 +54,7 @@ export class FixupTask {
     ) {
         this.id = Date.now().toString(36).replaceAll(/\d+/g, '')
         this.instruction = instruction.replace(/^\/(edit|fix)/, '').trim()
+        this.originalRange = selectionRange
     }
 
     /**

vscode/src/non-stop/tracked-range.test.ts

@@ -3,7 +3,7 @@ import assert from 'assert'
 import { describe, expect, it } from 'vitest'
 import { Position, Range } from 'vscode'
 
-import { updateRange, updateRangeMultipleChanges, UpdateRangeOptions } from './tracked-range'
+import { updateFixedRange, updateRange, updateRangeMultipleChanges, UpdateRangeOptions } from './tracked-range'
 
 // Creates a position.
 function pos(line: number, character: number): Position {
@@ -222,6 +222,30 @@ describe('Tracked Range', () => {
     })
 })
 
+// Given a spec with a tracked range in [], an edited range in (),
+// replaces () with the specified text; applies range tracking;
+// and serializes the resulting text and tracked range.
+function trackFixed(spec: string, replacement: string): string {
+    const scenario = parse(spec)
+    const editedText = edit(scenario.text, scenario.edited, replacement)
+    const updatedRange = updateFixedRange(scenario.tracked, { range: scenario.edited, text: replacement })
+    return updatedRange ? show(editedText, updatedRange) : editedText
+}
+
+describe('Tracked Range Fixed', () => {
+    it('should preserve the fixed range for changes before the range', () => {
+        expect(trackFixed('()[hello world]', 'I am here, ')).toBe('[I am here, ]hello world')
+    })
+
+    it('should preserve the fixed range for changes after the range', () => {
+        expect(trackFixed('[hello world]( )', ', I am here')).toBe('[hello world], I am here')
+    })
+
+    it('should preserve the fixed range for changes within the range', () => {
+        expect(trackFixed('[hello ()orld]', 'w')).toBe('[hello worl]d')
+    })
+})
+
 describe('Multi-change Tracked Range', () => {
     it('should stack edits on separate lines', () => {
         const range = rng(pos(2, 7), pos(17, 13))

vscode/src/non-stop/tracked-range.ts

@@ -29,7 +29,8 @@ export interface UpdateRangeOptions {
 export function updateRangeMultipleChanges(
     range: vscode.Range,
     changes: TextChange[],
-    options: UpdateRangeOptions = {}
+    options: UpdateRangeOptions = {},
+    rangeUpdater = updateRange
 ): vscode.Range {
     changes.sort((a, b) => (b.range.start.isBefore(a.range.start) ? -1 : 1))
     for (let i = 0; i < changes.length - 1; i++) {
@@ -39,7 +40,7 @@ export function updateRangeMultipleChanges(
         )
     }
     for (const change of changes) {
-        range = updateRange(range, change, options)
+        range = rangeUpdater(range, change, options)
     }
     return range
 }
@@ -151,3 +152,42 @@ export function updateRange(range: vscode.Range, change: TextChange, options: Up
     }
     return range
 }
+
+/**
+ * Given a range and an edit, shifts the range for the edit.
+ * Only handles edits that are outside of the range, as it is purely focused on shifting a fixed range in a document.
+ * Does not expand or shrink the original rank.
+ */
+export function updateFixedRange(range: vscode.Range, change: TextChange): vscode.Range {
+    const lines = change.text.split(/\r\n|\r|\n/m)
+    const insertedLastLine = lines.at(-1)?.length
+    if (insertedLastLine === undefined) {
+        throw new TypeError('unreachable') // Any string .split produces a non-empty array.
+    }
+    const insertedLineBreaks = lines.length - 1
+
+    // The only case where a range should be shifted is when the change happens before the range.
+    // In this case, we just need to adjust the start and end position depending on if the incoming change added or removed text.
+    if (change.range.end.isBefore(range.start)) {
+        range = range.with(
+            range.start.translate(
+                change.range.start.line - change.range.end.line + insertedLineBreaks,
+                change.range.end.line === range.start.line
+                    ? insertedLastLine +
+                          -change.range.end.character +
+                          (insertedLineBreaks === 0 ? change.range.start.character : 0)
+                    : 0
+            ),
+            range.end.translate(
+                change.range.start.line - change.range.end.line + insertedLineBreaks,
+                change.range.end.line === range.end.line
+                    ? insertedLastLine -
+                          change.range.end.character +
+                          (insertedLineBreaks === 0 ? change.range.start.character : 0)
+                    : 0
+            )
+        )
+    }
+
+    return range
+}