Edit: Support trimming active selections to a tree-sitter node

vscode/src/commands/execute/doc.ts

@@ -8,9 +8,29 @@ import { getEditor } from '../../editor/active-editor'
 import type { EditCommandResult } from '../../main'
 import type { CodyCommandArgs } from '../types'
 
-import { getEditLineSelection } from '../../edit/utils/edit-selection'
+import {
+    getEditAdjustedUserSelection,
+    getEditDefaultProvidedRange,
+} from '../../edit/utils/edit-selection'
 import { execQueryWrapper } from '../../tree-sitter/query-sdk'
 
+function getDefaultDocumentableRange(editor: vscode.TextEditor): {
+    range?: vscode.Range
+    insertionPoint?: vscode.Position
+} {
+    const defaultRange = getEditDefaultProvidedRange(editor.document, editor.selection)
+
+    if (defaultRange) {
+        return {
+            range: defaultRange,
+            insertionPoint: defaultRange.start,
+        }
+    }
+
+    // No usable selection, fallback to expanding the range to the nearest block.
+    return {}
+}
+
 /**
  * Gets the symbol range and preferred insertion point for documentation
  * at the given document position.
@@ -27,39 +47,49 @@ function getDocumentableRange(editor: vscode.TextEditor): {
     range?: vscode.Range
     insertionPoint?: vscode.Position
 } {
-    const { document, selection } = editor
-    if (!selection.isEmpty) {
-        const lineSelection = getEditLineSelection(editor.document, editor.selection)
-        // The user has made an active selection, use that as the documentable range
-        return {
-            range: lineSelection,
-            insertionPoint: lineSelection.start,
-        }
-    }
+    const { document } = editor
+    const adjustedSelection = getEditAdjustedUserSelection(document, editor.selection)
 
     /**
      * Attempt to get the range of a documentable symbol at the current cursor position.
      * If present, use this for the edit instead of expanding the range to the nearest block.
      */
     const [documentableNode] = execQueryWrapper({
         document,
-        position: editor.selection.active,
+        position: adjustedSelection.start,
         queryWrapper: 'getDocumentableNode',
     })
+
     if (!documentableNode) {
-        return {}
+        return getDefaultDocumentableRange(editor)
     }
 
     const { range: documentableRange, insertionPoint: documentableInsertionPoint } = documentableNode
     if (!documentableRange) {
-        // No user-provided selection, no documentable range found.
+        // No documentable range found.
         // Fallback to expanding the range to the nearest block.
-        return {}
+        return getDefaultDocumentableRange(editor)
     }
 
     const {
         node: { startPosition, endPosition },
     } = documentableRange
+    const range = new vscode.Range(
+        startPosition.row,
+        startPosition.column,
+        endPosition.row,
+        endPosition.column
+    )
+
+    // If the users' adjusted selection aligns with the start of the node and is contained within the node,
+    // It is probable that the user would benefit from expanding to this node completely
+    const selectionMatchesNode =
+        adjustedSelection.start.isEqual(range.start) && range.contains(adjustedSelection.end)
+    if (!selectionMatchesNode && !editor.selection.isEmpty) {
+        // We found a documentable range, but the users' adjusted selection does not match it.
+        // We have to use the users' selection here, as it's possible they do not want the documentable node.
+        return getDefaultDocumentableRange(editor)
+    }
 
     const insertionPoint = documentableInsertionPoint
         ? new vscode.Position(documentableInsertionPoint.node.startPosition.row + 1, 0)
@@ -69,12 +99,7 @@ function getDocumentableRange(editor: vscode.TextEditor): {
           )
 
     return {
-        range: new vscode.Range(
-            startPosition.row,
-            startPosition.column,
-            endPosition.row,
-            endPosition.column
-        ),
+        range,
         insertionPoint,
     }
 }

vscode/src/commands/execute/test-edit.ts

@@ -5,7 +5,10 @@ import type { URI } from 'vscode-uri'
 
 import { defaultCommands } from '.'
 import { type ExecuteEditArguments, executeEdit } from '../../edit/execute'
-import { getEditLineSelection } from '../../edit/utils/edit-selection'
+import {
+    getEditAdjustedUserSelection,
+    getEditDefaultProvidedRange,
+} from '../../edit/utils/edit-selection'
 import { getEditor } from '../../editor/active-editor'
 import type { EditCommandResult } from '../../main'
 import { execQueryWrapper } from '../../tree-sitter/query-sdk'
