src/extension.ts

@@ -61,6 +61,7 @@ const previewDecorationType = vscode.window.createTextEditorDecorationType({
     textDecoration: 'none; display: none;', // Hide the original text
 });
 
+
 async function autocompleteCommand(textEditor: vscode.TextEditor, cancellationToken?: vscode.CancellationToken) {
     const document = textEditor.document;
     const position = textEditor.selection.active;
@@ -113,82 +114,93 @@ async function autocompleteCommand(textEditor: vscode.TextEditor, cancellationTo
                 let completionText = response.data.response;
                 completionText = completionText.replace(/<fim_middle>|<fim_suffix>|<fim_prefix>/g, '').trim();
 
-                // Split the completion text by new lines
+                let previewInserted = false;
-                const newLines = completionText.split('\n');
+                let originalContent: string;
+                let previewStartLine: number;
 
-                // Calculate the number of new lines in the completion, plus 2 extra lines
+                // Store the original content and insert preview
-                const totalNewLines = newLines.length + 2;
+const storeAndInsertPreview = async () => {
+    previewStartLine = startLine;
+    const endLine = document.lineCount - 1;
+    const endCharacter = document.lineAt(endLine).text.length;
+    const fullRange = new vscode.Range(startLine, 0, endLine, endCharacter);
+    originalContent = document.getText(fullRange);
 
-                // Create preview decorations and insert new lines
+    const previewContent = completionText + '\n'.repeat(1);
-                const previewRanges: vscode.DecorationOptions[] = [];
     const edit = new vscode.WorkspaceEdit();
-                const document = textEditor.document;
+    // Overwrite the context by replacing instead of inserting
-
+    edit.replace(document.uri, fullRange, previewContent);
-                for (let i = 0; i < totalNewLines; i++) {
+    await vscode.workspace.applyEdit(edit);
-                    const lineContent = i < newLines.length ? newLines[i] : '';
-                    const position = new vscode.Position(startLine + i, 0);
-                    
-                    // Insert a new line
-                    edit.insert(document.uri, position, '\n');
-                    
-                    // Create a range for the newly inserted line
-                    const range = new vscode.Range(position, position.translate(1, 0));
 
+    // Set decorations on the newly inserted lines
+    const previewRanges: vscode.DecorationOptions[] = [];
+    const previewLines = previewContent.split('\n');
+    for (let i = 0; i < previewLines.length; i++) {
+        const range = new vscode.Range(startLine + i, 0, startLine + i + 1, 0);
         previewRanges.push({
             range,
             renderOptions: {
                 after: {
-                                contentText: lineContent,
+                    contentText: previewLines[i],
                 }
             }
         });
     }
-
-                // Apply the edit to insert new lines
-                await vscode.workspace.applyEdit(edit);
-
-                // Set decorations on the newly inserted lines
     textEditor.setDecorations(previewDecorationType, previewRanges);
+    previewInserted = true;
+};
 
-                let previewInserted = true;
 
                 // Handle preview acceptance or dismissal
                 const disposable = vscode.workspace.onDidChangeTextDocument(async (event) => {
-                if (event.document.uri.toString() === document.uri.toString()) {
+                    if (event.document.uri.toString() === document.uri.toString() && previewInserted) {
                         const change = event.contentChanges[0];
 
                         // Dismiss preview with Backspace
                         if (change && change.text === '' && change.rangeLength === 1) {
-                        textEditor.setDecorations(previewDecorationType, []);
+                            await restoreOriginalContent();
-                        disposable.dispose();
-                        previewInserted = false;
                         }
 
                         // Accept preview with Tab key
-                    if (change && change.text === '\t' && previewInserted) {
+                        if (change && change.text === '\t') {
-                        textEditor.setDecorations(previewDecorationType, []);
+                            await acceptPreview();
-                        const edit = new vscode.WorkspaceEdit();
-
-                        // Insert the completion text
-                        const insertText = '\n'.repeat(1) + completionText + '\n'.repeat(1); // Add 2 extra newlines
-
-                        // Replace the entire range from the start of the context to the current position
-                        const replaceRange = new vscode.Range(startLine, 0, position.line, position.character);
-                        edit.replace(document.uri, replaceRange, insertText);
-
-                        await vscode.workspace.applyEdit(edit);
-                        await document.save(); // Optionally save after inserting
-
-                        // Move the cursor to the end of the inserted completion
-                        const newPosition = new vscode.Position(startLine + totalNewLines - 2, 0);
-                        textEditor.selection = new vscode.Selection(newPosition, newPosition);
-
-                        disposable.dispose();
-                        previewInserted = false;
                         }
                     }
                 });
 
+                const restoreOriginalContent = async () => {
+                    const edit = new vscode.WorkspaceEdit();
+                    const endLine = document.lineCount - 1;
+                    const endCharacter = document.lineAt(endLine).text.length;
+                    const fullRange = new vscode.Range(previewStartLine, 0, endLine, endCharacter);
+                    edit.replace(document.uri, fullRange, originalContent);
+                    await vscode.workspace.applyEdit(edit);
+                    textEditor.setDecorations(previewDecorationType, []);
+                    previewInserted = false;
+                    disposable.dispose();
+                };
+
+                const acceptPreview = async () => {
+                    textEditor.setDecorations(previewDecorationType, []);
+                    const edit = new vscode.WorkspaceEdit();
+                    const endLine = document.lineCount - 1;
+                    const endCharacter = document.lineAt(endLine).text.length;
+                    const fullRange = new vscode.Range(previewStartLine, 0, endLine, endCharacter);
+                    edit.replace(document.uri, fullRange, completionText);
+                    await vscode.workspace.applyEdit(edit);
+                    await document.save();
+                    
+                    // Move the cursor to the end of the inserted completion
+                    const newPosition = new vscode.Position(previewStartLine + completionText.split('\n').length - 1, 0);
+                    textEditor.selection = new vscode.Selection(newPosition, newPosition);
+                    
+                    previewInserted = false;
+                    disposable.dispose();
+                };
+
+                // Call this function to initiate the preview
+                await storeAndInsertPreview();
+
             } catch (err: any) {
                 vscode.window.showErrorMessage(
                     "Fabelous Autocoder encountered an error: " + err.message