fix(language): ref resolution failure causes exception in validator

packages/ide/vscode/package.json

@@ -1,7 +1,7 @@
 {
     "name": "zenstack-v3",
     "publisher": "zenstack",
-    "version": "3.0.6",
+    "version": "3.0.8",
     "displayName": "ZenStack V3 Language Tools",
     "description": "VSCode extension for ZenStack (v3) ZModel language",
     "private": true,

packages/language/src/module.ts

@@ -9,6 +9,7 @@ import {
 } from 'langium/lsp';
 import { ZModelGeneratedModule, ZModelGeneratedSharedModule, ZModelLanguageMetaData } from './generated/module';
 import { ZModelValidator, registerValidationChecks } from './validator';
+import { ZModelDocumentBuilder } from './zmodel-document-builder';
 import { ZModelLinker } from './zmodel-linker';
 import { ZModelScopeComputation, ZModelScopeProvider } from './zmodel-scope';
 import { ZModelWorkspaceManager } from './zmodel-workspace-manager';
@@ -49,6 +50,7 @@ export type ZModelSharedServices = LangiumSharedServices;
 
 export const ZModelSharedModule: Module<ZModelSharedServices, DeepPartial<ZModelSharedServices>> = {
     workspace: {
+        DocumentBuilder: (services) => new ZModelDocumentBuilder(services),
         WorkspaceManager: (services) => new ZModelWorkspaceManager(services),
     },
 };

packages/language/src/utils.ts

@@ -1,4 +1,3 @@
-import { invariant } from '@zenstackhq/common-helpers';
 import { AstUtils, URI, type AstNode, type LangiumDocument, type LangiumDocuments, type Reference } from 'langium';
 import fs from 'node:fs';
 import path from 'path';
@@ -172,7 +171,10 @@ export function getRecursiveBases(
     }
     seen.add(decl);
     decl.mixins.forEach((mixin) => {
-        const baseDecl = mixin.ref;
+        // avoid using mixin.ref since this function can be called before linking
+        const baseDecl = decl.$container.declarations.find(
+            (d): d is TypeDef => isTypeDef(d) && d.name === mixin.$refText,
+        );
         if (baseDecl) {
             if (!includeDelegate && isDelegateModel(baseDecl)) {
                 return;
@@ -517,13 +519,15 @@ export function getAllFields(
 
     const fields: DataField[] = [];
     for (const mixin of decl.mixins) {
-        invariant(mixin.ref, `Mixin ${mixin.$refText} is not resolved`);
-        fields.push(...getAllFields(mixin.ref, includeIgnored, seen));
+        if (mixin.ref) {
+            fields.push(...getAllFields(mixin.ref, includeIgnored, seen));
+        }
     }
 
     if (isDataModel(decl) && decl.baseModel) {
-        invariant(decl.baseModel.ref, `Base model ${decl.baseModel.$refText} is not resolved`);
-        fields.push(...getAllFields(decl.baseModel.ref, includeIgnored, seen));
+        if (decl.baseModel.ref) {
+            fields.push(...getAllFields(decl.baseModel.ref, includeIgnored, seen));
+        }
     }
 
     fields.push(...decl.fields.filter((f) => includeIgnored || !hasAttribute(f, '@ignore')));
@@ -541,13 +545,15 @@ export function getAllAttributes(
 
     const attributes: DataModelAttribute[] = [];
     for (const mixin of decl.mixins) {
-        invariant(mixin.ref, `Mixin ${mixin.$refText} is not resolved`);
-        attributes.push(...getAllAttributes(mixin.ref, seen));
+        if (mixin.ref) {
+            attributes.push(...getAllAttributes(mixin.ref, seen));
+        }
     }
 
     if (isDataModel(decl) && decl.baseModel) {
-        invariant(decl.baseModel.ref, `Base model ${decl.baseModel.$refText} is not resolved`);
-        attributes.push(...getAllAttributes(decl.baseModel.ref, seen));
+        if (decl.baseModel.ref) {
+            attributes.push(...getAllAttributes(decl.baseModel.ref, seen));
+        }
     }
 
     attributes.push(...decl.attributes);

packages/language/src/validator.ts

@@ -1,6 +1,4 @@
-/* eslint-disable @typescript-eslint/no-unused-expressions */
-
-import type { AstNode, LangiumDocument, ValidationAcceptor, ValidationChecks } from 'langium';
+import type { ValidationAcceptor, ValidationChecks } from 'langium';
 import type {
     Attribute,
     DataModel,
@@ -50,54 +48,39 @@ export function registerValidationChecks(services: ZModelServices) {
 export class ZModelValidator {
     constructor(protected readonly services: ZModelServices) {}
 
-    private shouldCheck(node: AstNode) {
-        let doc: LangiumDocument | undefined;
-        let currNode: AstNode | undefined = node;
-        while (currNode) {
-            if (currNode.$document) {
-                doc = currNode.$document;
-                break;
-            }
-            currNode = currNode.$container;
-        }
-
-        return doc?.parseResult.lexerErrors.length === 0 && doc?.parseResult.parserErrors.length === 0;
-    }
-
     checkModel(node: Model, accept: ValidationAcceptor): void {
-        this.shouldCheck(node) &&
-            new SchemaValidator(this.services.shared.workspace.LangiumDocuments).validate(node, accept);
+        new SchemaValidator(this.services.shared.workspace.LangiumDocuments).validate(node, accept);
     }
 
     checkDataSource(node: DataSource, accept: ValidationAcceptor): void {
-        this.shouldCheck(node) && new DataSourceValidator().validate(node, accept);
+        new DataSourceValidator().validate(node, accept);
     }
 
     checkDataModel(node: DataModel, accept: ValidationAcceptor): void {
-        this.shouldCheck(node) && new DataModelValidator().validate(node, accept);
+        new DataModelValidator().validate(node, accept);
     }
 
     checkTypeDef(node: TypeDef, accept: ValidationAcceptor): void {
-        this.shouldCheck(node) && new TypeDefValidator().validate(node, accept);
+        new TypeDefValidator().validate(node, accept);
     }
 
     checkEnum(node: Enum, accept: ValidationAcceptor): void {
-        this.shouldCheck(node) && new EnumValidator().validate(node, accept);
+        new EnumValidator().validate(node, accept);
     }
 
     checkAttribute(node: Attribute, accept: ValidationAcceptor): void {
-        this.shouldCheck(node) && new AttributeValidator().validate(node, accept);
+        new AttributeValidator().validate(node, accept);
     }
 
     checkExpression(node: Expression, accept: ValidationAcceptor): void {
-        this.shouldCheck(node) && new ExpressionValidator().validate(node, accept);
+        new ExpressionValidator().validate(node, accept);
     }
 
     checkFunctionInvocation(node: InvocationExpr, accept: ValidationAcceptor): void {
-        this.shouldCheck(node) && new FunctionInvocationValidator().validate(node, accept);
+        new FunctionInvocationValidator().validate(node, accept);
     }
 
     checkFunctionDecl(node: FunctionDecl, accept: ValidationAcceptor): void {
-        this.shouldCheck(node) && new FunctionDeclValidator().validate(node, accept);
+        new FunctionDeclValidator().validate(node, accept);
     }
 }

packages/language/src/zmodel-document-builder.ts

@@ -0,0 +1,22 @@
+import { DefaultDocumentBuilder, type BuildOptions, type LangiumDocument } from 'langium';
+
+export class ZModelDocumentBuilder extends DefaultDocumentBuilder {
+    override buildDocuments(documents: LangiumDocument[], options: BuildOptions, cancelToken: any): Promise<void> {
+        return super.buildDocuments(
+            documents,
+            {
+                ...options,
+                validation:
+                    // force overriding validation options
+                    options.validation === false || options.validation === undefined
+                        ? options.validation
+                        : {
+                              stopAfterLexingErrors: true,
+                              stopAfterParsingErrors: true,
+                              stopAfterLinkingErrors: true,
+                          },
+            },
+            cancelToken,
+        );
+    }
+}

packages/language/test/expression-validation.test.ts

@@ -2,7 +2,7 @@ import { describe, it } from 'vitest';
 import { loadSchema, loadSchemaWithError } from './utils';
 
 describe('Expression Validation Tests', () => {
-    it('should reject model comparison', async () => {
+    it('should reject model comparison1', async () => {
         await loadSchemaWithError(
             `
             model User {
@@ -15,14 +15,15 @@ describe('Expression Validation Tests', () => {
                 id Int @id
                 title String
                 author User @relation(fields: [authorId], references: [id])
+                authorId Int
                 @@allow('all', author == this)
             }
         `,
             'comparison between models is not supported',
         );
     });
 
-    it('should reject model comparison', async () => {
+    it('should reject model comparison2', async () => {
         await loadSchemaWithError(
             `
             model User {