Adding Smart Selection to FixupRecipe (#1317)

Original PR -> #1083 (Closed
because @abeatrix added some really amazing folding range functions to
get the resize the selection range)

Related Issues -> #585 and
#223

## Problem Summary
Currently the selection made by users is not always fully inclusive of
the entire range where changes might need to be made by Cody to satisfy
the request of fixup by the user. Because we're limiting the changes to
only the region that was selected by the user, that creates quite a few
problems:
1. Empty Selections don't work( shown by  @sqs ) 
<img width="814" alt="image"
src="https://github.com/sourcegraph/cody/assets/11155207/115da45a-b142-4d40-94f5-baf91f2c8a64">

2. Variable renaming within a function would miss some parts of the
function with the original variable name that wasn't selected
4. Bad selection boundaries cause issues with indentation with the
changes are applied coz the diffs function doesn't work perfectly in
some edge cases
5. Refactors -> If you only selected the return value then any refactor
would miss other parts of the function where changes need to be applied
6. Imperfect selection boundaries sometimes lead to repetition of code
outside the selection range and sometimes it doesn't generate outputs
7. Hard mode Problem: Its hard to do a mega refactor where change
request is on a small selection but that edits multiple files and other
places in the codebase to make it all work together in synergy.

## Original solution
Originally i leveraged an LLM call to decide the folding range and that
was a cool algorithm but it had the latency of an extra LLM call. Now
that @abeatrix added some really cool folding range functions I can just
leverage them to get a better range for the selection.

## Video

### Before


https://github.com/sourcegraph/cody/assets/11155207/53d63ee7-8650-4c03-b935-dfcb64d38c01


https://github.com/sourcegraph/cody/assets/11155207/5ab30335-6c2d-4142-849d-47234d6f4100


## After


https://drive.google.com/file/d/1RVShpNGiWK4wHJW6N_tHEGOdkzSwxM6d/view?usp=sharing


## Test plan

<!-- Required. See
https://docs.sourcegraph.com/dev/background-information/testing_principles.
-->
Tested in my local machine a few times. Works perfectly on edge cases
too.

---------

Co-authored-by: Dominic Cooney <dominic.cooney@gmail.com>

lib/shared/src/chat/recipes/fixup.ts

@@ -1,5 +1,5 @@
 import { ContextMessage, getContextMessageWithResponse } from '../../codebase-context/messages'
-import { VsCodeFixupTaskRecipeData } from '../../editor'
+import { VsCodeFixupController, VsCodeFixupTaskRecipeData } from '../../editor'
 import { IntentClassificationOption } from '../../intent-detector'
 import { MAX_CURRENT_FILE_TOKENS, MAX_HUMAN_INPUT_TOKENS } from '../../prompt/constants'
 import { populateCodeContextTemplate, populateCurrentEditorDiagnosticsTemplate } from '../../prompt/templates'
@@ -14,7 +14,7 @@ import { Recipe, RecipeContext, RecipeID } from './recipe'
  * This is either provided by the user, or inferred from their instructions
  */
 export type FixupIntent = 'add' | 'edit' | 'document'
