fix: vscode url handling, loading imported models

Author: ymc9

packages/ide/vscode/src/extension/zenstack-auth-provider.ts

@@ -11,6 +11,7 @@ interface JWTClaims {
 export const AUTH_PROVIDER_ID = 'ZenStack';
 export const AUTH_URL = 'https://accounts.zenstack.dev';
 export const API_URL = 'https://api.zenstack.dev';
+const EXTENSION_ID = 'zenstack.zenstack-v3';
 
 export class ZenStackAuthenticationProvider implements vscode.AuthenticationProvider, vscode.Disposable {
     private _onDidChangeSessions =
@@ -111,19 +112,21 @@ export class ZenStackAuthenticationProvider implements vscode.AuthenticationProv
                         reject(new Error('User Cancelled'));
                     });
 
-                    const appName =
-                        vscode.env.appName === 'Visual Studio Code' ? 'vscode' : vscode.env.appName.toLowerCase();
-                    const redirectUrl = `${API_URL}/oauth/oauth_callback?extId=zenstack.zenstack-v3&vscodeapp=${appName}`;
+                    const appName = vscode.env.uriScheme;
+                    const redirectUrl = `${API_URL}/oauth/oauth_callback?vscodeapp=${appName}&extId=${EXTENSION_ID}`;
                     const signInUrl = vscode.Uri.parse(new URL('/sign-in', AUTH_URL).toString()).with({
-                        query: `redirect_url=${redirectUrl}`,
+                        query: `redirect_url=${encodeURIComponent(redirectUrl)}`,
                     });
 
-                    console.log('ZenStack sign-in URL:', signInUrl.toString());
+                    const finalUrl = signInUrl.toString(true);
+                    console.log('ZenStack sign-in URL:', finalUrl);
+
                     // Store the state and resolve function for later use
                     this.pendingAuth = { resolve, reject, scopes };
 
                     // Open the ZenStack sign-in page in the user's default browser
-                    vscode.env.openExternal(signInUrl).then(
+                    // @ts-ignore - vscode issue: https://github.com/microsoft/vscode/issues/85930
+                    vscode.env.openExternal(finalUrl).then(
                         () => {
                             console.log('Opened ZenStack sign-in page in browser');
                             progress.report({ message: 'Waiting for return from browser...' });

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

@@ -1,4 +1,7 @@
 import { createZModelLanguageServices } from '@zenstackhq/language';
+import type { Model } from '@zenstackhq/language/ast';
+import { getDocument, resolveImport } from '@zenstackhq/language/utils';
+import { URI, type LangiumDocument, type LangiumDocuments } from 'langium';
 import { startLanguageServer } from 'langium/lsp';
 import { NodeFileSystem } from 'langium/node';
 import { createConnection, ProposedFeatures } from 'vscode-languageserver/node.js';
@@ -15,5 +18,76 @@ const { shared } = createZModelLanguageServices(
     true,
 );
 
+// Add custom LSP request handlers
+connection.onRequest('zenstack/getAllImportedZModelURIs', async (params: { textDocument: { uri: string } }) => {
+    try {
+        const uri = URI.parse(params.textDocument.uri);
+        const document = await shared.workspace.LangiumDocuments.getOrCreateDocument(uri);
+
+        // Ensure the document is parsed and built
+        if (!document.parseResult) {
+            await shared.workspace.DocumentBuilder.build([document]);
+        }
+
+        const langiumDocuments = shared.workspace.LangiumDocuments;
+
+        // load all imports
+        const importedURIs = eagerLoadAllImports(document, langiumDocuments);
+
+        const importedDocuments = await Promise.all(
+            importedURIs.map((uri) => langiumDocuments.getOrCreateDocument(uri)),
+        );
+
+        // build the document together with standard library, plugin modules, and imported documents
+        await shared.workspace.DocumentBuilder.build([document, ...importedDocuments], {
+            validation: true,
+        });
+
+        let hasSyntaxErrors = false;
+        for (const docUri of [uri, ...importedURIs]) {
+            const doc = await langiumDocuments.getOrCreateDocument(docUri);
+            if (
+                doc.parseResult.lexerErrors.length > 0 ||
+                doc.parseResult.parserErrors.length > 0 ||
+                doc.diagnostics?.some((e) => e.severity === 1)
+            ) {
+                hasSyntaxErrors = true;
+                break;
+            }
+        }
+
+        return {
+            hasSyntaxErrors,
+            importedURIs,
+        };
+    } catch (error) {
+        console.error('Error getting imported ZModel file:', error);
+        return {
+            hasSyntaxErrors: true,
+            importedURIs: [],
+        };
+    }
+});
+
+function eagerLoadAllImports(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);
+                eagerLoadAllImports(importedDoc, documents, uris);
+            }
+        }
+    }
+
+    return Array.from(uris)
+        .filter((x) => uriString != x)
+        .map((e) => URI.parse(e));
+}
+
 // Start the language server with the shared services
 startLanguageServer(shared);