[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

@@ -19,6 +19,7 @@ import type {
   ESQLAstQueryExpression,
 } from '@elastic/esql/types';
 import { EDITOR_MARKER } from '../constants';
+import { TRAILING_COMMA_REGEX } from './shared';
 
 export function isMarkerNode(node: ESQLAstItem | undefined): boolean {
   if (Array.isArray(node)) {
@@ -80,23 +81,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 {
@@ -277,7 +316,6 @@ export function correctQuerySyntax(_query: string) {
   return query;
 }
 
-const PROMQL_TRAILING_COMMA_REGEX = /,\s*$/;
 const PROMQL_TRAILING_COLON_REGEX = /:\s*$/;
 
 /**
@@ -290,7 +328,7 @@ export function correctPromqlQuerySyntax(input: string): string {
 
   if (
     !query.includes(EDITOR_MARKER) &&
-    (PROMQL_TRAILING_COMMA_REGEX.test(query) || PROMQL_TRAILING_COLON_REGEX.test(query))
+    (TRAILING_COMMA_REGEX.test(query) || PROMQL_TRAILING_COLON_REGEX.test(query))
   ) {
     query += ` ${EDITOR_MARKER}`;
   }

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

@@ -13,9 +13,7 @@ import { suggestForExpression } from '../suggestion_engine';
 import type { ExpressionContext } from '../types';
 import type { ISuggestionItem } from '../../../../../registry/types';
 import { buildExpressionFunctionParameterContext } from '../utils';
-
-/** Matches comma followed by optional whitespace at end of text */
-const STARTING_NEW_PARAM_REGEX = /,\s*$/;
+import { TRAILING_COMMA_REGEX } from '../../../shared';
 
 /** Suggests completions when cursor is inside a function call (e.g., CONCAT(field1, /)) */
 export async function suggestInFunction(ctx: ExpressionContext): Promise<ISuggestionItem[]> {
@@ -32,7 +30,12 @@ export async function suggestInFunction(ctx: ExpressionContext): Promise<ISugges
   } = ctx;
 
   const functionExpression = expressionRoot as ESQLFunction;
-  const paramContext = buildExpressionFunctionParameterContext(functionExpression, context);
+  const startingNewParam = TRAILING_COMMA_REGEX.test(innerText);
+  const paramContext = buildExpressionFunctionParameterContext(
+    functionExpression,
+    context,
+    startingNewParam
+  );
 
   if (!paramContext) {
     return [];
@@ -69,7 +72,7 @@ function determineTargetExpression(
   innerText: string
 ): ESQLSingleAstItem | undefined {
   const { args } = functionExpression;
-  const startingNewParam = STARTING_NEW_PARAM_REGEX.test(innerText);
+  const startingNewParam = TRAILING_COMMA_REGEX.test(innerText);
   const firstArgEmpty = isFirstArgumentEmpty(args, innerText);
 
   const lastArg = args[args.length - 1];

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/autocomplete/fields_list.ts

@@ -19,6 +19,7 @@ import type { ICommandCallbacks, ICommandContext, ISuggestionItem } from '../../
 import { getAssignmentExpressionRoot } from '../expressions';
 import { suggestForExpression } from './expressions';
 import { withAutoSuggest } from './helpers';
+import { TRAILING_COMMA_REGEX } from '../shared';
 import type { ExpressionContextOptions } from './expressions/types';
 
 const ENDS_WITH_WHITESPACE_REGEX = /\s$/;
@@ -54,7 +55,7 @@ export async function suggestFieldsList(
   const innerText = query.substring(0, cursorPosition);
   const lastField = fieldList[fieldList.length - 1];
 
-  const endsWithComma = /,\s*$/.test(innerText);
+  const endsWithComma = TRAILING_COMMA_REGEX.test(innerText);
   const withinFunction =
     lastField && isFunctionExpression(lastField) && within(innerText.length, lastField);
   const startingNewExpression = endsWithComma && !withinFunction;

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

@@ -15,6 +15,7 @@ import type {
   PromQLLabel,
   PromQLSelector,
 } from '@elastic/esql';
+import { TRAILING_COMMA_REGEX } from '../../shared';
 import type { PromQLFunctionParamType } from '../../../types';
 import type { CursorMatch, PromqlDetailedPosition } from './types';
 import {
@@ -41,7 +42,6 @@ import {
 } from '../../promql';
 import { promqlOperatorDefinitions } from '../../../generated/promql_operators';
 
-const TRAILING_COMMA_WITH_SPACES_REGEX = /,\s*$/;
 const SELECTOR_DURATION_START_REGEX = /^\s*\[$/;
 
 // Builds a regex pattern like `\+|\-|\*|==|and|or|...` from all binary operator definitions.
@@ -132,7 +132,7 @@ export function getQueryPosition(
       node.type === 'identifier' &&
       parent?.type === 'label-map' &&
       innermostFunc &&
-      TRAILING_COMMA_WITH_SPACES_REGEX.test(textBeforeCursor);
+      TRAILING_COMMA_REGEX.test(textBeforeCursor);
 
     if (!skipLabel) {
       const labelPos = resolveLabelPosition(match, cursor, textBeforeCursor, getSignatureTypes());

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/shared.ts

@@ -14,6 +14,9 @@ import type { ESQLUserDefinedColumn, ICommandContext } from '../../registry/type
 import { getLastNonWhitespaceChar } from './autocomplete/helpers';
 import type { SupportedDataType } from '../types';
 
+/** Matches a trailing comma (with optional whitespace) at end of text — used to detect "starting a new argument/expression" */
+export const TRAILING_COMMA_REGEX = /,\s*$/;
+
 export const techPreviewLabel = i18n.translate(
   'kbn-esql-language.esql.autocomplete.techPreviewLabel',
   {

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');
           }