@@ -20,12 +23,8 @@ import { isTestFileForOriginal } from '../utils/test-commands'
  * using a tree-sitter query. If found, returns the range for the symbol.
  */
 function getTestableRange(editor: vscode.TextEditor): vscode.Range | undefined {
-    const { document, selection } = editor
-    if (!selection.isEmpty) {
-        const lineSelection = getEditLineSelection(editor.document, editor.selection)
-        // The user has made an active selection, use that as the testable range
-        return lineSelection
-    }
+    const { document } = editor
+    const adjustedSelection = getEditAdjustedUserSelection(document, editor.selection)
 
     /**
      * Attempt to get the range of a testable symbol at the current cursor position.
@@ -37,21 +36,37 @@ function getTestableRange(editor: vscode.TextEditor): vscode.Range | undefined {
         queryWrapper: 'getTestableNode',
     })
     if (!testableNode) {
-        return undefined
+        return getEditDefaultProvidedRange(editor.document, editor.selection)
     }
 
     const { range: testableRange } = testableNode
     if (!testableRange) {
         // No user-provided selection, no testable range found.
         // Fallback to expanding the range to the nearest block.
-        return undefined
+        return getEditDefaultProvidedRange(editor.document, editor.selection)
     }
 
     const {
         node: { startPosition, endPosition },
     } = testableRange
+    const range = new vscode.Range(
+        startPosition.row,
+        startPosition.column,
+        endPosition.row,
+        endPosition.column
+    )
+
+    // If the users' adjusted selection aligns with the start of the node and is contained within the node,
+    // It is probable that the user would benefit from expanding to this node completely
+    const selectionMatchesNode =
+        adjustedSelection.start.isEqual(range.start) && range.contains(adjustedSelection.end)
+    if (!selectionMatchesNode && !editor.selection.isEmpty) {
+        // We found a testable range, but the users' adjusted selection does not match it.
+        // We have to use the users' selection here, as it's possible they do not want the testable node.
+        return getEditDefaultProvidedRange(editor.document, editor.selection)
+    }
 
-    return new vscode.Range(startPosition.row, startPosition.column, endPosition.row, endPosition.column)
+    return range
 }
 
 /**

vscode/src/edit/utils/edit-selection.test.ts

@@ -0,0 +1,34 @@
+import { describe, expect, it } from 'vitest'
+import { Position, Selection } from 'vscode'
+import { document } from '../../completions/test-helpers'
+import { getEditAdjustedUserSelection } from './edit-selection'
+
+describe('getEditAdjustedUserSelection', () => {
+    it('should return the original selection if it is empty', () => {
+        const doc = document('Hello world!')
+        const selection = new Selection(new Position(0, 0), new Position(0, 0))
+        const result = getEditAdjustedUserSelection(doc, selection)
+        expect(result).toEqual(selection)
+    })
+
+    it('should trim whitespace from the start and end of a selection', () => {
+        const doc = document('  Hello World  \n\n')
+        const selection = new Selection(doc.positionAt(0), doc.positionAt(doc.getText().length))
+        const updatedSelectionRange = getEditAdjustedUserSelection(doc, selection)
+        expect(doc.getText(updatedSelectionRange)).toBe('Hello World')
+    })
+
+    it('should expand selection to include full text of a line', () => {
+        const doc = document('  Hello World  \n\n')
+        const selection = new Selection(doc.positionAt(4), doc.positionAt(10)) // partial selection on "llo wo"
+        const updatedSelectionRange = getEditAdjustedUserSelection(doc, selection)
+        expect(doc.getText(updatedSelectionRange)).toBe('Hello World')
+    })
+
+    it('should handle selections that span multiple lines', () => {
+        const doc = document('  Hello\n  World  \n')
+        const selection = new Selection(doc.positionAt(0), doc.positionAt(doc.getText().length))
+        const updatedSelectionRange = getEditAdjustedUserSelection(doc, selection)
+        expect(doc.getText(updatedSelectionRange)).toBe('Hello\n  World')
+    })
+})

vscode/src/edit/utils/edit-selection.ts

@@ -77,3 +77,53 @@ export function getEditLineSelection(
     const endChar = document.lineAt(selection.end.line).text.length
     return new vscode.Range(selection.start.line, startChar, selection.end.line, endChar)
 }
+
+/**
+ * Given a provided user selection, adjust this to match an more optimal, realistic selection.
+ * This helps reduce the chance of user errors where an Edit may be triggered with either:
+ * - Too little text selected (e.g. the user selected all but a few chars of a function)
+ * - Too much text selected (e.g. a user selected a function, and some extra lines of whitespace)
+ *
+ * We expand, and then trim the selection to match only the intended non-whitespace characters
+ */
+export function getEditAdjustedUserSelection(
+    document: vscode.TextDocument,
+    selection: vscode.Selection
+): vscode.Range {
+    if (selection.isEmpty) {
+        // No selection provided, do nothing
+        return selection
+    }
+
+    // Expand the selection to include all non-whitespace characters from the selected lines
+    const lineSelection = getEditLineSelection(document, selection)
+    const text = document.getText(lineSelection)
+
+    // Trim any additional whitespace characters (e.g. additional lines) from the selection
+    const trimmedText = text.trim()
+    const startOffset = document.offsetAt(lineSelection.start) + text.indexOf(trimmedText)
+    const endOffset = startOffset + trimmedText.length
+
+    return new vscode.Selection(document.positionAt(startOffset), document.positionAt(endOffset))
+}
+
+/**
+ * Returns a range that represents the default provided range for an edit operation.
+ *
+ * If the user has made an active selection, the function will return the line-expanded
+ * selection range. Otherwise, it will return `undefined`, indicating that no default
+ * range is available.
+ */
+export function getEditDefaultProvidedRange(
+    document: vscode.TextDocument,
+    selection: vscode.Selection
+): vscode.Range | undefined {
+    if (!selection.isEmpty) {
+        // User has made an active selection.
+        // We have to use their selection as the intended range.
+        return getEditLineSelection(document, selection)
+    }
+
+    // No usable selection, fallback to the range expansion behaviour for Edit
+    return
+}

vscode/src/testutils/mocks.ts

@@ -696,6 +696,7 @@ export const vsCodeMocks = {
     FileSystemError,
     FileType,
     Range,
+    Selection,
     Position,
     InlineCompletionItem,
     EventEmitter,