feat(auto-edits): fix the partial decoration issue when not enough lines in the editor

vscode/src/autoedits/analytics-logger/analytics-logger.ts

@@ -86,7 +86,7 @@ const validRequestTransitions = {
     started: ['contextLoaded', 'discarded'],
     contextLoaded: ['loaded', 'discarded'],
     loaded: ['postProcessed', 'discarded'],
-    postProcessed: ['suggested'],
+    postProcessed: ['suggested', 'discarded'],
     suggested: ['read', 'accepted', 'rejected'],
     read: ['accepted', 'rejected'],
     accepted: [],
@@ -173,6 +173,8 @@ export const autoeditDiscardReason = {
     suffixOverlap: 5,
     emptyPredictionAfterInlineCompletionExtraction: 6,
     noActiveEditor: 7,
+    conflictingDecorationWithEdits: 8,
+    noEnoughLinesEditor: 9,
 } as const
 
 /** We use numeric keys to send these to the analytics backend */

vscode/src/autoedits/renderer/decorators/base.ts

@@ -31,6 +31,16 @@ export interface AutoEditsDecorator extends vscode.Disposable {
      *        and how they should be decorated in the editor.
      */
     setDecorations(decorationInfo: DecorationInfo): void
+
+    /**
+     * Checks if the decorator can render decorations for the given decoration information.
+     *
+     * This method verifies if the current editor state allows for the decorations to be
+     * rendered properly. Some conditions that might prevent rendering include:
+     * - Insufficient lines in the editor
+     * @returns true if decorations can be rendered, false otherwise
+     */
+    canRenderDecoration(decorationInfo: DecorationInfo): boolean
 }
 
 /**

vscode/src/autoedits/renderer/decorators/default-decorator.ts

@@ -11,15 +11,40 @@ interface AddedLinesDecorationInfo {
     lineText: string
 }
 
+interface DiffDecorationAddedLinesInfo {
+    /** Information about lines that have been added */
+    addedLinesDecorationInfo: AddedLinesDecorationInfo[]
+    /** Starting line number for the decoration */
+    startLine: number
+    /** Column position for the replacement text */
+    replacerCol: number
+}
+
+/**
+ * Information about diff decorations to be applied to lines in the editor
+ */
+interface DiffDecorationInfo {
+    /** Ranges of text that have been removed */
+    removedRangesInfo: vscode.Range[]
+    /** Information about lines that have been added */
+    addedLinesInfo?: DiffDecorationAddedLinesInfo
+}
+
 export class DefaultDecorator implements AutoEditsDecorator {
     private readonly decorationTypes: vscode.TextEditorDecorationType[]
+    private readonly editor: vscode.TextEditor
+
+    // Decoration types
     private readonly removedTextDecorationType: vscode.TextEditorDecorationType
     private readonly modifiedTextDecorationType: vscode.TextEditorDecorationType
     private readonly suggesterType: vscode.TextEditorDecorationType
-    private readonly hideRemainderDecorationType: vscode.TextEditorDecorationType
     private readonly addedLinesDecorationType: vscode.TextEditorDecorationType
     private readonly insertMarkerDecorationType: vscode.TextEditorDecorationType
-    private readonly editor: vscode.TextEditor
+
+    /**
+     * Pre-computed information about diff decorations to be applied to lines in the editor.
+     */
+    private diffDecorationInfo: DiffDecorationInfo | undefined
 
     constructor(editor: vscode.TextEditor) {
         this.editor = editor
@@ -34,9 +59,6 @@ export class DefaultDecorator implements AutoEditsDecorator {
             before: { color: GHOST_TEXT_COLOR },
             after: { color: GHOST_TEXT_COLOR },
         })
-        this.hideRemainderDecorationType = vscode.window.createTextEditorDecorationType({
-            opacity: '0',
-        })
         this.addedLinesDecorationType = vscode.window.createTextEditorDecorationType({
             backgroundColor: 'red', // SENTINEL (should not actually appear)
             before: {
@@ -55,12 +77,28 @@ export class DefaultDecorator implements AutoEditsDecorator {
             this.removedTextDecorationType,
             this.modifiedTextDecorationType,
             this.suggesterType,
-            this.hideRemainderDecorationType,
             this.addedLinesDecorationType,
             this.insertMarkerDecorationType,
         ]
     }
 
+    public canRenderDecoration(decorationInfo: DecorationInfo): boolean {
+        if (!this.diffDecorationInfo) {
+            this.diffDecorationInfo = this.getDiffDecorationsInfo(decorationInfo)
+        }
+        const { addedLinesInfo } = this.diffDecorationInfo
+        if (addedLinesInfo) {
+            // Check if there are enough lines in the editor to render the diff decorations
+            if (
+                addedLinesInfo.startLine + addedLinesInfo.addedLinesDecorationInfo.length >
+                this.editor.document.lineCount
+            ) {
+                return false
+            }
+        }
+        return true
+    }
+
     private clearDecorations(): void {
         for (const decorationType of this.decorationTypes) {
             this.editor.setDecorations(decorationType, [])
@@ -90,6 +128,27 @@ export class DefaultDecorator implements AutoEditsDecorator {
     }
 
     private renderDiffDecorations(decorationInfo: DecorationInfo): void {
+        if (!this.diffDecorationInfo) {
+            this.diffDecorationInfo = this.getDiffDecorationsInfo(decorationInfo)
+        }
+        this.editor.setDecorations(
+            this.modifiedTextDecorationType,
+            this.diffDecorationInfo.removedRangesInfo
+        )
+        const addedLinesInfo = this.diffDecorationInfo.addedLinesInfo
+
+        if (!addedLinesInfo) {
+            return
+        }
+
+        this.renderAddedLinesDecorations(
+            addedLinesInfo.addedLinesDecorationInfo,
+            addedLinesInfo.startLine,
+            addedLinesInfo.replacerCol
+        )
+    }
+
+    private getDiffDecorationsInfo(decorationInfo: DecorationInfo): DiffDecorationInfo {
         const { modifiedLines, addedLines, unchangedLines } = decorationInfo
 
         // Display the removed range decorations
@@ -119,7 +178,6 @@ export class DefaultDecorator implements AutoEditsDecorator {
                 })
             }
         }
-        this.editor.setDecorations(this.modifiedTextDecorationType, removedRanges)
 
         // Handle fully added lines
         for (const addedLine of addedLines) {
@@ -151,13 +209,22 @@ export class DefaultDecorator implements AutoEditsDecorator {
         // Sort addedLinesInfo by line number in ascending order
         addedLinesInfo.sort((a, b) => a.afterLine - b.afterLine)
         if (addedLinesInfo.length === 0) {
-            return
+            return { removedRangesInfo: removedRanges }
         }
-        // todo (hitesh): handle case when too many lines to fit in the editor
-        const oldLines = addedLinesInfo.map(info => this.editor.document.lineAt(info.afterLine))
+        const oldLines = addedLinesInfo
+            .filter(info => info.afterLine < this.editor.document.lineCount)
+            .map(info => this.editor.document.lineAt(info.afterLine))
         const replacerCol = Math.max(...oldLines.map(line => line.range.end.character))
         const startLine = Math.min(...oldLines.map(line => line.lineNumber))
-        this.renderAddedLinesDecorations(addedLinesInfo, startLine, replacerCol)
+
+        return {
+            removedRangesInfo: removedRanges,
+            addedLinesInfo: {
+                addedLinesDecorationInfo: addedLinesInfo,
+                startLine,
+                replacerCol,
+            },
+        }
     }
 
     private renderAddedLinesDecorations(

vscode/src/autoedits/renderer/decorators/inline-diff-decorator.ts

@@ -19,6 +19,11 @@ export class InlineDiffDecorator implements vscode.Disposable, AutoEditsDecorato
         this.editor.setDecorations(this.addedTextDecorationType, added)
     }
 
+    public canRenderDecoration(decorationInfo: DecorationInfo): boolean {
+        // Inline decorator can render any decoration, so it should always return true.
+        return true
+    }
+
     /**
      * Process modified lines to create decorations for inserted and deleted text within those lines.
      */

vscode/src/autoedits/renderer/manager.ts

@@ -7,7 +7,11 @@ import {
     getLatestVisibilityContext,
     isCompletionVisible,
 } from '../../completions/is-completion-visible'
-import { type AutoeditRequestID, autoeditAnalyticsLogger } from '../analytics-logger'
+import {
+    type AutoeditRequestID,
+    autoeditAnalyticsLogger,
+    autoeditDiscardReason,
+} from '../analytics-logger'
 import { autoeditsProviderConfig } from '../autoedits-config'
 import { autoeditsOutputChannelLogger } from '../output-channel-logger'
 import type { CodeToReplaceData } from '../prompt/prompt-utils'
@@ -174,16 +178,34 @@ export class AutoEditsDefaultRendererManager implements AutoEditsRendererManager
         await this.rejectActiveEdit()
 
         const request = autoeditAnalyticsLogger.getRequest(requestId)
+        if (!request) {
+            return
+        }
+        if (this.hasConflictingDecorations(request.document, request.codeToReplaceData.range)) {
+            autoeditAnalyticsLogger.markAsDiscarded({
+                requestId,
+                discardReason: autoeditDiscardReason.conflictingDecorationWithEdits,
+            })
+            return
+        }
+
+        this.decorator = this.createDecorator(vscode.window.activeTextEditor!)
         if (
-            !request ||
-            this.hasConflictingDecorations(request.document, request.codeToReplaceData.range)
+            'decorationInfo' in request &&
+            request.decorationInfo &&
+            !this.decorator.canRenderDecoration(request.decorationInfo)
         ) {
+            // If the decorator cannot render the decoration properly, dispose of it and return early.
+            this.decorator.dispose()
+            this.decorator = null
+            autoeditAnalyticsLogger.markAsDiscarded({
+                requestId,
+                discardReason: autoeditDiscardReason.noEnoughLinesEditor,
+            })
             return
         }
 
         this.activeRequestId = requestId
-        this.decorator = this.createDecorator(vscode.window.activeTextEditor!)
-
         autoeditAnalyticsLogger.markAsSuggested(requestId)
 
         // Clear any existing timeouts, only one suggestion can be shown at a time