-const FixupIntentClassification: IntentClassificationOption<FixupIntent>[] = [
+export const FixupIntentClassification: IntentClassificationOption<FixupIntent>[] = [
     {
         id: 'edit',
         rawCommand: '/edit',
@@ -52,21 +52,16 @@ export class Fixup implements Recipe {
             return null
         }
 
-        const fixupTask = await fixupController.getTaskRecipeData(taskId)
-        if (!fixupTask) {
-            await context.editor.showWarningMessage('Select some code to fixup.')
-            return null
-        }
+        const intent = await this.fetchRecipeIntent(fixupController, taskId, context)
+        const enableSmartSelection = intent === 'edit'
+        const fixupTask = await fixupController.getTaskRecipeData(taskId, { enableSmartSelection })
 
-        const quarterFileContext = Math.floor(MAX_CURRENT_FILE_TOKENS / 4)
-        if (truncateText(fixupTask.selectedText, MAX_CURRENT_FILE_TOKENS) !== fixupTask.selectedText) {
-            const msg = "The amount of text selected exceeds Cody's current capacity."
-            await context.editor.showWarningMessage(msg)
+        if (!fixupTask || !intent) {
             return null
         }
 
-        const intent = await this.getIntent(fixupTask, context)
         const promptText = this.getPrompt(fixupTask, intent)
+        const quarterFileContext = Math.floor(MAX_CURRENT_FILE_TOKENS / 4)
 
         return Promise.resolve(
             new Interaction(
@@ -83,24 +78,18 @@ export class Fixup implements Recipe {
         )
     }
 
-    private async getIntent(task: VsCodeFixupTaskRecipeData, context: RecipeContext): Promise<FixupIntent> {
-        if (task.selectedText.trim().length === 0) {
-            // Nothing selected, assume this is always 'add'.
-            return 'add'
+    public async fetchRecipeIntent(
+        fixupController: VsCodeFixupController,
+        taskId: string,
+        context: RecipeContext
+    ): Promise<FixupIntent | null> {
+        try {
+            const intent = await fixupController.getTaskIntent(taskId, context.intentDetector)
+            return intent
+        } catch (error) {
+            await context.editor.showWarningMessage((error as Error).message)
+            return null
         }
-
-        /**
-         * TODO(umpox): We should probably find a shorter way of detecting intent when possible.
-         * Possible methods:
-         * - Input -> Match first word against update|fix|add|delete verbs
-         * - Context -> Infer intent from context, e.g. Current file is a test -> Test intent, Current selection is a comment symbol -> Documentation intent
-         */
-        const intent = await context.intentDetector.classifyIntentFromOptions(
-            task.instruction,
-            FixupIntentClassification,
-            'edit'
-        )
-        return intent
     }
 
     public getPrompt(task: VsCodeFixupTaskRecipeData, intent: FixupIntent): string {

lib/shared/src/editor/index.ts

@@ -1,6 +1,8 @@
 import { URI } from 'vscode-uri'
 
 import { CodyPrompt } from '../chat/prompts'
+import { FixupIntent } from '../chat/recipes/fixup'
+import { IntentDetector } from '../intent-detector'
 
 export interface ActiveTextEditor {
     content: string
@@ -66,7 +68,11 @@ export interface VsCodeFixupTaskRecipeData {
 }
 
 export interface VsCodeFixupController {
-    getTaskRecipeData(taskId: string): Promise<VsCodeFixupTaskRecipeData | undefined>
+    getTaskRecipeData(
+        taskId: string,
+        options: { enableSmartSelection?: boolean }
+    ): Promise<VsCodeFixupTaskRecipeData | undefined>
+    getTaskIntent(taskId: string, intentDetector: IntentDetector): Promise<FixupIntent>
 }
 
 export interface VsCodeCommandsController {
@@ -105,7 +111,6 @@ export interface Editor<
     getActiveTextEditor(): ActiveTextEditor | null
     getActiveTextEditorSelection(): ActiveTextEditorSelection | null
     getActiveTextEditorSmartSelection(): Promise<ActiveTextEditorSelection | null>
-
     getActiveInlineChatTextEditor(): ActiveTextEditor | null
     getActiveInlineChatSelection(): ActiveTextEditorSelection | null
 

vscode/src/non-stop/FixupController.ts

@@ -1,7 +1,12 @@
 import * as vscode from 'vscode'
 
+import { FixupIntent, FixupIntentClassification } from '@sourcegraph/cody-shared/src/chat/recipes/fixup'
 import { VsCodeFixupController, VsCodeFixupTaskRecipeData } from '@sourcegraph/cody-shared/src/editor'
+import { IntentDetector } from '@sourcegraph/cody-shared/src/intent-detector'
+import { MAX_CURRENT_FILE_TOKENS } from '@sourcegraph/cody-shared/src/prompt/constants'
+import { truncateText } from '@sourcegraph/cody-shared/src/prompt/truncation'
 
+import { getSmartSelection } from '../editor/utils'
 import { logDebug } from '../log'
 import { countCode } from '../services/InlineAssist'
 import { telemetryService } from '../services/telemetry'
@@ -178,6 +183,88 @@ export class FixupController
         return undefined
     }
 
+    /**
+     * Retrieves the intent for a specific task based on the selected text and other contextual information.
+     *
+     * @param taskId - The ID of the task for which the intent is to be determined.
+     * @param intentDetector - The detector used to classify the intent from available options.
+     *
+     * @returns A promise that resolves to a `FixupIntent` which can be one of the intents like 'add', 'edit', etc.
+     *
+     * @throws
+     * - Will throw an error if no code is selected for fixup.
+     * - Will throw an error if the selected text exceeds the defined maximum limit.
+     *
+     * @todo (umpox): Explore shorter and more efficient ways to detect intent.
+     * Possible methods:
+     * - Input -> Match first word against update|fix|add|delete verbs
+     * - Context -> Infer intent from context, e.g. Current file is a test -> Test intent, Current selection is a comment symbol -> Documentation intent
+     */
+    public async getTaskIntent(taskId: string, intentDetector: IntentDetector): Promise<FixupIntent> {
+        const task = this.tasks.get(taskId)
+        if (!task) {
+            throw new Error('Select some code to fixup.')
+        }
+        const document = await vscode.workspace.openTextDocument(task.fixupFile.uri)
+        const selectedText = document.getText(task.selectionRange)
+        if (truncateText(selectedText, MAX_CURRENT_FILE_TOKENS) !== selectedText) {
+            const msg = "The amount of text selected exceeds Cody's current capacity."
+            throw new Error(msg)
+        }
+
+        if (selectedText.trim().length === 0) {
+            // Nothing selected, assume this is always 'add'.
+            return 'add'
+        }
+
+        const intent = await intentDetector.classifyIntentFromOptions(
+            task.instruction,
+            FixupIntentClassification,
+            'edit'
+        )
+        return intent
+    }
+
+    /**
+     * This function retrieves a "smart" selection a FixupTask.
+     * The idea of a "smart" selection is to look at both the start and end positions of the current selection,
+     * and attempt to expand those positions to encompass more meaningful chunks of code, such as folding regions.
+     *
+     * The function does the following:
+     * 1. Finds the document URI from it's fileName
+     * 2. If the selection starts in a folding range, moves the selection start position back to the start of that folding range.
+     * 3. If the selection ends in a folding range, moves the selection end positionforward to the end of that folding range.
+     *
+     * @returns A Promise that resolves to an `vscode.Range` which represents the combined "smart" selection.
+     */
+    private async getFixupTaskSmartSelection(task: FixupTask, selectionRange: vscode.Range): Promise<vscode.Range> {
+        const fileName = task.fixupFile.uri.fsPath
+        const documentUri = vscode.Uri.file(fileName)
+
+        // Retrieve the start position of the current selection
+        const activeCursorStartPosition = selectionRange.start
+        // If we find a new expanded selection position then we set it as the new start position
+        // and if we don't then we fallback to the original selection made by the user
+        const newSelectionStartingPosition =
+            (await getSmartSelection(documentUri, activeCursorStartPosition.line))?.start || selectionRange.start
+
+        // Retrieve the ending line of the current selection
+        const activeCursorEndPosition = selectionRange.end
+        // If we find a new expanded selection position then we set it as the new ending position
+        // and if we don't then we fallback to the original selection made by the user
+        const newSelectionEndingPosition =
+            (await getSmartSelection(documentUri, activeCursorEndPosition.line))?.end || selectionRange.end
+
+        // Create a new range that starts from the beginning of the folding range at the start position
+        // and ends at the end of the folding range at the end position.
+        return new vscode.Range(
+            newSelectionStartingPosition.line,
+            newSelectionStartingPosition.character,
+            newSelectionEndingPosition.line,
+            newSelectionEndingPosition.character
+        )
+    }
+
     private async applyTask(task: FixupTask): Promise<void> {
         if (task.state !== CodyTaskState.ready) {
             return
@@ -314,12 +401,20 @@ export class FixupController
     }
 
     // Called by the non-stop recipe to gather current state for the task.
-    public async getTaskRecipeData(id: string): Promise<VsCodeFixupTaskRecipeData | undefined> {
+    public async getTaskRecipeData(
+        id: string,
+        options: { enableSmartSelection?: boolean }
+    ): Promise<VsCodeFixupTaskRecipeData | undefined> {
         const task = this.tasks.get(id)
         if (!task) {
             return undefined
         }
         const document = await vscode.workspace.openTextDocument(task.fixupFile.uri)
+        if (options.enableSmartSelection && task.selectionRange) {
+            const newRange = await this.getFixupTaskSmartSelection(task, task.selectionRange)
+            task.selectionRange = newRange
+        }
+
         const precedingText = document.getText(
             new vscode.Range(
                 task.selectionRange.start.translate({ lineDelta: -Math.min(task.selectionRange.start.line, 50) }),