[ES|QL] Fix nested editor marker leaks in autocomplete and introduce a shared autocomplete parsing boundary

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

@@ -80,23 +80,61 @@ const removeMarkerNode = (node: ESQLAstExpression) => {
   });
 };
 
-export function removeMarkerArgFromArgsList<T extends ESQLSingleAstItem | ESQLAstAllCommands>(
+function removeMarkerArgFromArgsList<T extends ESQLSingleAstItem | ESQLAstAllCommands>(
   node: T | undefined
 ) {
   if (!node) {
     return;
   }
   if (node.type === 'command' || node.type === 'option' || node.type === 'function') {
-    return {
+    const cleanedNode = {
       ...node,
+      text: node.text.replace(EDITOR_MARKER, ''),
       args: node.args.filter(isNotMarkerNodeOrArray).map(mapToNonMarkerNode),
     };
+
+    if (cleanedNode.type === 'command' && 'expression' in cleanedNode && cleanedNode.expression) {
+      cleanedNode.expression = isMarkerNode(cleanedNode.expression)
+        ? undefined
+        : (mapToNonMarkerNode(cleanedNode.expression) as ESQLAstExpression);
+    }
+
+    return cleanedNode;
   }
   return node;
 }
 
 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,
+      text: arg.text.replace(EDITOR_MARKER, ''),
+      args: arg.args.filter(isNotMarkerNodeOrArray).map(mapToNonMarkerNode) as typeof arg.args,
+    } as typeof arg;
+  }
+
+  if ('values' in arg) {
+    return {
+      ...arg,
+      text: arg.text.replace(EDITOR_MARKER, ''),
+      values: arg.values
+        .filter((value) => !isMarkerNode(value))
+        .map((value) => mapToNonMarkerNode(value) as ESQLAstExpression) as typeof arg.values,
+    } as typeof arg;
+  }
+
+  if ('text' in arg && typeof arg.text === 'string' && arg.text.includes(EDITOR_MARKER)) {
+    return {
+      ...arg,
+      text: arg.text.replace(EDITOR_MARKER, ''),
+    } 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/autocomplete/expressions/positions/in_function.ts

@@ -32,7 +32,12 @@ export async function suggestInFunction(ctx: ExpressionContext): Promise<ISugges
   } = ctx;
 
   const functionExpression = expressionRoot as ESQLFunction;
-  const paramContext = buildExpressionFunctionParameterContext(functionExpression, context);
+  const startingNewParam = STARTING_NEW_PARAM_REGEX.test(innerText);
+  const paramContext = buildExpressionFunctionParameterContext(
+    functionExpression,
+    context,
+    startingNewParam
+  );
 
   if (!paramContext) {
     return [];

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

@@ -13,7 +13,6 @@ import { nullCheckOperators, inOperators } from '../../../all_operators';
 import type { ExpressionContext, FunctionParameterContext } from './types';
 import type { ICommandContext, ISuggestionItem } from '../../../../registry/types';
 import { getFunctionDefinition } from '../..';
-import { EDITOR_MARKER } from '../../../constants';
 import { resolveArgumentTypes } from '../../expressions';
 import type { SupportedDataType } from '../../../types';
 import {
@@ -63,7 +62,8 @@ export function matchesSpecialFunction(name: string, expected: SpecialFunctionNa
  */
 export function buildExpressionFunctionParameterContext(
   fn: ESQLFunction,
-  context?: ICommandContext
+  context?: ICommandContext,
+  shouldGetNextArgument = false
 ): FunctionParameterContext | null {
   const fnDefinition = getFunctionDefinition(fn.name);
 
@@ -76,7 +76,6 @@ export function buildExpressionFunctionParameterContext(
     unmappedFieldsStrategy: context?.unmappedFieldsStrategy,
   });
 
-  const shouldGetNextArgument = fn.text.includes(EDITOR_MARKER);
   let argIndex = Math.max(fn.args.length, 0);
   if (!shouldGetNextArgument && argIndex) {
     argIndex -= 1;

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/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', {
@@ -38,9 +37,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);
       };
 
@@ -124,12 +119,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);
@@ -188,10 +179,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/fuse/autocomplete.ts

@@ -19,7 +19,6 @@ import {
   noneValueCompleteItem,
 } from '../complete_items';
 import { withAutoSuggest } from '../../definitions/utils/autocomplete/helpers';
-import { EDITOR_MARKER } from '../../definitions/constants';
 import { ESQL_STRING_TYPES } from '../../definitions/types';
 import { columnExists, findFinalWord } from '../../definitions/utils/autocomplete/helpers';
 import type { ICommandCallbacks } from '../types';
@@ -40,6 +39,9 @@ enum FusePosition {
   WITH = 'with',
 }
 
+// After `KEY BY field,`, suggest the next KEY BY field.
+const KEY_BY_TRAILING_COMMA_REGEX = /\bkey\s+by(?:\s+\S+,)+\s*$/i;
+
 function getPosition(innerText: string, command: ESQLAstFuseCommand): FusePosition {
   const { scoreBy, keyBy, groupBy, withOption } = extractFuseArgs(command);
 
@@ -54,7 +56,7 @@ function getPosition(innerText: string, command: ESQLAstFuseCommand): FusePositi
   if (
     (keyBy && keyBy.incomplete) ||
     immediatelyAfterOptionFieldsList(innerText, 'key by') ||
-    keyBy?.text.includes(EDITOR_MARKER)
+    KEY_BY_TRAILING_COMMA_REGEX.test(innerText)
   ) {
     return FusePosition.KEY_BY;
   }