Fix editor marker cleanup for autocomplete

Author: bartoval

src/platform/packages/shared/kbn-esql-language/src/__tests__/commands/autocomplete.ts

@@ -15,7 +15,6 @@
 import { uniq } from 'lodash';
 import type { LicenseType } from '@kbn/licensing-types';
 import type { EsqlFieldType } from '@kbn/esql-types';
-import { Parser } from '@elastic/esql';
 import type { ESQLAstAllCommands } from '@elastic/esql/types';
 import type {
   ICommandCallbacks,
@@ -45,7 +44,7 @@ import { FunctionDefinitionTypes } from '../../commands/definitions/types';
 import { mockContext, getMockCallbacks } from './context_fixtures';
 import { getSafeInsertText } from '../../commands/definitions/utils';
 import { timeUnitsToSuggest } from '../../commands/definitions/constants';
-import { correctQuerySyntax, findAstPosition } from '../../commands/definitions/utils/ast';
+import { findAutocompleteAstPosition } from '../../language/shared/parse_for_autocomplete_query';
 import { FUNCTIONS_TO_IGNORE } from '../../commands/registry/eval/autocomplete';
 import { attachReplacementRanges } from '../../language/autocomplete/utils/prefix_range';
 
@@ -87,24 +86,20 @@ export const suggest = async (
   ) => Promise<ISuggestionItem[]>,
   offset?: number
 ): Promise<ISuggestionItem[]> => {
-  const innerText = query.substring(0, offset ?? query.length);
-  const correctedQuery = correctQuerySyntax(innerText);
-  const { root } = Parser.parse(correctedQuery, { withFormatting: true });
-  const headerConstruction = root?.header?.find((cmd) => cmd.name === commandName);
-
   const cursorPosition = offset ?? query.length;
+  const { innerText, root, command } = findAutocompleteAstPosition(query, cursorPosition);
+  const headerConstruction = root?.header?.find((cmd) => cmd.name === commandName);
+  const targetCommand = headerConstruction ?? command;
 
-  const command = headerConstruction ?? findAstPosition(root, cursorPosition).command;
-
-  if (!command) {
+  if (!targetCommand) {
     throw new Error(`${commandName.toUpperCase()} command not found in the parsed query`);
   }
 
   const contextWithRoot = { ...context, rootAst: root };
 
   const suggestions = await autocomplete(
     query,
-    command,
+    targetCommand,
     mockCallbacks,
     contextWithRoot,
     cursorPosition

src/platform/packages/shared/kbn-esql-language/src/commands/definitions/utils/ast.ts

@@ -96,7 +96,27 @@ export function removeMarkerArgFromArgsList<T extends ESQLSingleAstItem | ESQLAs
 }
 
 export function mapToNonMarkerNode(arg: ESQLAstItem): ESQLAstItem {
-  return Array.isArray(arg) ? arg.filter(isNotMarkerNodeOrArray).map(mapToNonMarkerNode) : arg;
+  if (Array.isArray(arg)) {
+    return arg.filter(isNotMarkerNodeOrArray).map(mapToNonMarkerNode);
+  }
+
+  if ('args' in arg) {
+    return {
+      ...arg,
+      args: arg.args.filter(isNotMarkerNodeOrArray).map(mapToNonMarkerNode) as typeof arg.args,
+    } as typeof arg;
+  }
+
+  if ('values' in arg) {
+    return {
+      ...arg,
+      values: arg.values
+        .filter((value) => !isMarkerNode(value))
+        .map((value) => mapToNonMarkerNode(value) as ESQLAstExpression) as typeof arg.values,
+    } as typeof arg;
+  }
+
+  return arg;
 }
 
 function cleanMarkerNode(node: ESQLSingleAstItem | undefined): ESQLSingleAstItem | undefined {

src/platform/packages/shared/kbn-esql-language/src/commands/definitions/utils/autocomplete/expressions/operators/utils.ts

@@ -9,7 +9,6 @@
 
 import { isList } from '@elastic/esql';
 import type { ESQLSingleAstItem } from '@elastic/esql/types';
-import { isMarkerNode } from '../../../ast';
 import { getOperatorSuggestion } from '../../../operators';
 import type { ISuggestionItem } from '../../../../../registry/types';
 import { logicalOperators } from '../../../../all_operators';
@@ -34,11 +33,7 @@ export function endsWithIsOrIsNotToken(innerText: string): boolean {
 }
 
 export function isOperandMissing(operand: ESQLSingleAstItem | undefined): boolean {
-  return (
-    !operand ||
-    isMarkerNode(operand) ||
-    (operand?.type === 'unknown' && operand?.incomplete === true)
-  );
+  return !operand || (operand?.type === 'unknown' && operand?.incomplete === true);
 }
 
 /** Returns true if we should suggest opening a list for the right operand */

src/platform/packages/shared/kbn-esql-language/src/commands/definitions/utils/expressions.test.ts

@@ -6,14 +6,15 @@
  * your election, the "Elastic License 2.0", the "GNU Affero General Public
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
-import { Parser } from '@elastic/esql';
+import { isAssignment, Parser } from '@elastic/esql';
 import type { SupportedDataType } from '../types';
 import { FunctionDefinitionTypes } from '../types';
 import type { ESQLColumnData } from '../../registry/types';
 import { Location } from '../../registry/types';
-import { buildPartialMatcher, getExpressionType } from './expressions';
+import { buildPartialMatcher, getAssignmentExpressionRoot, getExpressionType } from './expressions';
 import { setTestFunctions } from './test_functions';
 import { TIME_SYSTEM_PARAMS } from './literals';
+import { getAutocompleteCursorContext } from '../../../language/shared/parse_for_autocomplete_query';
 
 describe('buildPartialMatcher', () => {
   it('should build a partial matcher', () => {
@@ -458,3 +459,41 @@ describe('getExpressionType', () => {
     );
   });
 });
+
+describe('getAssignmentExpressionRoot', () => {
+  it('returns undefined for an incomplete assignment after autocomplete parsing cleanup', () => {
+    const query = 'FROM employees | EVAL total = ';
+    const { astContext } = getAutocompleteCursorContext(query, query.length);
+
+    if (astContext.type !== 'expression') {
+      throw new Error(`Expected expression context for query: ${query}`);
+    }
+
+    const assignment = astContext.command.args[astContext.command.args.length - 1];
+
+    if (!assignment || !isAssignment(assignment)) {
+      throw new Error(`Expected assignment expression for query: ${query}`);
+    }
+
+    expect(getAssignmentExpressionRoot(assignment)).toBeUndefined();
+  });
+
+  it('returns the RHS root for a complete assignment after autocomplete parsing cleanup', () => {
+    const query = 'FROM employees | EVAL total = salary';
+    const { astContext } = getAutocompleteCursorContext(query, query.length);
+
+    if (astContext.type !== 'expression') {
+      throw new Error(`Expected expression context for query: ${query}`);
+    }
+
+    const assignment = astContext.command.args[astContext.command.args.length - 1];
+
+    if (!assignment || !isAssignment(assignment)) {
+      throw new Error(`Expected assignment expression for query: ${query}`);
+    }
+
+    const root = getAssignmentExpressionRoot(assignment);
+
+    expect(root).toMatchObject({ type: 'column', name: 'salary' });
+  });
+});

src/platform/packages/shared/kbn-esql-language/src/commands/definitions/utils/expressions.ts

@@ -24,9 +24,9 @@ import { getColumnForASTNode } from './shared';
 import type { ESQLColumnData } from '../../registry/types';
 import { UnmappedFieldsStrategy } from '../../registry/types';
 import { TIME_SYSTEM_PARAMS } from './literals';
-import { isMarkerNode } from './ast';
 import { getUnmappedFieldType } from './settings';
 import { getPromqlFunctionDefinition } from './promql';
+import { isMarkerNode } from './ast';
 
 // #region type detection
 

src/platform/packages/shared/kbn-esql-language/src/commands/registry/completion/autocomplete.ts

@@ -19,7 +19,6 @@ import type {
 import { suggestForExpression } from '../../definitions/utils';
 import type { MapParameters } from '../../definitions/utils/autocomplete/map_expression';
 import { getCommandMapExpressionSuggestions } from '../../definitions/utils/autocomplete/map_expression';
-import { EDITOR_MARKER } from '../../definitions/constants';
 import {
   pipeCompleteItem,
   assignCompletionItem,
@@ -85,7 +84,7 @@ function getPosition(
     return { position: CompletionPosition.AFTER_COMMAND };
   }
 
-  const expressionRoot = prompt?.text !== EDITOR_MARKER ? prompt : undefined;
+  const expressionRoot = prompt;
 
   // (function, literal, or existing column) - handle as primaryExpression
   if (isFunctionExpression(expressionRoot) || isLiteral(prompt) || isExistingColumn) {

src/platform/packages/shared/kbn-esql-language/src/commands/registry/dissect/autocomplete.ts

@@ -8,14 +8,13 @@
  */
 import { i18n } from '@kbn/i18n';
 import type { ESQLAstAllCommands } from '@elastic/esql/types';
-import { Parser } from '@elastic/esql';
 import { withAutoSuggest } from '../../definitions/utils/autocomplete/helpers';
 import type { ICommandCallbacks } from '../types';
 import { pipeCompleteItem, colonCompleteItem, semiColonCompleteItem } from '../complete_items';
 import { type ISuggestionItem, type ICommandContext } from '../types';
 import { buildConstantsDefinitions } from '../../definitions/utils/literals';
 import { ESQL_STRING_TYPES } from '../../definitions/types';
-import { correctQuerySyntax, findAstPosition } from '../../definitions/utils/ast';
+import { findAutocompleteAstPosition } from '../../../language/shared/parse_for_autocomplete_query';
 
 const appendSeparatorCompletionItem: ISuggestionItem = withAutoSuggest({
   detail: i18n.translate('kbn-esql-language.esql.definitions.appendSeparatorDoc', {
@@ -39,9 +38,7 @@ export async function autocomplete(
   const commandArgs = command.args.filter((arg) => !Array.isArray(arg) && arg.type !== 'unknown');
 
   // If cursor is inside a string literal, don't suggest anything
-  const correctedQuery = correctQuerySyntax(innerText);
-  const { root } = Parser.parse(correctedQuery, { withFormatting: true });
-  const { node } = findAstPosition(root, innerText.length);
+  const { node } = findAutocompleteAstPosition(query, cursorPosition);
 
   if (node?.type === 'literal' && node.literalType === 'keyword') {
     return [];

src/platform/packages/shared/kbn-esql-language/src/commands/registry/fork/autocomplete.test.ts

@@ -29,8 +29,7 @@ import {
   ESQL_STRING_TYPES,
   ESQL_NUMBER_TYPES,
 } from '../../definitions/types';
-import { correctQuerySyntax, findAstPosition } from '../../definitions/utils/ast';
-import { Parser } from '@elastic/esql';
+import { findAutocompleteAstPosition } from '../../../language/shared/parse_for_autocomplete_query';
 
 const allEvalFnsForWhere = getFunctionSignaturesByReturnType(Location.WHERE, 'any', {
   scalar: true,
@@ -428,11 +427,8 @@ describe('FORK Autocomplete', () => {
 
       it('suggests pipe after complete subcommands', async () => {
         const assertSuggestsPipe = async (query: string) => {
-          const correctedQuery = correctQuerySyntax(query);
-          const { root } = Parser.parse(correctedQuery, { withFormatting: true });
-
           const cursorPosition = query.length;
-          const { command } = findAstPosition(root, cursorPosition);
+          const { command } = findAutocompleteAstPosition(query, cursorPosition);
           if (!command) {
             throw new Error('Command not found in the parsed query');
           }

src/platform/packages/shared/kbn-esql-language/src/commands/registry/from/autocomplete.test.ts

@@ -15,8 +15,7 @@ import {
 } from '../options/recommended_queries';
 import type { ICommandCallbacks } from '../types';
 import { autocomplete } from './autocomplete';
-import { correctQuerySyntax, findAstPosition } from '../../definitions/utils/ast';
-import { Parser } from '@elastic/esql';
+import { findAutocompleteAstPosition } from '../../../language/shared/parse_for_autocomplete_query';
 
 const metadataFields = [...METADATA_FIELDS].sort();
 
@@ -80,12 +79,8 @@ describe('FROM Autocomplete', () => {
       };
 
       const suggest = async (query: string) => {
-        const correctedQuery = correctQuerySyntax(query);
-        const { root } = Parser.parse(correctedQuery, { withFormatting: true });
-
         const cursorPosition = query.length;
-        const { command } = findAstPosition(root, cursorPosition);
-
+        const { command } = findAutocompleteAstPosition(query, cursorPosition);
         return autocomplete(query, command!, mockCallbacks, mockContext, cursorPosition);
       };
 
@@ -120,12 +115,8 @@ describe('FROM Autocomplete', () => {
 
     test('suggests comma or pipe after complete index name', async () => {
       const suggest = async (query: string) => {
-        const correctedQuery = correctQuerySyntax(query);
-        const { root } = Parser.parse(correctedQuery, { withFormatting: true });
-
         const cursorPosition = query.length;
-        const { command } = findAstPosition(root, cursorPosition);
-
+        const { command } = findAutocompleteAstPosition(query, cursorPosition);
         return autocomplete(query, command!, mockCallbacks, mockContext, cursorPosition);
       };
       const suggestions = (await suggest('from index')).map((s) => s.text);
@@ -176,10 +167,8 @@ describe('FROM Autocomplete', () => {
       );
       // View names appear when typing (fragment "my_")
       const getSuggestions = async (query: string) => {
-        const correctedQuery = correctQuerySyntax(query);
-        const { root } = Parser.parse(correctedQuery, { withFormatting: true });
         const cursorPosition = query.length;
-        const { command } = findAstPosition(root, cursorPosition);
+        const { command } = findAutocompleteAstPosition(query, cursorPosition);
         return autocomplete(query, command!, mockCallbacks, contextWithViews, cursorPosition);
       };
       const suggestions = (await getSuggestions('FROM my_')).map((s) => s.text);

src/platform/packages/shared/kbn-esql-language/src/commands/registry/grok/autocomplete.ts

@@ -8,14 +8,13 @@
  */
 import { i18n } from '@kbn/i18n';
 import type { ESQLAstAllCommands } from '@elastic/esql/types';
-import { Parser } from '@elastic/esql';
 import { withAutoSuggest } from '../../definitions/utils/autocomplete/helpers';
 import { commaCompleteItem, pipeCompleteItem } from '../complete_items';
 import type { ICommandCallbacks } from '../types';
 import type { ISuggestionItem, ICommandContext } from '../types';
 import { buildConstantsDefinitions } from '../../definitions/utils/literals';
 import { ESQL_STRING_TYPES } from '../../definitions/types';
-import { correctQuerySyntax, findAstPosition } from '../../definitions/utils/ast';
+import { findAutocompleteAstPosition } from '../../../language/shared/parse_for_autocomplete_query';
 
 export async function autocomplete(
   query: string,
@@ -28,9 +27,7 @@ export async function autocomplete(
   const commandArgs = command.args.filter((arg) => !Array.isArray(arg) && arg.type !== 'unknown');
 
   // If cursor is inside a string literal, don't suggest anything
-  const correctedQuery = correctQuerySyntax(innerText);
-  const { root } = Parser.parse(correctedQuery, { withFormatting: true });
-  const { node } = findAstPosition(root, innerText.length);
+  const { node } = findAutocompleteAstPosition(query, cursorPosition);
 
   if (node?.type === 'literal' && node.literalType === 'keyword') {
     return [];