use inline completion context

Author: umpox

vscode/src/autoedits/autoedits-provider.ts

@@ -45,7 +45,7 @@ import {
     autoeditTriggerKind,
     getTimeNowInMillis,
 } from './analytics-logger'
-import { AutoeditCompletionItem } from './autoedit-completion-item'
+import type { AutoeditCompletionItem } from './autoedit-completion-item'
 import { autoeditsOnboarding } from './autoedit-onboarding'
 import { autoeditsProviderConfig } from './autoedits-config'
 import { FilterPredictionBasedOnRecentEdits } from './filter-prediction-edits'
@@ -71,7 +71,11 @@ import type { AutoEditRenderOutput } from './renderer/render-output'
 import { type AutoeditRequestManagerParams, RequestManager } from './request-manager'
 import { shrinkPredictionUntilSuffix } from './shrink-prediction'
 import { SmartThrottleService } from './smart-throttle'
-import { areSameUriDocs, isDuplicatingTextFromRewriteArea } from './utils'
+import {
+    areSameUriDocs,
+    getDocumentTextWithInlineCompletionContext,
+    isDuplicatingTextFromRewriteArea,
+} from './utils'
 
 const AUTOEDIT_CONTEXT_STRATEGY = 'auto-edit'
 const RESET_SUGGESTION_ON_CURSOR_CHANGE_AFTER_INTERVAL_MS = 60 * 1000
@@ -328,27 +332,6 @@ export class AutoeditsProvider implements vscode.InlineCompletionItemProvider, v
         let stopLoading: (() => void) | undefined
         const startedAt = getTimeNowInMillis()
 
-        if (inlineCompletionContext.selectedCompletionInfo !== undefined) {
-            const { range, text } = inlineCompletionContext.selectedCompletionInfo
-            const completion = new AutoeditCompletionItem({
-                id: null,
-                insertText: text,
-                range,
-                withoutCurrentLinePrefix: { insertText: text, range },
-            })
-            // User has a currently selected item in the autocomplete widget.
-            // Instead of attempting to suggest an auto-edit, just show the selected item
-            // as the completion. This is to avoid an undesirable edit conflicting with the acceptance
-            // of the item shown in the widget.
-            // TODO: We should consider the optimal solution here, it may be better to show an
-            // inline completion (not an edit) that includes the currently selected item.
-            return {
-                items: [completion],
-                inlineCompletionItems: [completion],
-                decoratedEditItems: [],
-            }
-        }
-
         try {
             const triggerKind =
                 this.lastManualTriggerTimestamp > performance.now() - 50
@@ -398,6 +381,7 @@ export class AutoeditsProvider implements vscode.InlineCompletionItemProvider, v
                 position,
                 maxPrefixLength: tokensToChars(autoeditsProviderConfig.tokenLimit.prefixTokens),
                 maxSuffixLength: tokensToChars(autoeditsProviderConfig.tokenLimit.suffixTokens),
+                context: inlineCompletionContext,
             })
 
             // Determine the code to replace for this specific request
@@ -409,6 +393,10 @@ export class AutoeditsProvider implements vscode.InlineCompletionItemProvider, v
                 position,
                 tokenBudget: autoeditsProviderConfig.tokenLimit,
             })
+            console.log('UMPOX GOT REQUEST CODE TO REPLACE DATA', {
+                rewrite: requestCodeToReplaceData.codeToRewrite,
+                inlineCompletionContext,
+            })
             const requestId = autoeditAnalyticsLogger.createRequest({
                 startedAt,
                 filePath: getCurrentFilePath(document).toString(),
@@ -478,6 +466,7 @@ export class AutoeditsProvider implements vscode.InlineCompletionItemProvider, v
                 prompt,
                 codeToReplaceData: requestCodeToReplaceData,
                 requestDocContext,
+                inlineCompletionContext,
                 abortSignal,
                 triggerKind,
             })
@@ -503,7 +492,10 @@ export class AutoeditsProvider implements vscode.InlineCompletionItemProvider, v
             const initialPrediction = predictionResult.response.prediction
             const predictionDocContext = predictionResult.docContext
             const predictionCodeToReplaceData = predictionResult.codeToReplaceData
