st-language: refined waitUntil with document URI to fire as soon as document reaches required state

added corresponding test

Author: sailingKieler

packages/langium/src/lsp/language-server.ts

@@ -731,8 +731,10 @@ export function createRequestHandler<P extends { textDocument: TextDocumentIdent
 
 async function waitUntilPhase<E>(services: LangiumSharedServices, cancelToken: CancellationToken, uri?: URI, targetState?: DocumentState): Promise<ResponseError<E> | undefined> {
     if (targetState !== undefined) {
+        const workspaceManager = services.workspace.WorkspaceManager;
         const documentBuilder = services.workspace.DocumentBuilder;
         try {
+            await workspaceManager.ready; // mandatory if awaiting the state of a document (uri !== undefined) while the LS is starting
             await documentBuilder.waitUntil(targetState, uri, cancelToken);
         } catch (err) {
             return responseError(err);

packages/langium/src/workspace/document-builder.ts

@@ -4,6 +4,7 @@
  * terms of the MIT License, which is available in the project root.
  ******************************************************************************/
 
+import { LSPErrorCodes, ResponseError } from 'vscode-languageserver-protocol';
 import { CancellationToken } from '../utils/cancellation.js';
 import { Disposable } from '../utils/disposable.js';
 import type { ServiceRegistry } from '../service-registry.js';
@@ -490,8 +491,8 @@ export class DefaultDocumentBuilder implements DocumentBuilder {
     }
 
     waitUntil(state: DocumentState, cancelToken?: CancellationToken): Promise<void>;
-    waitUntil(state: DocumentState, uri?: URI, cancelToken?: CancellationToken): Promise<URI | undefined>;
-    waitUntil(state: DocumentState, uriOrToken?: URI | CancellationToken, cancelToken?: CancellationToken): Promise<URI | undefined | void> {
+    waitUntil(state: DocumentState, uri?: URI, cancelToken?: CancellationToken): Promise<URI>;
+    waitUntil(state: DocumentState, uriOrToken?: URI | CancellationToken, cancelToken?: CancellationToken): Promise<URI | void> {
         let uri: URI | undefined = undefined;
         if (uriOrToken && 'path' in uriOrToken) {
             uri = uriOrToken;
@@ -500,26 +501,67 @@ export class DefaultDocumentBuilder implements DocumentBuilder {
         }
         cancelToken ??= CancellationToken.None;
         if (uri) {
-            const document = this.langiumDocuments.getDocument(uri);
-            if (document && document.state >= state) {
-                return Promise.resolve(uri);
-            }
+            return this.awaitDocumentState(state, uri, cancelToken);
+
+        } else {
+            return this.awaitBuilderState(state, cancelToken);
+        }
+    }
+
+    protected awaitDocumentState(state: DocumentState, uri: URI, cancelToken: CancellationToken): Promise<URI> {
+        const document = this.langiumDocuments.getDocument(uri);
+        if (!document) {
+            return Promise.reject(
+                new ResponseError(
+                    LSPErrorCodes.ServerCancelled,
+                    `No document found for URI: ${uri.toString()}`
+                )
+            );
+
+        } else if (document.state >= state) {
+            return Promise.resolve(uri);
+
+        } else if (cancelToken.isCancellationRequested) {
+            return Promise.reject(OperationCancelled);
+
+        } else if (this.currentState >= state && state > document.state) {
+            // this would imply that the document has been excluded from linking or validation, for example;
+            // this should never occur, the LS need to make sure that the affected document is properly built,
+            //  alternatively, the build state requirement need to be relaxed.
+            return Promise.reject(
+                new ResponseError(
+                    LSPErrorCodes.RequestFailed,
+                    `Document state of ${uri.toString()} is ${DocumentState[document.state]}, requiring ${DocumentState[state]}, but workspace state is already ${DocumentState[this.currentState]}. Returning undefined.`
+                )
+            );
         }
+        return new Promise((resolve, reject) => {
+            const buildDisposable = this.onDocumentPhase(state, (doc) => {
+                if (doc.uri.toString() === uri.toString()) {
+                    buildDisposable.dispose();
+                    cancelDisposable.dispose();
+                    resolve(doc.uri);
+                }
+            });
+            const cancelDisposable = cancelToken!.onCancellationRequested(() => {
+                buildDisposable.dispose();
+                cancelDisposable.dispose();
+                reject(OperationCancelled);
+            });
+        });
+    }
+
+    protected awaitBuilderState(state: DocumentState, cancelToken: CancellationToken): Promise<void> {
         if (this.currentState >= state) {
-            return Promise.resolve(undefined);
+            return Promise.resolve();
         } else if (cancelToken.isCancellationRequested) {
             return Promise.reject(OperationCancelled);
         }
         return new Promise((resolve, reject) => {
             const buildDisposable = this.onBuildPhase(state, () => {
                 buildDisposable.dispose();
                 cancelDisposable.dispose();
-                if (uri) {
-                    const document = this.langiumDocuments.getDocument(uri);
-                    resolve(document?.uri);
-                } else {
-                    resolve(undefined);
-                }
+                resolve();
             });
             const cancelDisposable = cancelToken!.onCancellationRequested(() => {
                 buildDisposable.dispose();

packages/langium/test/workspace/document-builder.test.ts

@@ -497,6 +497,31 @@ describe('DefaultDocumentBuilder', () => {
         ).resolves.toEqual(documentUri);
     });
 
+    test('`waitUntil` on document fires as soon as document reaches required state.', async () => {
+        const services = await createServices();
+        const documentFactory = services.shared.workspace.LangiumDocumentFactory;
+        const documents = services.shared.workspace.LangiumDocuments;
+        const builder = services.shared.workspace.DocumentBuilder;
+
+        const document = documentFactory.fromString<Model>(`
+            foo 1 A
+            foo 11 B
+            bar A
+            bar B
+        `, URI.parse('file:///test1.txt'));
+        documents.addDocument(document);
+        const document2 = documentFactory.fromString<Model>('', URI.parse('file:///test2.txt'));
+        documents.addDocument(document2);
+
+        const states: DocumentState[] = [];
+        builder.waitUntil(DocumentState.Linked, document.uri).then(() => {
+            states.push(document.state, document2.state);
+        });
+
+        await builder.build(documents.all.toArray());
+        expect(states).toEqual([ DocumentState.Linked, DocumentState.ComputedScopes ]);
+    });
+
     test('`onDocumentPhase` always triggers before the respective `onBuildPhase`', async () => {
         const services = await createServices();
         const documentFactory = services.shared.workspace.LangiumDocumentFactory;