package.json

@@ -1,6 +1,6 @@
 {
 	"name": "fabelous-autocoder",
-	"version": "0.1.7",
+	"version": "0.2.0",
 	"displayName": "Fabelous Autocoder",
 	"description": "A simple to use Ollama autocompletion Plugin",
 	"icon": "icon.png",
@@ -110,10 +110,9 @@
 		},
 		"commands": [
 			{
-			  "command": "fabelous-autocoder.autocomplete",
+				"command": "fabelous-autocoder.autocomplete",
-			  "title": "Fabelous autocompletion"
+				"title": "Fabelous Autocompletion"
-			}
+			}]
-		  ]
 	},
 	"scripts": {
 		"vscode:prepublish": "npm run compile",

src/extension.ts

@@ -1,253 +1,298 @@
-import * as vscode from "vscode";
+import * as vscode from 'vscode';
-import axios from "axios";
+import axios from 'axios';
 
-let VSConfig: vscode.WorkspaceConfiguration;
+let config: {
-let apiEndpoint: string;
+    apiEndpoint: string;
-let apiAuthentication: string;
+    apiAuthentication: string;
-let apiModel: string;
+    apiModel: string;
-let apiTemperature: number;
+    apiTemperature: number;
-let numPredict: number;
+    numPredict: number;
-let promptWindowSize: number;
+    promptWindowSize: number;
-let completionKeys: string;
+    completionKeys: string[];
-let responsePreview: boolean | undefined;
+    responsePreview: boolean;
-let responsePreviewMaxTokens: number;
+    responsePreviewMaxTokens: number;
-let responsePreviewDelay: number;
+    responsePreviewDelay: number;
-let continueInline: boolean | undefined;
+    continueInline: boolean;
-let keepAlive: number | undefined;
+    keepAlive: number;
-let topP: number | undefined;
+    topP: number;
+};
 
-function updateVSConfig() {
+let previewDecorationType: vscode.TextEditorDecorationType;
-    VSConfig = vscode.workspace.getConfiguration("fabelous-autocoder");
+let activeCompletionManager: CompletionManager | null = null;
-    apiEndpoint = VSConfig.get("endpoint") || "http://localhost:11434/api/generate";
+
-    apiAuthentication = VSConfig.get("authentication") || "";
+function updateConfig() {
-    apiModel = VSConfig.get("model") || "fabelous-coder:latest";
+    const vsConfig = vscode.workspace.getConfiguration('fabelous-autocoder');
-    numPredict = VSConfig.get("max tokens predicted") || 1000;
+    config = {
-    promptWindowSize = VSConfig.get("prompt window size") || 2000;
+        apiEndpoint: vsConfig.get('endpoint') || 'http://localhost:11434/api/generate',
-    completionKeys = VSConfig.get("completion keys") || " ";
+        apiAuthentication: vsConfig.get('authentication') || '',
-    responsePreview = VSConfig.get("response preview");
+        apiModel: vsConfig.get('model') || 'fabelous-coder:latest',
-    responsePreviewMaxTokens = VSConfig.get("preview max tokens") || 50;
+        apiTemperature: vsConfig.get('temperature') || 0.7,
-    responsePreviewDelay = VSConfig.get("preview delay") || 0;
+        numPredict: vsConfig.get('max tokens predicted') || 1000,
-    continueInline = VSConfig.get("continue inline");
+        promptWindowSize: vsConfig.get('prompt window size') || 2000,
-    apiTemperature = VSConfig.get("temperature") || 0.7;
+        completionKeys: (vsConfig.get('completion keys') as string || ' ').split(''),
-    keepAlive = VSConfig.get("keep alive") || 30;
+        responsePreview: vsConfig.get('response preview') || false,
-    topP = VSConfig.get("top p") || 1;
+        responsePreviewMaxTokens: vsConfig.get('preview max tokens') || 50,
+        responsePreviewDelay: vsConfig.get('preview delay') || 0,
+        continueInline: vsConfig.get('continue inline') || false,
+        keepAlive: vsConfig.get('keep alive') || 30,
+        topP: vsConfig.get('top p') || 1,
+    };
 }
 
-updateVSConfig();
+function createPreviewDecorationType() {
-vscode.workspace.onDidChangeConfiguration(updateVSConfig);
+    previewDecorationType = vscode.window.createTextEditorDecorationType({
+        after: {
+            color: '#888888',
+            fontStyle: 'italic',
+        },
+        textDecoration: 'none; display: none;',
+    });
+}
 
 function getContextLines(document: vscode.TextDocument, position: vscode.Position): string {
-    const lines = [];
     const startLine = Math.max(0, position.line - 1);
     const endLine = position.line;
-
+    return document.getText(new vscode.Range(startLine, 0, endLine, position.character));
-    for (let i = startLine; i <= endLine; i++) {
-        lines.push(document.lineAt(i).text);
-    }
-
-    return lines.join("\n");
 }
 
-
 function createFIMPrompt(prefix: string, language: string): string {
     return `<fim_prefix>${prefix}<fim_middle><fim_suffix>${language}\n`;
 }
 
-const previewDecorationType = vscode.window.createTextEditorDecorationType({
+async function generateCompletion(prompt: string, cancellationToken: vscode.CancellationToken): Promise<string> {
-    color: '#888888', // Grayed-out preview text
+    const axiosCancelToken = new axios.CancelToken((c) => {
-    fontStyle: 'italic',
+        cancellationToken.onCancellationRequested(() => c('Request cancelled'));
-    rangeBehavior: vscode.DecorationRangeBehavior.ClosedOpen, // Ensure proper handling of multiline decorations
+    });
-});
 
-async function autocompleteCommand(textEditor: vscode.TextEditor, cancellationToken?: vscode.CancellationToken) {
+    const response = await axios.post(config.apiEndpoint, {
-    const document = textEditor.document;
+        model: config.apiModel,
-    const position = textEditor.selection.active;
+        prompt: prompt,
-
+        stream: false,
-    const context = getContextLines(document, position);
+        raw: true,
-    const fimPrompt = createFIMPrompt(context, document.languageId);
+        options: {
-
+            num_predict: config.numPredict,
-    vscode.window.withProgress(
+            temperature: config.apiTemperature,
-        {
+            stop: ['<fim_suffix>', '```'],
-            location: vscode.ProgressLocation.Notification,
+            keep_alive: config.keepAlive,
-            title: "Fabelous Autocoder",
+            top_p: config.topP,
-            cancellable: true,
-        },
-        async (progress, progressCancellationToken) => {
-            try {
-                progress.report({ message: "Starting model..." });
-
-                let axiosCancelPost: () => void;
-                const axiosCancelToken = new axios.CancelToken((c) => {
-                    axiosCancelPost = () => {
-                        c("Autocompletion request terminated by user cancel");
-                    };
-                    if (cancellationToken) cancellationToken.onCancellationRequested(axiosCancelPost);
-                    progressCancellationToken.onCancellationRequested(axiosCancelPost);
-                    vscode.workspace.onDidCloseTextDocument(axiosCancelPost);
-                });
-
-                const response = await axios.post(apiEndpoint, {
-                    model: apiModel,
-                    prompt: fimPrompt,
-                    stream: false,
-                    raw: true,
-                    options: {
-                        num_predict: numPredict,
-                        temperature: apiTemperature,
-                        stop: ["<fim_suffix>", "```"]
-                    }
-                }, {
-                    cancelToken: axiosCancelToken,
-                    headers: {
-                        'Authorization': apiAuthentication
-                    }
-                });
-
-                progress.report({ message: "Generating..." });
-
-                let completionText = response.data.response;
-                completionText = completionText.replace(/<fim_middle>|<fim_suffix>|<fim_prefix>/g, '').trim();
-
-                // Split the completion text by new lines
-                const lines = completionText.split('\n');
-
-                // Create a decoration for each line of the response
-                const previewRanges = lines.map((line: string, idx: number) => {
-                    const linePos = new vscode.Position(position.line + idx, 0);
-                    const range = new vscode.Range(linePos, linePos); // Set range at the start of each new line
-                    return {
-                        range,
-                        renderOptions: {
-                            before: {
-                                contentText: line,
-                                color: '#888888',
-                                fontStyle: 'italic',
-                            }
-                        }
-                    };
-                });
-
-                // Apply the decorations for multiline preview
-                textEditor.setDecorations(previewDecorationType, previewRanges);
-
-                const disposable = vscode.workspace.onDidChangeTextDocument(async (event) => {
-                    if (event.document.uri.toString() === document.uri.toString()) {
-                        const change = event.contentChanges[0];
-
-                        // Handle Backspace to decline the preview
-                        if (change && change.text === '' && change.rangeLength === 1) {
-                            textEditor.setDecorations(previewDecorationType, []); // Remove preview decorations
-                            disposable.dispose();
-                        }
-
-                        // Handle Ctrl + Enter (or Cmd + Enter on macOS) to accept the preview
-                        const isCtrlOrCmdPressed = event.contentChanges.some(
-                            (change) => {
-                                const isMac = process.platform === 'darwin';
-                                const isCtrlOrCmd = isMac ? change.text.includes('\u0010') : change.text.includes('\n');
-                                return isCtrlOrCmd;
-                            }
-                        );
-
-                        if (isCtrlOrCmdPressed) {
-                            // Remove the preview decoration before applying the final completion
-                            textEditor.setDecorations(previewDecorationType, []);
-
-                            const edit = new vscode.WorkspaceEdit();
-                            const insertPosition = new vscode.Position(position.line, 0);
-
-                            // Insert the completion only once
-                            if (!document.getText().includes(completionText)) {
-                                edit.insert(document.uri, insertPosition, '\n' + completionText);
-                                await vscode.workspace.applyEdit(edit);
-                            }
-
-                            const newPosition = new vscode.Position(position.line + lines.length, lines[lines.length - 1].length);
-                            textEditor.selection = new vscode.Selection(newPosition, newPosition);
-
-                            disposable.dispose(); // Clean up the listener after accepting the completion
-                        }
-                    }
-                });
-
-            } catch (err: any) {
-                vscode.window.showErrorMessage(
-                    "Fabelous Autocoder encountered an error: " + err.message
-                );
-                console.log(err);
-            }
         }
-    );
+    }, {
+        cancelToken: axiosCancelToken,
+        headers: {
+            'Authorization': config.apiAuthentication
+        }
+    });
+
+    return response.data.response.replace(/<fim_middle>|<fim_suffix>|<fim_prefix>/g, '').trim();
 }
 
+class CompletionManager {
+    private textEditor: vscode.TextEditor;
+    private document: vscode.TextDocument;
+    private startPosition: vscode.Position;
+    private completionText: string;
 
+    constructor(textEditor: vscode.TextEditor, startPosition: vscode.Position, completionText: string) {
+        this.textEditor = textEditor;
+        this.document = textEditor.document;
+        this.startPosition = startPosition;
+        this.completionText = completionText;
+    }
+
+    public async showPreview() {
+        this.completionText = '\n' + this.completionText;
+        const completionLines = this.completionText.split('\n');
+
+        const emptyLine = ''; // Empty line for spacing
+        const previewLines = [emptyLine, ...completionLines, emptyLine];
+
+        const previewRanges: vscode.DecorationOptions[] = previewLines.map((line, index) => {
+            const actualLineNumber = this.startPosition.line + index;
+            const totalLines = this.textEditor.document.lineCount;
+            const lineNumber = Math.min(totalLines - 1, actualLineNumber);
+            return {
+                range: new vscode.Range(
+                    new vscode.Position(lineNumber, 0),
+                    new vscode.Position(lineNumber, Number.MAX_VALUE)
+                ),
+                renderOptions: {
+                    after: {
+                        contentText: line.length > 0 ? ` ${line}` : '',
+                    },
+                },
+            };
+        });
+
+        this.textEditor.setDecorations(previewDecorationType, previewRanges);
+    }
+
+    public async acceptCompletion() {
+        const edit = new vscode.WorkspaceEdit();
+        const startLine = Math.max(0, this.startPosition.line - 1);
+        const range = new vscode.Range(
+            new vscode.Position(startLine, 0),
+            this.startPosition.translate(0, Number.MAX_VALUE)
+        );
+        edit.replace(this.document.uri, range, this.completionText);
+        await vscode.workspace.applyEdit(edit);
+        this.clearPreview();
+    }
+    
+    public clearPreview() {
+        this.textEditor.setDecorations(previewDecorationType, []);
+    }
+
+    public declineCompletion() {
+        this.clearPreview();
+    }
+}
+
+async function autocompleteCommand(textEditor: vscode.TextEditor, edit: vscode.TextEditorEdit, ...args: any[]) {
+    const cancellationTokenSource = new vscode.CancellationTokenSource();
+    const cancellationToken = cancellationTokenSource.token;
+
+    try {
+        const document = textEditor.document;
+        const position = textEditor.selection.active;
+        const context = getContextLines(document, position);
+        const fimPrompt = createFIMPrompt(context, document.languageId);
+
+        const completionText = await vscode.window.withProgress({
+            location: vscode.ProgressLocation.Notification,
+            title: 'Fabelous Autocoder',
+            cancellable: true,
+        }, async (progress, progressCancellationToken) => {
+            progress.report({ message: 'Generating...' });
+            return await generateCompletion(fimPrompt, progressCancellationToken);
+        });
+
+        console.log('Completion generated:', completionText);
+
+
+        const completionManager = new CompletionManager(textEditor, position, completionText);
+        await completionManager.showPreview();
+        activeCompletionManager = completionManager;
+
+        let isDisposed = false;
+
+        const dispose = () => {
+            if (!isDisposed) {
+                console.log('Disposing listeners');
+                disposable.dispose();
+                typeDisposable.dispose();
+                activeCompletionManager = null;
+                isDisposed = true;
+            }
+        };
+
+        const disposable = vscode.Disposable.from(
+            vscode.window.onDidChangeTextEditorSelection(async (event) => {
+                if (event.textEditor !== textEditor) return;
+
+                if (event.kind === vscode.TextEditorSelectionChangeKind.Keyboard) {
+                    console.log('Accepting completion');
+                    await completionManager.acceptCompletion();
+                    dispose();
+                }
+            }),
+            vscode.window.onDidChangeActiveTextEditor(() => {
+                console.log('Active editor changed, clearing preview');
+                completionManager.clearPreview();
+                dispose();
+            }),
+            vscode.workspace.onDidChangeTextDocument((event) => {
+                if (event.document === document) {
+                    console.log('Document changed, clearing preview');
+                    completionManager.clearPreview();
+                    dispose();
+                }
+            })
+        );
+
+        const typeDisposable = vscode.commands.registerCommand('type', async (args) => {
+            if (args.text === '\b') { // Backspace key
+                console.log('Declining completion');
+                completionManager.declineCompletion();
+                dispose();
+            }
+        });
+
+    } catch (err: any) {
+        console.error('Error in autocompleteCommand:', err);
+        vscode.window.showErrorMessage(`Fabelous Autocoder encountered an error: ${err.message}`);
+    } finally {
+        cancellationTokenSource.dispose();
+    }
+}
+
+async function acceptCompletion() {
+    if (activeCompletionManager) {
+        await activeCompletionManager.acceptCompletion();
+        const editor = vscode.window.activeTextEditor;
+        if (editor) {
+            const lastLine = editor.document.lineAt(editor.document.lineCount - 1);
+            const newPosition = new vscode.Position(lastLine.lineNumber, lastLine.text.length);
+            editor.selection = new vscode.Selection(newPosition, newPosition);
+        }
+        activeCompletionManager = null;
+    }
+}
+
+async function handleTab() {
+    if (activeCompletionManager) {
+        await acceptCompletion();
+    } else {
+        await vscode.commands.executeCommand('tab');
+    }
+}
 
 async function provideCompletionItems(document: vscode.TextDocument, position: vscode.Position, cancellationToken: vscode.CancellationToken) {
-    const item = new vscode.CompletionItem("Fabelous autocompletion");
+    const item = new vscode.CompletionItem('Fabelous autocompletion');
     item.insertText = new vscode.SnippetString('${1:}');
+    item.documentation = new vscode.MarkdownString('Press `Enter` to get an autocompletion from Fabelous Autocoder');
 
-    if (responsePreview) {
+    if (config.responsePreview) {
-        await new Promise(resolve => setTimeout(resolve, responsePreviewDelay * 1000));
+        await new Promise(resolve => setTimeout(resolve, config.responsePreviewDelay * 1000));
         if (cancellationToken.isCancellationRequested) {
-            return [ item ];
+            return [item];
         }
 
         const context = getContextLines(document, position);
         const fimPrompt = createFIMPrompt(context, document.languageId);
 
         try {
-            const response_preview = await axios.post(apiEndpoint, {
+            const result = await generateCompletion(fimPrompt, cancellationToken);
-                model: apiModel,
+            const preview = (result as any).preview;
-                prompt: fimPrompt,
+            if (preview) {
-                stream: false,
+                item.detail = preview.split('\n')[0];
-                raw: true,
+            }
-                options: {
-                    num_predict: responsePreviewMaxTokens,
-                    temperature: apiTemperature,
-                    stop: ['<fim_suffix>', '\n', '```'],
-                    ...(keepAlive && { keep_alive: keepAlive }),
-                    ...(topP && { top_p: topP }),
-                }
-            }, {
-                cancelToken: new axios.CancelToken((c) => {
-                    cancellationToken.onCancellationRequested(() => c("Autocompletion request terminated by completion cancel"));
-                })
-            });
         } catch (error) {
-            console.error("Error fetching preview:", error);
+            console.error('Error fetching preview:', error);
         }
     }
 
-    item.documentation = new vscode.MarkdownString('Press `Enter` to get an autocompletion from Fabelous Autocoder');
+    if (config.continueInline || !config.responsePreview) {
-    if (continueInline || !responsePreview) {
         item.command = {
             command: 'fabelous-autocoder.autocomplete',
             title: 'Fabelous Autocomplete',
-            arguments: [cancellationToken]
+            arguments: []
         };
     }
+
     return [item];
 }
+export function activate(context: vscode.ExtensionContext) {
+    updateConfig();
+    createPreviewDecorationType();
 
-function activate(context: vscode.ExtensionContext) {
+    context.subscriptions.push(
-    const completionProvider = vscode.languages.registerCompletionItemProvider("*", {
+        vscode.workspace.onDidChangeConfiguration(updateConfig),
-        provideCompletionItems
+        vscode.languages.registerCompletionItemProvider('*', { provideCompletionItems }, ...config.completionKeys),
-    },
+        vscode.commands.registerTextEditorCommand('fabelous-autocoder.autocomplete', autocompleteCommand),
-        ...completionKeys.split("")
+        vscode.commands.registerCommand('fabelous-autocoder.acceptCompletion', acceptCompletion),
+        vscode.commands.registerCommand('fabelous-autocoder.handleTab', handleTab)
     );
-    const externalAutocompleteCommand = vscode.commands.registerTextEditorCommand(
-        "fabelous-autocoder.autocomplete",
-        (textEditor, _, cancellationToken?) => {
-            autocompleteCommand(textEditor, cancellationToken);
-        }
-    );
-    context.subscriptions.push(completionProvider);
-    context.subscriptions.push(externalAutocompleteCommand);
 }
 
-function deactivate() { }
 
-module.exports = {
+export function deactivate() {}
-    activate,
-    deactivate,
-};