-
+            console.log('UMPOX GOT PREDICTION CODE TO REPLACE DATA', {
+                rewrite: predictionCodeToReplaceData.codeToRewrite,
+                inlineCompletionContext,
+            })
             autoeditAnalyticsLogger.markAsLoaded({
                 requestId,
                 cacheId: predictionResult.cacheId,
@@ -610,6 +602,7 @@ export class AutoeditsProvider implements vscode.InlineCompletionItemProvider, v
                     position,
                     decorationInfo,
                     codeToReplaceData: predictionCodeToReplaceData,
+                    inlineCompletionContext,
                 },
                 this.capabilities
             )
@@ -851,6 +844,7 @@ export class AutoeditsProvider implements vscode.InlineCompletionItemProvider, v
         position,
         codeToReplaceData,
         requestDocContext,
+        inlineCompletionContext,
         prompt,
         abortSignal,
         triggerKind,
@@ -860,6 +854,7 @@ export class AutoeditsProvider implements vscode.InlineCompletionItemProvider, v
         position: vscode.Position
         codeToReplaceData: CodeToReplaceData
         requestDocContext: DocumentContext
+        inlineCompletionContext: vscode.InlineCompletionContext
         prompt: AutoeditsPrompt
         abortSignal: AbortSignal
         triggerKind: AutoeditTriggerKindMetadata
@@ -868,7 +863,7 @@ export class AutoeditsProvider implements vscode.InlineCompletionItemProvider, v
             requestId,
             requestUrl: autoeditsProviderConfig.url,
             documentUri: document.uri.toString(),
-            documentText: document.getText(),
+            documentText: getDocumentTextWithInlineCompletionContext(document, inlineCompletionContext),
             documentVersion: document.version,
             codeToReplaceData,
             requestDocContext,
@@ -897,6 +892,7 @@ export class AutoeditsProvider implements vscode.InlineCompletionItemProvider, v
                             document,
                             codeToReplaceData,
                             requestDocContext,
