[ES|QL] fix query highlight when wrapped in multi-line

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

@@ -11,11 +11,12 @@ import { ESQLState } from './esql_state';
 
 /** @internal **/
 export class ESQLLineTokens implements monaco.languages.ILineTokens {
-  endState: monaco.languages.IState;
+  endState: ESQLState;
   tokens: monaco.languages.IToken[];
 
-  constructor(tokens: monaco.languages.IToken[]) {
+  constructor(tokens: monaco.languages.IToken[], line: number) {
     this.endState = new ESQLState();
+    this.endState.setLineNumber(line);
     this.tokens = tokens;
   }
 }

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

@@ -10,8 +10,20 @@ import { monaco } from '../../../monaco_imports';
 
 /** @internal **/
 export class ESQLState implements monaco.languages.IState {
+  private lastLine: number = 0;
+
+  setLineNumber(n: number) {
+    this.lastLine = n;
+  }
+
+  getLineNumber() {
+    return this.lastLine;
+  }
+
   clone(): monaco.languages.IState {
-    return new ESQLState();
+    const newState = new ESQLState();
+    newState.setLineNumber(this.lastLine);
+    return newState;
   }
 
   equals(other: monaco.languages.IState): boolean {

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

@@ -24,10 +24,16 @@ export class ESQLTokensProvider implements monaco.languages.TokensProvider {
     return new ESQLState();
   }
 
-  tokenize(line: string, state: monaco.languages.IState): monaco.languages.ILineTokens {
+  tokenize(line: string, prevState: ESQLState): monaco.languages.ILineTokens {
     const errorStartingPoints: number[] = [];
     const errorListener = new ANTLREErrorListener();
-    const inputStream = CharStreams.fromString(line);
+    // This has the drawback of not styling any ESQL wrong query as
+    // | from ...
+    const cleanedLine =
+      prevState.getLineNumber() && line.trimStart()[0] === '|'
+        ? line.trimStart().substring(1)
+        : line;
+    const inputStream = CharStreams.fromString(cleanedLine);
     const lexer = getLexer(inputStream, errorListener);
 
     let done = false;
@@ -63,6 +69,6 @@ export class ESQLTokensProvider implements monaco.languages.TokensProvider {
 
     myTokens.sort((a, b) => a.startIndex - b.startIndex);
 
-    return new ESQLLineTokens(myTokens);
+    return new ESQLLineTokens(myTokens, prevState.getLineNumber() + 1);
   }
 }