feat: intellisense > prioritize column items over others (#465)

This is a small workaround to prevent the Monaco editor from
re-arranging the IntelliSense completion items based on the user's
search term.

Firstly, I trigger provideCompletionItems on every key press by setting
incomplete=true so I can change the filterText dynamically. To avoid
calling the language server on each key press, I implemented a caching
mechanism that prevents unnecessary calls.

Additionally, I prepend the user search term to column type items only
when the search term is present in the filter text.

Since the Monaco editor re-arranges items where the search term appears
first in the filterText, this ensures that the column items are brought
up in the IntelliSense order.

Plus, added sorting logic for intellisense items.

.github/workflows/pr.yml

@@ -34,6 +34,10 @@ jobs:
       - name: test samples
         run: yarn test:samples
 
+      - name: unit tests
+        working-directory: package
+        run: yarn test
+
       - name: integration tests
         working-directory: package
         run: yarn test:integration

package/jest.config.ts

@@ -3,4 +3,5 @@ module.exports = {
     transform: {
         '^.+\\.ts$': 'ts-jest',
     },
+    setupFilesAfterEnv: ['./jest.setup.ts'],
 };

package/jest.setup.ts

@@ -0,0 +1,4 @@
+require('@kusto/language-service/bridge.min');
+require('@kusto/language-service/Kusto.JavaScript.Client.min');
+require('@kusto/language-service/newtonsoft.json.min');
+require('@kusto/language-service-next/Kusto.Language.Bridge.min');

package/package.json

@@ -42,6 +42,7 @@
         "@babel/core": "^7.22.20",
         "@babel/preset-env": "^7.22.20",
         "@babel/preset-typescript": "^7.22.15",
+        "@faker-js/faker": "^8.4.1",
         "@playwright/test": "^1.44.0",
         "@rollup/plugin-alias": "^5.0.0",
         "@rollup/plugin-babel": "^6.0.3",
@@ -51,6 +52,7 @@
         "@rollup/plugin-terser": "^0.4.3",
         "@rollup/plugin-virtual": "^3.0.1",
         "@tsconfig/node20": "^20.1.2",
+        "@types/jest": "^29.5.12",
         "@types/lodash-es": "^4.17.9",
         "@types/node": "^20.6.3",
         "@types/xregexp": "^4.4.0",

package/src/completionCacheManager/completionCacheManager.spec.ts

@@ -0,0 +1,48 @@
+import { test, describe, expect, beforeEach } from '@jest/globals';
+import { CompletionCacheManager, GetFromLanguageService, createCompletionCacheManager } from './completionCacheManager';
+import * as monaco from 'monaco-editor';
+import * as ls from 'vscode-languageserver-types';
+
+describe('completionCacheManager', () => {
+    let completionCacheManager: CompletionCacheManager;
+    const mockGetFromLanguageService: jest.MockedFunction<GetFromLanguageService> = jest.fn();
+    const resource = {} as monaco.Uri;
+    const position = {} as ls.Position;
+
+    beforeEach(() => {
+        mockGetFromLanguageService.mockClear();
+        completionCacheManager = createCompletionCacheManager(mockGetFromLanguageService);
+    });
+
+    test('calls language service on single character', async () => {
+        await completionCacheManager.getCompletionItems('a', resource, position);
+        expect(mockGetFromLanguageService).toHaveBeenCalledTimes(1);
+    });
+
+    describe('on subsequent calls', () => {
+        beforeEach(async () => {
+            await completionCacheManager.getCompletionItems('ab', resource, position);
+            mockGetFromLanguageService.mockClear();
+        });
+
+        test('does not call language service again when the previous word is contained in the current word', async () => {
+            await completionCacheManager.getCompletionItems('abc', resource, position);
+            expect(mockGetFromLanguageService).toHaveBeenCalledTimes(0);
+        });
+
+        test('calls language service for different non-contained words', async () => {
+            await completionCacheManager.getCompletionItems('ba', resource, position);
+            expect(mockGetFromLanguageService).toHaveBeenCalledTimes(1);
+        });
+
+        test('calls language service when current word is undefined', async () => {
+            await completionCacheManager.getCompletionItems(undefined, resource, position);
+            expect(mockGetFromLanguageService).toHaveBeenCalledTimes(1);
+        });
+
+        test('calls language service when the current word is contained in the previous word', async () => {
+            await completionCacheManager.getCompletionItems('a', resource, position);
+            expect(mockGetFromLanguageService).toHaveBeenCalledTimes(1);
+        });
+    });
+});

package/src/completionCacheManager/completionCacheManager.ts

@@ -0,0 +1,22 @@
+import * as monaco from 'monaco-editor/esm/vs/editor/editor.api';
+import * as ls from 'vscode-languageserver-types';
+
+export const createCompletionCacheManager = (getFromLanguageService: GetFromLanguageService) => {
+    let completionList: Promise<ls.CompletionList> | undefined;
+    let lastWord: string | undefined;
+
+    return {
+        getCompletionItems: (word: string | undefined, resource: monaco.Uri, position: ls.Position) => {
+            const shouldGetItems = !lastWord || !word || !word?.includes(lastWord);
+            if (shouldGetItems) {
+                completionList = getFromLanguageService(resource, position);
+            }
+
+            lastWord = word;
+            return completionList;
+        },
+    };
+};
+
+export type CompletionCacheManager = ReturnType<typeof createCompletionCacheManager>;
+export type GetFromLanguageService = (resource: monaco.Uri, position: ls.Position) => Promise<ls.CompletionList>;

package/src/languageFeatures.ts

@@ -7,6 +7,8 @@ import type { LanguageServiceDefaults, LanguageSettings, OnDidProvideCompletionI
 import type { Schema } from './languageServiceManager/schema';
 import type { ClassifiedRange } from './languageServiceManager/kustoLanguageService';
 import { AugmentedWorkerAccessor } from './kustoMode';
+import { CompletionCacheManager, createCompletionCacheManager } from './completionCacheManager/completionCacheManager';
+import { createCompletionFilteredText } from './languageFeatures.utils';
 
 // --- diagnostics ---
 
@@ -782,7 +784,17 @@ function toTextEdit(textEdit: ls.TextEdit): monaco.editor.ISingleEditOperation {
 const DEFAULT_DOCS_BASE_URL = 'https://learn.microsoft.com/azure/data-explorer/kusto/query';
 
 export class CompletionAdapter implements monaco.languages.CompletionItemProvider {
-    constructor(private _worker: AugmentedWorkerAccessor, private languageSettings: LanguageSettings) {}
+    private readonly languageSettings: LanguageSettings;
+    private completionCacheManager: CompletionCacheManager;
+
+    constructor(workerAccessor: AugmentedWorkerAccessor, languageSettings: LanguageSettings) {
+        this.languageSettings = languageSettings;
+        const getFromLanguageService = async (resource: monaco.Uri, position: ls.Position) => {
+            const worker = await workerAccessor(resource);
+            return worker.doComplete(resource.toString(), position);
+        };
+        this.completionCacheManager = createCompletionCacheManager(getFromLanguageService);
+    }
 
     public get triggerCharacters(): string[] {
         return [' ', '.', '('];
@@ -802,24 +814,22 @@ export class CompletionAdapter implements monaco.languages.CompletionItemProvide
             wordInfo.endColumn
         );
         const resource = model.uri;
+        const word = model?.getWordAtPosition(position)?.word;
         const onDidProvideCompletionItems: OnDidProvideCompletionItems =
             this.languageSettings.onDidProvideCompletionItems;
 
-        return this._worker(resource)
-            .then((worker) => {
-                return worker.doComplete(resource.toString(), fromPosition(position));
-            })
+        return this.completionCacheManager
+            .getCompletionItems(word, resource, fromPosition(position))
             .then((info) => (onDidProvideCompletionItems ? onDidProvideCompletionItems(info) : info))
             .then((info) => {
-                if (!info) {
-                    return;
-                }
+                if (!info) return;
+
                 let items: monaco.languages.CompletionItem[] = info.items.map((entry) => {
                     let item: monaco.languages.CompletionItem = {
                         label: entry.label,
                         insertText: entry.insertText,
                         sortText: entry.sortText,
-                        filterText: entry.filterText,
+                        filterText: createCompletionFilteredText(word, entry),
                         // TODO: Is this cast safe?
                         documentation: this.formatDocLink((entry.documentation as undefined | ls.MarkupContent)?.value),
                         detail: entry.detail,
@@ -838,7 +848,7 @@ export class CompletionAdapter implements monaco.languages.CompletionItemProvide
                 });
 
                 return {
-                    isIncomplete: info.isIncomplete,
+                    incomplete: true,
                     suggestions: items,
                 };
             });

package/src/languageFeatures.utils.spec.ts

@@ -0,0 +1,62 @@
+import { test, describe, expect } from '@jest/globals';
+import { createCompletionFilteredText } from './languageFeatures.utils';
+import { completionItemBuilder } from '../tests/unit/builders/CompletionListBuilder';
+
+describe('createCompletionFilteredText', () => {
+    test('does not prepend prefix to result when search word is undefined', () => {
+        const completionItem = completionItemBuilder().withForcePrecedence(true).build();
+
+        const filterText = createCompletionFilteredText(undefined, completionItem);
+        expect(filterText).toBe(completionItem.filterText);
+    });
+
+    describe('when the search word is included in the original filter text', () => {
+        const originalFilterText = 'StartTime';
+        const searchWord = 'time';
+
+        test('prepends prefix to result when forcePrecedence=true', () => {
+            const completionItem = completionItemBuilder()
+                .withFilterText(originalFilterText)
+                .withForcePrecedence(true)
+                .build();
+
+            const filterText = createCompletionFilteredText(searchWord, completionItem);
+            expect(filterText).toBe(`${searchWord}${completionItem.filterText}`);
+        });
+
+        test('does not prepend prefix to result when forcePrecedence=false', () => {
+            const completionItem = completionItemBuilder()
+                .withFilterText(originalFilterText)
+                .withForcePrecedence(false)
+                .build();
+
+            const filterText = createCompletionFilteredText(searchWord, completionItem);
+            expect(filterText).toBe(completionItem.filterText);
+        });
+    });
+
+    describe('when the search word is NOT included in the original filter text', () => {
+        const originalFilterText = 'StartTime';
+        const searchWord = 'end';
+
+        test('does not prepend prefix to result when forcePrecedence=true', () => {
+            const completionItem = completionItemBuilder()
+                .withFilterText(originalFilterText)
+                .withForcePrecedence(true)
+                .build();
+
+            const filterText = createCompletionFilteredText(searchWord, completionItem);
+            expect(filterText).toBe(completionItem.filterText);
+        });
+
+        test('does not prepend prefix to result when forcePrecedence=false', () => {
+            const completionItem = completionItemBuilder()
+                .withFilterText(originalFilterText)
+                .withForcePrecedence(false)
+                .build();
+
+            const filterText = createCompletionFilteredText(searchWord, completionItem);
+            expect(filterText).toBe(completionItem.filterText);
+        });
+    });
+});

package/src/languageFeatures.utils.ts

@@ -0,0 +1,13 @@
+import * as ls from 'vscode-languageserver-types';
+
+export function createCompletionFilteredText(
+    searchWord: string | undefined,
+    completionItem: ls.CompletionItem
+): string {
+    if (!searchWord) return completionItem.filterText;
+
+    const containedInFilterText = completionItem.filterText.toLowerCase().includes(searchWord.toLowerCase());
+    const shouldPrependSearchWord = completionItem.data.forcePrecedence && containedInFilterText;
+
+    return shouldPrependSearchWord ? `${searchWord}${completionItem.filterText}` : completionItem.filterText;
+}

package/src/languageServiceManager/competionItemSort.ts

@@ -1,3 +1,31 @@
+import k2 = Kusto.Language.Editor;
+
 export const createSortingText = (priority: number) => {
     return priority.toString().padStart(10, '0');
 };
+
+export function sortByMatchTextKeepingKindOrder(items: k2.CompletionItem[]) {
+    const groupedByKind = groupContiguousByKind(items);
+    const sortedGroupedItems = sortGroupedItems(groupedByKind);
+    return sortedGroupedItems.flat();
+}
+
+function sortGroupedItems(groupedItems: k2.CompletionItem[][]) {
+    return groupedItems.map((group) => group.sort((i1, i2) => i1.MatchText.localeCompare(i2.MatchText)));
+}
+
+function groupContiguousByKind(items: k2.CompletionItem[]) {
+    return items.reduce((result: k2.CompletionItem[][], item: k2.CompletionItem) => {
+        const lastGroup = last(result);
+
+        const shouldCreateNewGroup = !lastGroup || last(lastGroup)?.Kind !== item.Kind;
+        if (shouldCreateNewGroup) result.push([]);
+        last(result).push(item);
+
+        return result;
+    }, []);
+}
+
+function last<T>(array: T[]): T | undefined {
+    return array.length > 0 ? array[array.length - 1] : undefined;
+}

package/src/languageServiceManager/completionItemSort.spec.ts

@@ -1,5 +1,7 @@
+import k2 = Kusto.Language.Editor;
 import { test, describe, expect } from '@jest/globals';
-import { createSortingText } from './competionItemSort';
+import { createSortingText, sortByMatchTextKeepingKindOrder } from './competionItemSort';
+import { kustoCompletionItemBuilder } from '../../tests/unit/builders/KustoCompletionItem';
 
 describe('createSortingText', () => {
     test('returns the correct sort text according to the given order', () => {
@@ -24,3 +26,37 @@ describe('createSortingText', () => {
         expect(sortedItems[5].label).toBe('sixth');
     });
 });
+
+describe('sortByMatchTextKeepingKindOrder', () => {
+    test('should sort by match text grouped by kind', () => {
+        const item1 = kustoCompletionItemBuilder()
+            .withMatchText('count_distinct')
+            .withKind(k2.CompletionKind.AggregateFunction)
+            .build();
+        const item2 = kustoCompletionItemBuilder()
+            .withMatchText('count')
+            .withKind(k2.CompletionKind.AggregateFunction)
+            .build();
+        const item3 = kustoCompletionItemBuilder()
+            .withMatchText('bin_at')
+            .withKind(k2.CompletionKind.BuiltInFunction)
+            .build();
+        const item4 = kustoCompletionItemBuilder()
+            .withMatchText('bin')
+            .withKind(k2.CompletionKind.BuiltInFunction)
+            .build();
+        const item5 = kustoCompletionItemBuilder()
+            .withMatchText('avg')
+            .withKind(k2.CompletionKind.AggregateFunction)
+            .build();
+        const items = [item1, item2, item3, item4, item5];
+
+        const sortedItems = sortByMatchTextKeepingKindOrder(items);
+
+        expect(sortedItems[0].MatchText).toBe('count');
+        expect(sortedItems[1].MatchText).toBe('count_distinct');
+        expect(sortedItems[2].MatchText).toBe('bin');
+        expect(sortedItems[3].MatchText).toBe('bin_at');
+        expect(sortedItems[4].MatchText).toBe('avg');
+    });
+});

package/src/languageServiceManager/kustoLanguageService.ts

@@ -20,7 +20,7 @@ import { Database, EntityGroup, getCslTypeNameFromClrType, getEntityDataTypeFrom
 import type { RenderOptions, VisualizationType, RenderOptionKeys, RenderInfo } from './renderInfo';
 import type { ClusterReference, DatabaseReference } from '../types';
 import { Mutable } from '../util';
-import { createSortingText } from './competionItemSort';
+import { createSortingText, sortByMatchTextKeepingKindOrder } from './competionItemSort';
 
 let List = System.Collections.Generic.List$1;
 
@@ -419,7 +419,8 @@ class KustoLanguageService implements LanguageService {
             });
         }
         const itemsAsArray = this.toArray<k2.CompletionItem>(completionItems.Items);
-        let items: ls.CompletionItem[] = itemsAsArray
+        const sortedArray = sortByMatchTextKeepingKindOrder(itemsAsArray);
+        let items: ls.CompletionItem[] = sortedArray
             .filter(
                 (item) =>
                     !(
@@ -465,6 +466,7 @@ class KustoLanguageService implements LanguageService {
                 lsItem.documentation = helpTopic
                     ? { value: this.formatHelpTopic(helpTopic), kind: ls.MarkupKind.Markdown }
                     : undefined;
+                lsItem.data = { forcePrecedence: kItem.Kind === k2.CompletionKind.Column };
                 return lsItem;
             });