fix: clean up plugin loading and fix VSCode extension

packages/cli/src/actions/action-utils.ts

@@ -1,4 +1,4 @@
-import { createZModelServices, loadDocument, type ZModelServices } from '@zenstackhq/language';
+import { loadDocument, type ZModelServices } from '@zenstackhq/language';
 import { isDataSource, isPlugin, Model } from '@zenstackhq/language/ast';
 import { getLiteral } from '@zenstackhq/language/utils';
 import { PrismaSchemaGenerator } from '@zenstackhq/sdk';
@@ -37,9 +37,7 @@ export function getSchemaFile(file?: string) {
 }
 
 export async function loadSchemaDocument(schemaFile: string) {
-    const { ZModelLanguage: services } = createZModelServices();
-    const pluginDocs = await getPluginDocuments(services, schemaFile);
-    const loadResult = await loadDocument(schemaFile, pluginDocs);
+    const loadResult = await loadDocument(schemaFile);
     if (!loadResult.success) {
         loadResult.errors.forEach((err) => {
             console.error(colors.red(err));

packages/ide/vscode/package.json

@@ -11,6 +11,7 @@
     },
     "scripts": {
         "build": "tsc --noEmit && tsup",
+        "watch": "tsup --watch",
         "lint": "eslint src --ext ts",
         "vscode:publish": "pnpm build && vsce publish --no-dependencies --follow-symlinks",
         "vscode:package": "pnpm build && vsce package --no-dependencies --follow-symlinks"

packages/ide/vscode/src/language-server/main.ts

@@ -7,10 +7,13 @@ import { createConnection, ProposedFeatures } from 'vscode-languageserver/node.j
 const connection = createConnection(ProposedFeatures.all);
 
 // Inject the shared services and language-specific services
-const { shared } = createZModelLanguageServices({
-    connection,
-    ...NodeFileSystem,
-});
+const { shared } = createZModelLanguageServices(
+    {
+        connection,
+        ...NodeFileSystem,
+    },
+    true,
+);
 
 // Start the language server with the shared services
 startLanguageServer(shared);

packages/language/src/document.ts

@@ -0,0 +1,202 @@
+import { isAstNode, URI, type AstNode, type LangiumDocument, type LangiumDocuments, type Mutable } from 'langium';
+import fs from 'node:fs';
+import path from 'node:path';
+import { fileURLToPath } from 'node:url';
+import { isDataSource, Model } from './ast';
+import { STD_LIB_MODULE_NAME } from './constants';
+import { createZModelServices } from './module';
+import { getDataModelAndTypeDefs, getDocument, hasAttribute, resolveImport, resolveTransitiveImports } from './utils';
+
+/**
+ * Loads ZModel document from the given file name. Include the additional document
+ * files if given.
+ */
+export async function loadDocument(
+    fileName: string,
+    additionalModelFiles: string[] = [],
+): Promise<
+    { success: true; model: Model; warnings: string[] } | { success: false; errors: string[]; warnings: string[] }
+> {
+    const { ZModelLanguage: services } = createZModelServices(false);
+    const extensions = services.LanguageMetaData.fileExtensions;
+    if (!extensions.includes(path.extname(fileName))) {
+        return {
+            success: false,
+            errors: ['invalid schema file extension'],
+            warnings: [],
+        };
+    }
+
+    if (!fs.existsSync(fileName)) {
+        return {
+            success: false,
+            errors: ['schema file does not exist'],
+            warnings: [],
+        };
+    }
+
+    // load standard library
+
+    // isomorphic __dirname
+    const _dirname = typeof __dirname !== 'undefined' ? __dirname : path.dirname(fileURLToPath(import.meta.url));
+    const stdLib = await services.shared.workspace.LangiumDocuments.getOrCreateDocument(
+        URI.file(path.resolve(path.join(_dirname, '../res', STD_LIB_MODULE_NAME))),
+    );
+
+    // load the document
+    const langiumDocuments = services.shared.workspace.LangiumDocuments;
+    const document = await langiumDocuments.getOrCreateDocument(URI.file(path.resolve(fileName)));
+
+    // load imports
+    const importedURIs = await loadImports(document, langiumDocuments);
+    const importedDocuments: LangiumDocument[] = [];
+    for (const uri of importedURIs) {
+        importedDocuments.push(await langiumDocuments.getOrCreateDocument(uri));
+    }
+
+    // build the document together with standard library, additional modules, and imported documents
+
+    // load additional model files
+    const additionalDocs = await Promise.all(
+        additionalModelFiles.map((file) =>
+            services.shared.workspace.LangiumDocuments.getOrCreateDocument(URI.file(path.resolve(file))),
+        ),
+    );
+
+    await services.shared.workspace.DocumentBuilder.build([stdLib, ...additionalDocs, document, ...importedDocuments], {
+        validation: {
+            stopAfterLexingErrors: true,
+            stopAfterParsingErrors: true,
+            stopAfterLinkingErrors: true,
+        },
+    });
+
+    const diagnostics = langiumDocuments.all
+        .flatMap((doc) => (doc.diagnostics ?? []).map((diag) => ({ doc, diag })))
+        .filter(({ diag }) => diag.severity === 1 || diag.severity === 2)
+        .toArray();
+
+    const errors: string[] = [];
+    const warnings: string[] = [];
+
+    if (diagnostics.length > 0) {
+        for (const { doc, diag } of diagnostics) {
+            const message = `${path.relative(process.cwd(), doc.uri.fsPath)}:${
+                diag.range.start.line + 1
+            }:${diag.range.start.character + 1} - ${diag.message}`;
+
+            if (diag.severity === 1) {
+                errors.push(message);
+            } else {
+                warnings.push(message);
+            }
+        }
+    }
+
+    if (errors.length > 0) {
+        return {
+            success: false,
+            errors,
+            warnings,
+        };
+    }
+
+    const model = document.parseResult.value as Model;
+
+    // merge all declarations into the main document
+    const imported = mergeImportsDeclarations(langiumDocuments, model);
+
+    // remove imported documents
+    imported.forEach((model) => {
+        langiumDocuments.deleteDocument(model.$document!.uri);
+        services.shared.workspace.IndexManager.remove(model.$document!.uri);
+    });
+
+    // extra validation after merging imported declarations
+    const additionalErrors = validationAfterImportMerge(model);
+    if (additionalErrors.length > 0) {
+        return {
+            success: false,
+            errors: additionalErrors,
+            warnings,
+        };
+    }
+
+    return {
+        success: true,
+        model: document.parseResult.value as Model,
+        warnings,
+    };
+}
+
+async function loadImports(document: LangiumDocument, documents: LangiumDocuments, uris: Set<string> = new Set()) {
+    const uriString = document.uri.toString();
+    if (!uris.has(uriString)) {
+        uris.add(uriString);
+        const model = document.parseResult.value as Model;
+        for (const imp of model.imports) {
+            const importedModel = resolveImport(documents, imp);
+            if (importedModel) {
+                const importedDoc = getDocument(importedModel);
+                await loadImports(importedDoc, documents, uris);
+            }
+        }
+    }
+    return Array.from(uris)
+        .filter((x) => uriString != x)
+        .map((e) => URI.parse(e));
+}
+
+function mergeImportsDeclarations(documents: LangiumDocuments, model: Model) {
+    const importedModels = resolveTransitiveImports(documents, model);
+
+    const importedDeclarations = importedModels.flatMap((m) => m.declarations);
+    model.declarations.push(...importedDeclarations);
+
+    // remove import directives
+    model.imports = [];
+
+    // fix $container, $containerIndex, and $containerProperty
+    linkContentToContainer(model);
+
+    return importedModels;
+}
+
+function linkContentToContainer(node: AstNode): void {
+    for (const [name, value] of Object.entries(node)) {
+        if (!name.startsWith('$')) {
+            if (Array.isArray(value)) {
+                value.forEach((item, index) => {
+                    if (isAstNode(item)) {
+                        (item as Mutable<AstNode>).$container = node;
+                        (item as Mutable<AstNode>).$containerProperty = name;
+                        (item as Mutable<AstNode>).$containerIndex = index;
+                    }
+                });
+            } else if (isAstNode(value)) {
+                (value as Mutable<AstNode>).$container = node;
+                (value as Mutable<AstNode>).$containerProperty = name;
+            }
+        }
+    }
+}
+
+function validationAfterImportMerge(model: Model) {
+    const errors: string[] = [];
+    const dataSources = model.declarations.filter((d) => isDataSource(d));
+    if (dataSources.length === 0) {
+        errors.push('Validation error: schema must have a datasource declaration');
+    } else {
+        if (dataSources.length > 1) {
+            errors.push('Validation error: multiple datasource declarations are not allowed');
+        }
+    }
+
+    // at most one `@@auth` model
+    const decls = getDataModelAndTypeDefs(model, true);
+    const authDecls = decls.filter((d) => hasAttribute(d, '@@auth'));
+    if (authDecls.length > 1) {
+        errors.push('Validation error: Multiple `@@auth` declarations are not allowed');
+    }
+    return errors;
+}