+                            inlineCompletionContext,
                             position,
                             options: {
                                 hotStreakEnabled: this.features.shouldHotStreak,
@@ -931,6 +927,7 @@ export class AutoeditsProvider implements vscode.InlineCompletionItemProvider, v
                 document,
                 codeToReplaceData,
                 requestDocContext,
+                inlineCompletionContext,
                 position,
                 options: {
                     hotStreakEnabled: this.features.shouldHotStreak,

vscode/src/autoedits/hot-streak/index.ts

@@ -26,6 +26,7 @@ export interface ProcessHotStreakResponsesParams {
     document: vscode.TextDocument
     codeToReplaceData: CodeToReplaceData
     requestDocContext: DocumentContext
+    inlineCompletionContext: vscode.InlineCompletionContext
     position: vscode.Position
     options: {
         // If hot-streak is actually enabled. If it is not, we will not attempt to emit
@@ -53,6 +54,7 @@ export async function* processHotStreakResponses({
     document: originalDocument,
     codeToReplaceData,
     requestDocContext,
+    inlineCompletionContext,
     position,
     options,
 }: ProcessHotStreakResponsesParams): AsyncGenerator<ProcessedHotStreakResponse> {
@@ -113,6 +115,7 @@ export async function* processHotStreakResponses({
                 position: hotStreakPosition,
                 maxPrefixLength: requestDocContext.maxPrefixLength,
                 maxSuffixLength: requestDocContext.maxSuffixLength,
+                context: inlineCompletionContext,
             })
 
             const lengthOfChunk = predictionChunk.range.end.line - predictionChunk.range.start.line - 1

vscode/src/autoedits/prompt/prompt-utils.ts

@@ -236,10 +236,21 @@ export function getCodeToReplaceData(options: CurrentFilePromptOptions): CodeToR
         remainingSuffixChars
     )
 
+    // Get the text from the document for each range
+    const codeToRewritePrefix =
+        document.getText(ranges.codeToRewritePrefix) + (docContext.injectedPrefix || '')
+    const codeToRewriteSuffix = document.getText(ranges.codeToRewriteSuffix)
+
+    console.log('UMPOX MADE CODE TO REWRITE', {
+        codeToRewrite: codeToRewritePrefix + codeToRewriteSuffix,
+        codeToRewritePrefix,
+        codeToRewriteSuffix,
+        prefixinject: docContext.injectedPrefix,
+    })
     return {
-        codeToRewrite: document.getText(ranges.codeToRewrite),
-        codeToRewritePrefix: document.getText(ranges.codeToRewritePrefix),
-        codeToRewriteSuffix: document.getText(ranges.codeToRewriteSuffix),
+        codeToRewrite: codeToRewritePrefix + codeToRewriteSuffix,
+        codeToRewritePrefix,
+        codeToRewriteSuffix,
         prefixInArea: document.getText(ranges.prefixInArea),
         suffixInArea: document.getText(ranges.suffixInArea),
         prefixBeforeArea: prefixBeforeArea.toString(),

vscode/src/autoedits/renderer/manager.ts

@@ -426,7 +426,14 @@ export class AutoEditsDefaultRendererManager
     }
 
     public getRenderOutput(
-        { requestId, prediction, codeToReplaceData, document, position }: GetRenderOutputArgs,
+        {
+            requestId,
+            prediction,
+            codeToReplaceData,
+            document,
+            position,
+            inlineCompletionContext,
+        }: GetRenderOutputArgs,
         capabilities?: AutoeditClientCapabilities
     ): AutoEditRenderOutput {
         const updatedPrediction = adjustPredictionIfInlineCompletionPossible(
@@ -438,6 +445,7 @@ export class AutoEditsDefaultRendererManager
         const { currentLinePrefix, currentLineSuffix } = getCurrentLinePrefixAndSuffix({
             document,
             position,
+            inlineCompletionContext,
         })
         const codeToRewriteAfterCurrentLine = codeToReplaceData.codeToRewriteSuffix.slice(
             currentLineSuffix.length + 1 // Additional char for newline

vscode/src/autoedits/renderer/render-output.ts

@@ -11,6 +11,7 @@ import { AutoeditCompletionItem } from '../autoedit-completion-item'
 import type { AutoeditClientCapabilities } from '../autoedits-provider'
 import { autoeditsOutputChannelLogger } from '../output-channel-logger'
 
+import { getNewLineChar } from '../../completions/text-processing'
 import type { AutoEditDecoration, AutoEditDecorations, DecorationInfo } from './decorators/base'
 import { cssPropertiesToString } from './decorators/utils'
 import { isOnlyAddingTextForModifiedLines, isOnlyRemovingTextForModifiedLines } from './diff-utils'
@@ -26,6 +27,7 @@ export interface GetRenderOutputArgs {
     position: vscode.Position
     decorationInfo: DecorationInfo
     codeToReplaceData: CodeToReplaceData
+    inlineCompletionContext: vscode.InlineCompletionContext
 }
 
 interface NoCompletionRenderOutput {
@@ -190,6 +192,7 @@ export class AutoEditsRenderOutput {
         args: GetRenderOutputArgs,
         capabilities: AutoeditClientCapabilities
     ): CompletionRenderOutput | CompletionWithDecorationsRenderOutput | null {
+        console.log('UMPOX CODE TO REWIRTE', args.codeToReplaceData.codeToRewrite)
         const completions = this.tryMakeInlineCompletions(args)
         if (!completions) {
             // Cannot render a completion
@@ -237,27 +240,67 @@ export class AutoEditsRenderOutput {
         return null
     }
 
+    private getInsertionPositionForInlineCompletion(
+        position: vscode.Position,
+        inlineCompletionContext: vscode.InlineCompletionContext
+    ) {
+        const selectedCompletion = inlineCompletionContext.selectedCompletionInfo
+        if (!selectedCompletion) {
+            // No selected completion so no need to modify the cursor position
+            return position
+        }
+
+        const newLineChar = getNewLineChar(selectedCompletion.text)
+        const linesToInsertFromSelectedCompletion = selectedCompletion.text.split(newLineChar)
+
+        if (linesToInsertFromSelectedCompletion.length === 1) {
+            return new vscode.Position(
+                selectedCompletion.range.start.line,
+                selectedCompletion.range.start.character + selectedCompletion.text.length
+            )
+        }
+
+        const lastInsertionLine =
+            linesToInsertFromSelectedCompletion[linesToInsertFromSelectedCompletion.length - 1]
+        return new vscode.Position(
+            selectedCompletion.range.start.line + linesToInsertFromSelectedCompletion.length - 1,
+            lastInsertionLine.length
+        )
+    }
+
     private tryMakeInlineCompletions({
         requestId,
         position,
         prediction,
         document,
         decorationInfo,
+        inlineCompletionContext,
     }: GetRenderOutputArgs): {
         type: 'full' | 'partial'
         inlineCompletionItems: AutoeditCompletionItem[]
         updatedDecorationInfo: DecorationInfo
         updatedPrediction: string
     } | null {
+        const otherPosition = this.getInsertionPositionForInlineCompletion(
+            position,
+            inlineCompletionContext
+        )
         const { insertText, usedChangeIds } = getCompletionText({
             prediction,
             cursorPosition: position,
             decorationInfo,
         })
 
+        const other = getCompletionText({
+            prediction,
+            cursorPosition: otherPosition,
+            decorationInfo,
+        })
+
         const { currentLinePrefix, currentLineSuffix } = getCurrentLinePrefixAndSuffix({
             document,
             position,
+            inlineCompletionContext,
         })
 
         if (insertText.length === 0) {
@@ -272,6 +315,15 @@ export class AutoEditsRenderOutput {
         }
 
         const completionText = currentLinePrefix + insertText
+        console.log('UMPOX GOT COMPLETION TEXT', {
+            used: {
+                completionText,
+                currentLinePrefix,
+                inlineCompletionContext,
+                insertText,
+            },
+            other,
+        })
         const inlineCompletionItems = [
             new AutoeditCompletionItem({
                 id: requestId,

vscode/src/autoedits/utils.ts

@@ -161,3 +161,23 @@ export function clip(line: number, min: number, max: number) {
 export function areSameUriDocs(a?: vscode.TextDocument, b?: vscode.TextDocument): boolean {
     return Boolean(a && b && a.uri.toString() === b.uri.toString())
 }
+
+export function getDocumentTextWithInlineCompletionContext(
+    document: vscode.TextDocument,
+    inlineCompletionContext: vscode.InlineCompletionContext
+): string {
+    const selectedCompletionContext = inlineCompletionContext.selectedCompletionInfo
+    if (!selectedCompletionContext) {
+        return document.getText()
+    }
+
+    const documentText = document.getText()
+    const startOffset = document.offsetAt(selectedCompletionContext.range.start)
+    const endOffset = document.offsetAt(selectedCompletionContext.range.end)
+
+    return (
+        documentText.substring(0, startOffset) +
+        selectedCompletionContext.text +
+        documentText.substring(endOffset)
+    )
+}

vscode/src/completions/get-current-doc-context.ts

@@ -352,15 +352,30 @@ function getLinesContext(params: GetLinesContextParams): LinesContext {
 export function getCurrentLinePrefixAndSuffix({
     document,
     position,
+    inlineCompletionContext,
 }: {
     document: vscode.TextDocument
     position: vscode.Position
+    inlineCompletionContext: vscode.InlineCompletionContext
 }) {
     const currentLineStartPosition = new vscode.Position(position.line, 0)
-    const currentLinePrefix = document.getText(new vscode.Range(currentLineStartPosition, position))
-
     const currentLineEndPosition = document.lineAt(position.line).range.end
-    const currentLineSuffix = document.getText(new vscode.Range(position, currentLineEndPosition))
 
+    const selectedCompletion = inlineCompletionContext.selectedCompletionInfo
+    if (selectedCompletion) {
+        console.log('UMPOX GOT SELECTION COMPLETION')
+        // The selected completion may replace existing characters in the document
+        // Use this so we get an accurate prefix, we can append the completion text afterwards
+        // const adjustedPosition = selectedCompletion.range.end
+        const currentLinePrefix =
+            document.getText(
+                new vscode.Range(currentLineStartPosition, selectedCompletion.range.start)
+            ) + selectedCompletion.text
+        const currentLineSuffix = document.getText(new vscode.Range(position, currentLineEndPosition))
+        return { currentLinePrefix, currentLineSuffix }
+    }
+
+    const currentLinePrefix = document.getText(new vscode.Range(currentLineStartPosition, position))
+    const currentLineSuffix = document.getText(new vscode.Range(position, currentLineEndPosition))
     return { currentLinePrefix, currentLineSuffix }
 }