[ES|QL] Support replacement range

packages/kbn-esql-validation-autocomplete/src/autocomplete/autocomplete.test.ts

@@ -1358,4 +1358,91 @@ describe('autocomplete', () => {
       ).map((s) => (s.text.toLowerCase().includes('null') ? s : attachTriggerCommand(s)))
     );
   });
+
+  describe('Replacement ranges are attached when needed', () => {
+    testSuggestions('FROM a | WHERE doubleField IS NOT N/', [
+      { text: 'IS NOT NULL', rangeToReplace: { start: 28, end: 35 } },
+      { text: 'IS NULL', rangeToReplace: { start: 35, end: 35 } },
+      '!= $0',
+      '< $0',
+      '<= $0',
+      '== $0',
+      '> $0',
+      '>= $0',
+      'IN $0',
+    ]);
+    testSuggestions('FROM a | WHERE doubleField IS N/', [
+      { text: 'IS NOT NULL', rangeToReplace: { start: 28, end: 31 } },
+      { text: 'IS NULL', rangeToReplace: { start: 28, end: 31 } },
+      { text: '!= $0', rangeToReplace: { start: 31, end: 31 } },
+      '< $0',
+      '<= $0',
+      '== $0',
+      '> $0',
+      '>= $0',
+      'IN $0',
+    ]);
+    testSuggestions('FROM a | EVAL doubleField IS NOT N/', [
+      { text: 'IS NOT NULL', rangeToReplace: { start: 27, end: 34 } },
+      'IS NULL',
+      '% $0',
+      '* $0',
+      '+ $0',
+      '- $0',
+      '/ $0',
+      '!= $0',
+      '< $0',
+      '<= $0',
+      '== $0',
+      '> $0',
+      '>= $0',
+      'IN $0',
+    ]);
+    testSuggestions('FROM a | SORT doubleField IS NOT N/', [
+      { text: 'IS NOT NULL', rangeToReplace: { start: 27, end: 34 } },
+      'IS NULL',
+      '% $0',
+      '* $0',
+      '+ $0',
+      '- $0',
+      '/ $0',
+      '!= $0',
+      '< $0',
+      '<= $0',
+      '== $0',
+      '> $0',
+      '>= $0',
+      'IN $0',
+    ]);
+    describe('dot-separated field names', () => {
+      testSuggestions(
+        'FROM a | KEEP field.nam/',
+        [{ text: 'field.name', rangeToReplace: { start: 15, end: 23 } }],
+        undefined,
+        [[{ name: 'field.name', type: 'double' }]]
+      );
+      // multi-line
+      testSuggestions(
+        'FROM a\n| KEEP field.nam/',
+        [{ text: 'field.name', rangeToReplace: { start: 15, end: 23 } }],
+        undefined,
+        [[{ name: 'field.name', type: 'double' }]]
+      );
+      // triple separator
+      testSuggestions(
+        'FROM a\n| KEEP field.name.f/',
+        [{ text: 'field.name.foo', rangeToReplace: { start: 15, end: 26 } }],
+        undefined,
+        [[{ name: 'field.name.foo', type: 'double' }]]
+      );
+      // whitespace — we can't support this case yet because
+      // we are relying on string checking instead of the AST :(
+      testSuggestions.skip(
+        'FROM a | KEEP field . n/',
+        [{ text: 'field . name', rangeToReplace: { start: 15, end: 23 } }],
+        undefined,
+        [[{ name: 'field.name', type: 'double' }]]
+      );
+    });
+  });
 });

packages/kbn-esql-validation-autocomplete/src/autocomplete/autocomplete.ts

