feat(amazonq): prefetch wip

server/aws-lsp-codewhisperer/src/language-server/inline-completion/codeWhispererServer.ts

@@ -19,6 +19,7 @@ import { autoTrigger, triggerType } from './auto-trigger/autoTrigger'
 import {
     CodeWhispererServiceToken,
     GenerateSuggestionsRequest,
+    getFileContext,
     Suggestion,
     SuggestionType,
 } from '../../shared/codeWhispererService'
@@ -57,38 +58,6 @@ import { RecentEditTracker, RecentEditTrackerDefaultConfig } from './codeEditTra
 const EMPTY_RESULT = { sessionId: '', items: [] }
 export const CONTEXT_CHARACTERS_LIMIT = 10240
 
-// Both clients (token, sigv4) define their own types, this return value needs to match both of them.
-const getFileContext = (params: {
-    textDocument: TextDocument
-    position: Position
-    inferredLanguageId: CodewhispererLanguage
-}): {
-    filename: string
-    programmingLanguage: {
-        languageName: CodewhispererLanguage
-    }
-    leftFileContent: string
-    rightFileContent: string
-} => {
-    const left = params.textDocument.getText({
-        start: { line: 0, character: 0 },
-        end: params.position,
-    })
-    const right = params.textDocument.getText({
-        start: params.position,
-        end: params.textDocument.positionAt(params.textDocument.getText().length),
-    })
-
-    return {
-        filename: params.textDocument.uri,
-        programmingLanguage: {
-            languageName: params.inferredLanguageId,
-        },
-        leftFileContent: left,
-        rightFileContent: right,
-    }
-}
-
 const emitServiceInvocationTelemetry = (telemetry: Telemetry, session: CodeWhispererSession) => {
     const duration = new Date().getTime() - session.startTime
     const data: CodeWhispererServiceInvocationEvent = {
@@ -458,7 +427,8 @@ export const CodewhispererServerFactory =
                     if (extraContext) {
                         requestContext.fileContext.leftFileContent = extraContext + '\n' + requestContext.fileContext.leftFileContent
                     }
-                    return codeWhispererService.generateSuggestionsAndPrefetch({
+
+                    return codeWhispererService.generateSuggestionsAndPrefetch(textDocument, {
                         ...requestContext,
                         predictionTypes : ['EDITS'],
                         fileContext: {
@@ -651,8 +621,8 @@ export const CodewhispererServerFactory =
                 firstCompletionDisplayLatency,
                 totalSessionDisplayTime,
                 typeaheadLength,
-                addedCharacterCount,
-                deletedCharacterCount,
+                // addedCharacterCount,
+                // deletedCharacterCount,
             } = params
 
             const session = sessionManager.getSessionById(sessionId)
@@ -704,8 +674,8 @@ export const CodewhispererServerFactory =
                 telemetryService,
                 session,
                 timeSinceLastUserModification,
-                addedCharacterCount,
-                deletedCharacterCount,
+                // addedCharacterCount,
+                // deletedCharacterCount,
                 streakLength
             )
         }

server/aws-lsp-codewhisperer/src/language-server/inline-completion/diffUtils.ts

@@ -5,6 +5,7 @@
 
 import * as diff from 'diff'
 import { CodeWhispererSupplementalContext, CodeWhispererSupplementalContextItem } from '../../shared/models/model'
+import { Position } from 'vscode-languageserver-textdocument'
 
 const supplementalContextMaxTotalLength: number = 8192
 const charactersLimit: number = 10000
@@ -177,3 +178,188 @@ function trimSupplementalContexts(
     const trimmedContexts = result.slice(0, i)
     return trimmedContexts
 }
+
+/** src: https://github.com/aws/aws-toolkit-vscode/blob/3921457b0a2094b831beea0d66cc2cbd2a833890/packages/amazonq/src/app/inline/EditRendering/diffUtils.ts#L18
+ * Apply a unified diff to original code to generate modified code
+ * @param originalCode The original code as a string
+ * @param unifiedDiff The unified diff content
+ * @returns The modified code after applying the diff
+ */
+export function applyUnifiedDiff(
+    docText: string,
+    unifiedDiff: string
+): { newCode: string; addedCharacterCount: number; deletedCharacterCount: number } {
+    try {
+        const { addedCharacterCount, deletedCharacterCount } = getAddedAndDeletedCharCount(unifiedDiff)
+        // First try the standard diff package
+        try {
+            const result = diff.applyPatch(docText, unifiedDiff)
+            if (result !== false) {
+                return {
+                    newCode: result,
+                    addedCharacterCount: addedCharacterCount,
+                    deletedCharacterCount: deletedCharacterCount,
+                }
+            }
+        } catch (error) {}
+
+        // Parse the unified diff to extract the changes
+        const diffLines = unifiedDiff.split('\n')
+        let result = docText
+
+        // Find all hunks in the diff
+        const hunkStarts = diffLines
+            .map((line, index) => (line.startsWith('@@ ') ? index : -1))
+            .filter(index => index !== -1)
+
+        // Process each hunk
+        for (const hunkStart of hunkStarts) {
+            // Parse the hunk header
+            const hunkHeader = diffLines[hunkStart]
+            const match = hunkHeader.match(/@@ -(\d+),(\d+) \+(\d+),(\d+) @@/)
+
+            if (!match) {
+                continue
+            }
+
+            const oldStart = parseInt(match[1])
+            const oldLines = parseInt(match[2])
+
+            // Extract the content lines for this hunk
+            let i = hunkStart + 1
+            const contentLines = []
+            while (i < diffLines.length && !diffLines[i].startsWith('@@')) {
+                contentLines.push(diffLines[i])
+                i++
+            }
+
+            // Build the old and new text
+            let oldText = ''
+            let newText = ''
+
+            for (const line of contentLines) {
+                if (line.startsWith('-')) {
+                    oldText += line.substring(1) + '\n'
+                } else if (line.startsWith('+')) {
+                    newText += line.substring(1) + '\n'
+                } else if (line.startsWith(' ')) {
+                    oldText += line.substring(1) + '\n'
+                    newText += line.substring(1) + '\n'
+                }
+            }
+
+            // Remove trailing newline if it was added
+            oldText = oldText.replace(/\n$/, '')
+            newText = newText.replace(/\n$/, '')
+
+            // Find the text to replace in the document
+            const docLines = docText.split('\n')
+            const startLine = oldStart - 1 // Convert to 0-based
+            const endLine = startLine + oldLines
+
+            // Extract the text that should be replaced
+            const textToReplace = docLines.slice(startLine, endLine).join('\n')
+
+            // Replace the text
+            result = result.replace(textToReplace, newText)
+        }
+        return {
+            newCode: result,
+            addedCharacterCount: addedCharacterCount,
+            deletedCharacterCount: deletedCharacterCount,
+        }
+    } catch (error) {
+        return {
+            newCode: docText, // Return original text if all methods fail
+            addedCharacterCount: 0,
+            deletedCharacterCount: 0,
+        }
+    }
+}
+
+// src https://github.com/aws/aws-toolkit-vscode/blob/3921457b0a2094b831beea0d66cc2cbd2a833890/packages/amazonq/src/app/inline/EditRendering/diffUtils.ts#L147
+export function getAddedAndDeletedCharCount(diff: string): {
+    addedCharacterCount: number
+    deletedCharacterCount: number
+} {
+    let addedCharacterCount = 0
+    let deletedCharacterCount = 0
+    let i = 0
+    const lines = diff.split('\n')
+    while (i < lines.length) {
+        const line = lines[i]
+        if (line.startsWith('+') && !line.startsWith('+++')) {
+            addedCharacterCount += line.length - 1
+        } else if (line.startsWith('-') && !line.startsWith('---')) {
+            const removedLine = line.substring(1)
+            deletedCharacterCount += removedLine.length
+
+            // Check if this is a modified line rather than a pure deletion
+            const nextLine = lines[i + 1]
+            if (nextLine && nextLine.startsWith('+') && !nextLine.startsWith('+++') && nextLine.includes(removedLine)) {
+                // This is a modified line, not a pure deletion
+                // We've already counted the deletion, so we'll just increment i to skip the next line
+                // since we'll process the addition on the next iteration
+                i += 1
+            }
+        }
+        i += 1
+    }
+    return {
+        addedCharacterCount,
+        deletedCharacterCount,
+    }
+}
+
+/**
+ * Calculates the end position of the actual edited content by finding the last changed part
+ */
+export function getEndOfEditPosition(originalCode: string, newCode: string): Position {
+    const changes = diff.diffLines(originalCode, newCode)
+    let lineOffset = 0
+
+    // Track the end position of the last added chunk
+    let lastChangeEndLine = 0
+    let lastChangeEndColumn = 0
+    let foundAddedContent = false
+
+    for (const part of changes) {
+        if (part.added) {
+            foundAddedContent = true
+
+            // Calculate lines in this added part
+            const lines = part.value.split('\n')
+            const linesCount = lines.length
+
+            // Update position to the end of this added chunk
+            lastChangeEndLine = lineOffset + linesCount - 1
+
+            // Get the length of the last line in this added chunk
+            lastChangeEndColumn = lines[linesCount - 1].length
+        }
+
+        // Update line offset (skip removed parts)
+        if (!part.removed) {
+            // Safely calculate line count from the part's value
+            const partLineCount = part.value.split('\n').length
+            lineOffset += partLineCount - 1
+        }
+    }
+
+    // If we found added content, return position at the end of the last addition
+    if (foundAddedContent) {
+        return {
+            line: lastChangeEndLine,
+            character: lastChangeEndColumn, // ?????????????????? is this correct?????????????
+        }
+        // return new vscode.Position(lastChangeEndLine, lastChangeEndColumn)
+    }
+
+    // Fallback to current cursor position if no changes were found
+    // const editor = vscode.window.activeTextEditor
+    // return editor ? editor.selection.active : new Position(0, 0)
+    return {
+        line: 0,
+        character: 0,
+    }
+}