Update the tests

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

@@ -587,7 +587,7 @@ export function addTypeHierarchyHandler(connection: Connection, sharedServices:
     connection.languages.typeHierarchy.onPrepare(
         createServerRequestHandler(
             async (services, document, params, cancelToken) => {
-                const result = await services.lsp.TypeHierarchyProvider?.prepareTypeHierarchy(document, params, cancelToken);
+                const result = await services.lsp?.TypeHierarchyProvider?.prepareTypeHierarchy(document, params, cancelToken);
                 return result ?? null;
             },
             sharedServices,
@@ -598,7 +598,7 @@ export function addTypeHierarchyHandler(connection: Connection, sharedServices:
     connection.languages.typeHierarchy.onSupertypes(
         createHierarchyRequestHandler(
             async (services, params, cancelToken) => {
-                const result = await services.lsp.TypeHierarchyProvider?.supertypes(params, cancelToken);
+                const result = await services.lsp?.TypeHierarchyProvider?.supertypes(params, cancelToken);
                 return result ?? null;
             },
             sharedServices
@@ -608,7 +608,7 @@ export function addTypeHierarchyHandler(connection: Connection, sharedServices:
     connection.languages.typeHierarchy.onSubtypes(
         createHierarchyRequestHandler(
             async (services, params, cancelToken) => {
-                const result = await services.lsp.TypeHierarchyProvider?.subtypes(params, cancelToken);
+                const result = await services.lsp?.TypeHierarchyProvider?.subtypes(params, cancelToken);
                 return result ?? null;
             },
             sharedServices
@@ -642,7 +642,7 @@ export function createHierarchyRequestHandler<P extends TypeHierarchySupertypesP
 }
 
 export function createServerRequestHandler<P extends { textDocument: TextDocumentIdentifier }, R, PR, E = void>(
-    serviceCall: (services: LangiumServices, document: LangiumDocument, params: P, cancelToken: CancellationToken) => HandlerResult<R, E>,
+    serviceCall: (services: LangiumCoreAndPartialLSPServices, document: LangiumDocument, params: P, cancelToken: CancellationToken) => HandlerResult<R, E>,
     sharedServices: LangiumSharedServices,
     targetState?: DocumentState
 ): ServerRequestHandler<P, R, PR, E> {
@@ -670,7 +670,7 @@ export function createServerRequestHandler<P extends { textDocument: TextDocumen
 }
 
 export function createRequestHandler<P extends { textDocument: TextDocumentIdentifier }, R, E = void>(
-    serviceCall: (services: LangiumServices, document: LangiumDocument, params: P, cancelToken: CancellationToken) => HandlerResult<R, E>,
+    serviceCall: (services: LangiumCoreAndPartialLSPServices, document: LangiumDocument, params: P, cancelToken: CancellationToken) => HandlerResult<R, E>,
     sharedServices: LangiumSharedServices,
     targetState?: DocumentState
 ): RequestHandler<P, R | null, E> {

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

@@ -9,21 +9,16 @@ import { Disposable } from '../utils/disposable.js';
 import type { ServiceRegistry } from '../service-registry.js';
 import type { LangiumSharedCoreServices } from '../services.js';
 import type { AstNode } from '../syntax-tree.js';
-import type { MaybePromise } from '../utils/promise-util.js';
-import type { Deferred } from '../utils/promise-util.js';
+import type { MaybePromise } from '../utils/promise-utils.js';
+import type { Deferred } from '../utils/promise-utils.js';
 import type { ValidationOptions } from '../validation/document-validator.js';
 import type { IndexManager } from '../workspace/index-manager.js';
 import type { LangiumDocument, LangiumDocuments, LangiumDocumentFactory } from './documents.js';
-import { CancellationToken, Disposable } from 'vscode-languageserver';
 import { MultiMap } from '../utils/collections.js';
-import type { MaybePromise } from '../utils/promise-utils.js';
-import { interruptAndCheck } from '../utils/promise-utils.js';
+import { OperationCancelled, interruptAndCheck } from '../utils/promise-utils.js';
 import { stream } from '../utils/stream.js';
 import type { URI } from '../utils/uri-utils.js';
-import type { ValidationOptions } from '../validation/document-validator.js';
 import { ValidationCategory } from '../validation/validation-registry.js';
-import type { IndexManager } from '../workspace/index-manager.js';
-import type { LangiumDocument, LangiumDocumentFactory, LangiumDocuments } from './documents.js';
 import { DocumentState } from './documents.js';
 
 export interface BuildOptions {
@@ -99,12 +94,13 @@ export interface DocumentBuilder {
     /**
      * Wait until the document specified by the {@link uri} has reached the specified state.
      *
-     * @param state The desired state. The promise won't resolve until the document has reached this state
-     * @param uri The specified URI that points to the document. If the URI does not exist, the promise will never resolve.
-     * @param cancelToken Optionally allows to cancel the wait operation, disposing any listeners in the process
+     * @param state The desired state. The promise won't resolve until the document has reached this state.
+     * @param uri The specified URI that points to the document. If the URI does not exist, the promise will resolve once the workspace has reached the specified state.
+     * @param cancelToken Optionally allows to cancel the wait operation, disposing any listeners in the process.
+     * @return The URI of the document that has reached the desired state, or `undefined` if the document does not exist.
      * @throws `OperationCancelled` if cancellation has been requested before the state has been reached
      */
-    waitUntil(state: DocumentState, uri?: URI, cancelToken?: CancellationToken): Promise<void>;
+    waitUntil(state: DocumentState, uri?: URI, cancelToken?: CancellationToken): Promise<URI | undefined>;
 }
 
 export type DocumentUpdateListener = (changed: URI[], deleted: URI[]) => void | Promise<void>
@@ -331,31 +327,36 @@ export class DefaultDocumentBuilder implements DocumentBuilder {
     }
 
     waitUntil(state: DocumentState, cancelToken?: CancellationToken): Promise<void>;
-    waitUntil(state: DocumentState, uri?: URI, cancelToken?: CancellationToken): Promise<void>;
-    waitUntil(state: DocumentState, uriOrToken?: URI | CancellationToken, 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> {
         let uri: URI | undefined = undefined;
         if (uriOrToken && 'path' in uriOrToken) {
             uri = uriOrToken;
         } else {
             cancelToken = uriOrToken;
         }
         cancelToken ??= CancellationToken.None;
-        if (this.currentState >= state) {
-            return Promise.resolve();
-        } else if (cancelToken.isCancellationRequested) {
-            return Promise.reject(OperationCancelled);
-        }
         if (uri) {
             const document = this.langiumDocuments.getDocument(uri);
             if (document && document.state > state) {
-                return Promise.resolve();
+                return Promise.resolve(uri);
             }
         }
+        if (this.currentState >= state) {
+            return Promise.resolve(undefined);
+        } else if (cancelToken.isCancellationRequested) {
+            return Promise.reject(OperationCancelled);
+        }
         return new Promise((resolve, reject) => {
             const buildDisposable = this.onBuildPhase(state, () => {
                 buildDisposable.dispose();
                 cancelDisposable.dispose();
-                resolve();
+                if (uri) {
+                    const document = this.langiumDocuments.getDocument(uri);
+                    resolve(document?.uri);
+                } else {
+                    resolve(undefined);
+                }
             });
             const cancelDisposable = cancelToken!.onCancellationRequested(() => {
                 buildDisposable.dispose();

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

@@ -11,7 +11,7 @@ import { setTextDocument } from 'langium/test';
 import { describe, expect, test } from 'vitest';
 import { CancellationTokenSource } from 'vscode-languageserver';
 import { TextDocument } from 'vscode-languageserver-textdocument';
-import { isOperationCancelled, DocumentState, EmptyFileSystem, URI, delayNextTick } from 'langium';
+import { isOperationCancelled, DocumentState, EmptyFileSystem, URI } from 'langium';
 import { createLangiumGrammarServices, createServicesForGrammar } from 'langium/grammar';
 import { setTextDocument } from 'langium/test';
 import { fail } from 'assert';
@@ -250,21 +250,19 @@ describe('DefaultDocumentBuilder', () => {
         documents.addDocument(document);
 
         const actual: string[] = [];
-        const expected: string[] = [];
         function wait(state: DocumentState): void {
-            expected.push('B' + state);
-            expected.push('W' + state);
             builder.onBuildPhase(state, async () => {
                 actual.push('B' + state);
-                await delayNextTick();
             });
             builder.waitUntil(state).then(() => actual.push('W' + state));
         }
-        for (let i = 2; i <= 6; i++) {
+        // Register listeners for all possible document states
+        // On each new state, there's supposed to be two new entries to the list
+        for (let i = DocumentState.IndexedContent; i <= DocumentState.Validated; i++) {
             wait(i);
         }
         await builder.build([document], { validation: true });
-        expect(actual).toEqual(expected);
+        expect(actual).toEqual(['B2', 'W2', 'B3', 'W3', 'B4', 'W4', 'B5', 'W5', 'B6', 'W6']);
     });
 
     test('`waitUntil` will correctly wait even though the build process has been cancelled', async () => {
@@ -280,10 +278,9 @@ describe('DefaultDocumentBuilder', () => {
         function wait(state: DocumentState): void {
             builder.onBuildPhase(state, async () => {
                 actual.push('B' + state);
-                await delayNextTick();
             });
         }
-        for (let i = 2; i <= 6; i++) {
+        for (let i = DocumentState.IndexedContent; i <= DocumentState.Validated; i++) {
             wait(i);
         }
         builder.waitUntil(DocumentState.ComputedScopes).then(() => cancelTokenSource.cancel());
@@ -293,6 +290,7 @@ describe('DefaultDocumentBuilder', () => {
         // Build twice but interrupt the first build after the computing scope phase
         try {
             await builder.build([document], { validation: true }, cancelTokenSource.token);
+            fail('The build is supposed to be cancelled');
         } catch {
             // build has been cancelled, ignore
         }
@@ -313,7 +311,7 @@ describe('DefaultDocumentBuilder', () => {
 
         const cancelTokenSource = new CancellationTokenSource();
         builder.waitUntil(DocumentState.IndexedReferences, cancelTokenSource.token).then(() => {
-            fail('This should have been cancelled');
+            fail('The test should fail here because the cancellation should reject the promise');
         }).catch(err => {
             expect(isOperationCancelled(err)).toBeTruthy();
         });