@@ -88,6 +88,7 @@ import {
 import { ESQLCallbacks } from '../shared/types';
 import {
   getFunctionsToIgnoreForStats,
+  getOverlapRange,
   getParamAtPosition,
   getQueryForFields,
   getSourcesFromCommands,
@@ -656,31 +657,52 @@ async function getExpressionSuggestionsByType(
     }
     // Suggest fields or variables
     if (argDef.type === 'column' || argDef.type === 'any') {
-      // ... | <COMMAND> <suggest>
       if ((!nodeArg || isNewExpression) && !endsWithNot) {
-        suggestions.push(
-          ...(await getFieldsOrFunctionsSuggestions(
-            argDef.innerTypes ?? ['any'],
-            command.name,
-            option?.name,
-            getFieldsByType,
-            {
-              // TODO instead of relying on canHaveAssignments and other command name checks
-              // we should have a more generic way to determine if a command can have functions.
-              // I think it comes down to the definition of 'column' since 'any' should always
-              // include functions.
-              functions: canHaveAssignments || command.name === 'sort',
-              fields: !argDef.constantOnly,
-              variables: anyVariables,
-              literals: argDef.constantOnly,
-            },
-            {
-              ignoreFields: isNewExpression
-                ? command.args.filter(isColumnItem).map(({ name }) => name)
-                : [],
-            }
-          ))
+        const fieldSuggestions = await getFieldsOrFunctionsSuggestions(
+          argDef.innerTypes || ['any'],
+          command.name,
+          option?.name,
+          getFieldsByType,
+          {
+            // TODO instead of relying on canHaveAssignments and other command name checks
+            // we should have a more generic way to determine if a command can have functions.
+            // I think it comes down to the definition of 'column' since 'any' should always
+            // include functions.
+            functions: canHaveAssignments || command.name === 'sort',
+            fields: !argDef.constantOnly,
+            variables: anyVariables,
+            literals: argDef.constantOnly,
+          },
+          {
+            ignoreFields: isNewExpression
+              ? command.args.filter(isColumnItem).map(({ name }) => name)
+              : [],
+          }
         );
+
+        /**
+         * @TODO — this string manipulation is crude and can't support all cases
+         * Checking for a partial word and computing the replacement range should
+         * really be done using the AST node, but we'll have to refactor further upstream
+         * to make that available. This is a quick fix to support the most common case.
+         */
+        const words = innerText.split(/\s+/);
+        const lastWord = words[words.length - 1];
+        if (lastWord !== '') {
+          // ... | <COMMAND> <word><suggest>
+          suggestions.push(
+            ...fieldSuggestions.map((suggestion) => ({
+              ...suggestion,
+              rangeToReplace: {
+                start: innerText.length - lastWord.length + 1,
+                end: innerText.length,
+              },
+            }))
+          );
+        } else {
+          // ... | <COMMAND> <suggest>
+          suggestions.push(...fieldSuggestions);
+        }
       }
     }
     if (argDef.type === 'function' || argDef.type === 'any') {
@@ -784,6 +806,7 @@ async function getExpressionSuggestionsByType(
               const nodeArgType = extractFinalTypeFromArg(nodeArg, references);
               suggestions.push(
                 ...(await getBuiltinFunctionNextArgument(
+                  innerText,
                   command,
                   option,
                   argDef,
@@ -859,7 +882,7 @@ async function getExpressionSuggestionsByType(
           // i.e. ... | <COMMAND> field >= <suggest>
           // i.e. ... | <COMMAND> field > 0 <suggest>
           // i.e. ... | <COMMAND> field + otherN <suggest>
-
+          // "FROM a | WHERE doubleField IS NOT N"
           if (nodeArgType) {
             if (isFunctionItem(nodeArg)) {
               if (nodeArg.name === 'not') {
@@ -879,6 +902,7 @@ async function getExpressionSuggestionsByType(
               } else {
                 suggestions.push(
                   ...(await getBuiltinFunctionNextArgument(
+                    innerText,
                     command,
                     option,
                     argDef,
@@ -988,6 +1012,7 @@ async function getExpressionSuggestionsByType(
 }
 
 async function getBuiltinFunctionNextArgument(
+  queryText: string,
   command: ESQLCommand,
   option: ESQLCommandOption | undefined,
   argDef: { type: string },
@@ -1072,7 +1097,16 @@ async function getBuiltinFunctionNextArgument(
       }
     }
   }
-  return suggestions;
+  return suggestions.map<SuggestionRawDefinition>((s) => {
+    const overlap = getOverlapRange(queryText, s.text);
+    return {
+      ...s,
+      rangeToReplace: {
+        start: overlap.start,
+        end: overlap.end,
+      },
+    };
+  });
 }
 
 function pushItUpInTheList(suggestions: SuggestionRawDefinition[], shouldPromote: boolean) {
@@ -1636,6 +1670,7 @@ async function getOptionArgsSuggestions(
       if (isFunctionItem(nodeArg) && !isFunctionArgComplete(nodeArg, references).complete) {
         suggestions.push(
           ...(await getBuiltinFunctionNextArgument(
+            innerText,
             command,
             option,
             { type: argDef?.type || 'any' },

packages/kbn-esql-validation-autocomplete/src/autocomplete/helper.ts

@@ -92,3 +92,42 @@ export function getSupportedTypesForBinaryOperators(
         .map(({ params }) => params[1].type)
     : [previousType];
 }
+
+/**
+ * Checks the suggestion text for overlap with the current query.
+ *
+ * This is useful to determine the range of the existing query that should be
+ * replaced if the suggestion is accepted.
+ *
+ * For example
+ * QUERY: FROM source | WHERE field IS NO
+ * SUGGESTION: IS NOT NULL
+ *
+ * The overlap is "IS NO" and the range to replace is "IS NO" in the query.
+ *
+ * @param query
+ * @param suggestionText
+ * @returns
+ */
+export function getOverlapRange(
+  query: string,
+  suggestionText: string
+): { start: number; end: number } {
+  let overlapLength = 0;
+
+  // Convert both strings to lowercase for case-insensitive comparison
+  const lowerQuery = query.toLowerCase();
+  const lowerSuggestionText = suggestionText.toLowerCase();
+
+  for (let i = 0; i <= lowerSuggestionText.length; i++) {
+    const substr = lowerSuggestionText.substring(0, i);
+    if (lowerQuery.endsWith(substr)) {
+      overlapLength = i;
+    }
+  }
+
+  return {
+    start: Math.min(query.length - overlapLength + 1, query.length),
+    end: query.length,
+  };
+}

packages/kbn-esql-validation-autocomplete/src/autocomplete/types.ts

@@ -57,6 +57,13 @@ export interface SuggestionRawDefinition {
     title: string;
     id: string;
   };
+  /**
+   * The range that should be replaced when the suggestion is applied
+   */
+  rangeToReplace?: {
+    start: number;
+    end: number;
+  };
 }
 
 export interface EditorContext {

packages/kbn-monaco/src/esql/language.ts

@@ -102,7 +102,8 @@ export const ESQLLang: CustomLangModuleType<ESQLCallbacks> = {
         );
         const suggestionEntries = await astAdapter.autocomplete(model, position, context);
         return {
-          suggestions: wrapAsMonacoSuggestions(suggestionEntries.suggestions),
+          // @ts-expect-error because of range typing: https://github.com/microsoft/monaco-editor/issues/4638
+          suggestions: wrapAsMonacoSuggestions(suggestionEntries),
         };
       },
     };

packages/kbn-monaco/src/esql/lib/converters/suggestions.ts

@@ -6,29 +6,34 @@
  * Side Public License, v 1.
  */
 
-import type { SuggestionRawDefinition } from '@kbn/esql-validation-autocomplete';
 import { monaco } from '../../../monaco_imports';
-import { MonacoAutocompleteCommandDefinition } from '../types';
+import {
+  MonacoAutocompleteCommandDefinition,
+  SuggestionRawDefinitionWithMonacoRange,
+} from '../types';
 
 export function wrapAsMonacoSuggestions(
-  suggestions: SuggestionRawDefinition[]
+  suggestions: SuggestionRawDefinitionWithMonacoRange[]
 ): MonacoAutocompleteCommandDefinition[] {
-  return suggestions.map(
-    ({ label, text, asSnippet, kind, detail, documentation, sortText, command }) => ({
-      label,
-      insertText: text,
-      kind:
-        kind in monaco.languages.CompletionItemKind
-          ? monaco.languages.CompletionItemKind[kind]
-          : monaco.languages.CompletionItemKind.Method, // fallback to Method
-      detail,
-      documentation,
-      sortText,
-      command,
-      insertTextRules: asSnippet
-        ? monaco.languages.CompletionItemInsertTextRule.InsertAsSnippet
-        : undefined,
-      range: undefined as unknown as monaco.IRange,
-    })
+  return suggestions.map<MonacoAutocompleteCommandDefinition>(
+    ({ label, text, asSnippet, kind, detail, documentation, sortText, command, range }) => {
+      const monacoSuggestion: MonacoAutocompleteCommandDefinition = {
+        label,
+        insertText: text,
+        kind:
+          kind in monaco.languages.CompletionItemKind
+            ? monaco.languages.CompletionItemKind[kind]
+            : monaco.languages.CompletionItemKind.Method, // fallback to Method
+        detail,
+        documentation,
+        sortText,
+        command,
+        insertTextRules: asSnippet
+          ? monaco.languages.CompletionItemInsertTextRule.InsertAsSnippet
+          : undefined,
+        range,
+      };
+      return monacoSuggestion;
+    }
   );
 }

packages/kbn-monaco/src/esql/lib/esql_ast_provider.ts

@@ -16,8 +16,9 @@ import { monaco } from '../../monaco_imports';
 import type { ESQLWorker } from '../worker/esql_worker';
 import { wrapAsMonacoMessages } from './converters/positions';
 import { getHoverItem } from './hover/hover';
-import { monacoPositionToOffset } from './shared/utils';
+import { monacoPositionToOffset, offsetRangeToMonacoRange } from './shared/utils';
 import { getSignatureHelp } from './signature';
+import { SuggestionRawDefinitionWithMonacoRange } from './types';
 
 export class ESQLAstAdapter {
   constructor(
@@ -66,17 +67,17 @@ export class ESQLAstAdapter {
     model: monaco.editor.ITextModel,
     position: monaco.Position,
     context: monaco.languages.CompletionContext
-  ) {
+  ): Promise<SuggestionRawDefinitionWithMonacoRange[]> {
     const getAstFn = await this.getAstWorker(model);
     const fullText = model.getValue();
     const offset = monacoPositionToOffset(fullText, position);
-    const suggestionEntries = await suggest(fullText, offset, context, getAstFn, this.callbacks);
-    return {
-      suggestions: suggestionEntries.map((suggestion) => ({
-        ...suggestion,
-        range: undefined as unknown as monaco.IRange,
-      })),
-    };
+    const suggestions = await suggest(fullText, offset, context, getAstFn, this.callbacks);
+    for (const s of suggestions) {
+      (s as SuggestionRawDefinitionWithMonacoRange).range = s.rangeToReplace
+        ? offsetRangeToMonacoRange(fullText, s.rangeToReplace)
+        : undefined;
+    }
+    return suggestions;
   }
 
   async codeAction(