VS Code extension (#197)

* Fix running the extension

I'm not sure how my setup was different but I was unable to get the extension to run locally without adding a binary. This mirrors what Svelte does so I'm assuming it's the way it's supposed to be loaded.

* Resolve TypeScript suggestions to the correct file

This fixes a couple of bugs related to suggestions. 1 was this does the whole `.ts` extension fakeout thing so that the TypeScript plugin thinks that Astro files are TypeScript. Secondly this fixes the caching of the Document, so that suggestions account for the current document text.

7 files changed, 182 insertions, 20 deletions

diff --git a/scripts/cmd/build.js b/scripts/cmd/build.js
index 81761030e..250bb1efd 100644
--- a/scripts/cmd/build.js
+++ b/scripts/cmd/build.js
@@ -8,7 +8,7 @@ import glob from 'tiny-glob';
 /** @type {import('esbuild').BuildOptions} */
 const defaultConfig = {
   bundle: true,
-  minify: true,
+  minify: false,
   format: 'esm',
   platform: 'node',
   target: 'node14',
diff --git a/tools/astro-languageserver/src/index.ts b/tools/astro-languageserver/src/index.ts
index c834beaf9..41f04d11a 100644
--- a/tools/astro-languageserver/src/index.ts
+++ b/tools/astro-languageserver/src/index.ts
@@ -29,7 +29,7 @@ export function startServer() {
         textDocumentSync: TextDocumentSyncKind.Incremental,
         foldingRangeProvider: true,
         completionProvider: {
-          resolveProvider: false,
+          resolveProvider: true,
           triggerCharacters: [
             '.',
             '"',
@@ -70,7 +70,9 @@ export function startServer() {
 
   connection.onDidCloseTextDocument((evt) => docManager.closeDocument(evt.textDocument.uri));
 
-  connection.onDidChangeTextDocument((evt) => docManager.updateDocument(evt.textDocument.uri, evt.contentChanges));
+  connection.onDidChangeTextDocument((evt) => {
+    docManager.updateDocument(evt.textDocument.uri, evt.contentChanges)
+  });
 
   connection.onDidChangeWatchedFiles((evt) => {
     const params = evt.changes
diff --git a/tools/astro-languageserver/src/plugins/typescript/SnapshotManager.ts b/tools/astro-languageserver/src/plugins/typescript/SnapshotManager.ts
index 47d44838d..4f9e865a1 100644
--- a/tools/astro-languageserver/src/plugins/typescript/SnapshotManager.ts
+++ b/tools/astro-languageserver/src/plugins/typescript/SnapshotManager.ts
@@ -37,7 +37,7 @@ export class SnapshotManager {
       }
       previousSnapshot.update(changes);
     } else {
-      const newSnapshot = createDocumentSnapshot(fileName);
+      const newSnapshot = createDocumentSnapshot(fileName, null);
 
       if (previousSnapshot) {
         newSnapshot.version = previousSnapshot.version + 1;
@@ -120,8 +120,8 @@ export interface DocumentSnapshot extends ts.IScriptSnapshot {
   getFullText(): string;
 }
 
-export const createDocumentSnapshot = (filePath: string, createDocument?: (_filePath: string, text: string) => Document): DocumentSnapshot => {
-  const text = ts.sys.readFile(filePath) ?? '';
+export const createDocumentSnapshot = (filePath: string, currentText: string | null, createDocument?: (_filePath: string, text: string) => Document): DocumentSnapshot => {
+  const text = currentText || (ts.sys.readFile(filePath) ?? '');
 
   if (isAstroFilePath(filePath)) {
     if (!createDocument) throw new Error('Astro documents require the "createDocument" utility to be provided');
diff --git a/tools/astro-languageserver/src/plugins/typescript/features/CompletionsProvider.ts b/tools/astro-languageserver/src/plugins/typescript/features/CompletionsProvider.ts
index 348f3e4ae..e56902e6e 100644
--- a/tools/astro-languageserver/src/plugins/typescript/features/CompletionsProvider.ts
+++ b/tools/astro-languageserver/src/plugins/typescript/features/CompletionsProvider.ts
@@ -25,7 +25,13 @@ export class CompletionsProviderImpl implements CompletionsProvider<CompletionEn
     const { tsDoc, lang } = await this.lang.getTypeScriptDoc(document);
     const fragment = await tsDoc.getFragment();
 
-    const { entries } = lang.getCompletionsAtPosition(fragment.filePath, document.offsetAt(position), {}) ?? { entries: [] };
+    const offset = document.offsetAt(position);
+    const entries = lang.getCompletionsAtPosition(fragment.filePath, offset, {
+      importModuleSpecifierPreference: 'relative',
+      importModuleSpecifierEnding: 'auto',
+      quotePreference: 'single'
+    })?.entries || [];
+
     const completionItems = entries
       .map((entry: ts.CompletionEntry) => this.toCompletionItem(fragment, entry, document.uri, position, new Set()))
       .filter((i) => i) as CompletionItem[];
@@ -37,12 +43,16 @@ export class CompletionsProviderImpl implements CompletionsProvider<CompletionEn
     const { data: comp } = completionItem;
     const { tsDoc, lang } = await this.lang.getTypeScriptDoc(document);
 
-    const filePath = tsDoc.filePath;
+    let filePath = tsDoc.filePath;
 
     if (!comp || !filePath) {
       return completionItem;
     }
 
+    if(filePath.endsWith('.astro')) {
+      filePath = filePath + '.ts';
+    }
+
     const fragment = await tsDoc.getFragment();
     const detail = lang.getCompletionEntryDetails(filePath, fragment.offsetAt(comp.position), comp.name, {}, comp.source, {});
 
diff --git a/tools/astro-languageserver/src/plugins/typescript/languageService.ts b/tools/astro-languageserver/src/plugins/typescript/languageService.ts
index 098c335e7..ddd8cb3cd 100644
--- a/tools/astro-languageserver/src/plugins/typescript/languageService.ts
+++ b/tools/astro-languageserver/src/plugins/typescript/languageService.ts
@@ -2,10 +2,10 @@
 
 import * as ts from 'typescript';
 import { basename } from 'path';
-import { ensureRealAstroFilePath, findTsConfigPath, isAstroFilePath, toVirtualAstroFilePath } from './utils';
+import { ensureRealAstroFilePath, findTsConfigPath } from './utils';
 import { Document } from '../../core/documents';
 import { createDocumentSnapshot, SnapshotManager, DocumentSnapshot } from './SnapshotManager';
-import { createAstroSys } from './astro-sys';
+import { createAstroModuleLoader } from './module-loader';
 
 const services = new Map<string, Promise<LanguageServiceContainer>>();
 
@@ -72,18 +72,19 @@ async function createLanguageService(tsconfigPath: string, workspaceRoot: string
 
   let projectVersion = 0;
   const snapshotManager = new SnapshotManager(project.fileNames, { exclude: ['node_modules', 'dist'], include: ['astro'] }, workspaceRoot || process.cwd());
-  const astroSys = createAstroSys(updateDocument);
+  
+  const astroModuleLoader = createAstroModuleLoader(getScriptSnapshot, {});
 
   const host: ts.LanguageServiceHost = {
     getNewLine: () => ts.sys.newLine,
     useCaseSensitiveFileNames: () => ts.sys.useCaseSensitiveFileNames,
-    readFile: astroSys.readFile,
-    writeFile: astroSys.writeFile,
-    fileExists: astroSys.fileExists,
-    directoryExists: astroSys.directoryExists,
-    getDirectories: astroSys.getDirectories,
-    readDirectory: astroSys.readDirectory,
-    realpath: astroSys.realpath,
+    readFile: astroModuleLoader.readFile,
+    writeFile: astroModuleLoader.writeFile,
+    fileExists: astroModuleLoader.fileExists,
+    directoryExists: astroModuleLoader.directoryExists,
+    getDirectories: astroModuleLoader.getDirectories,
+    readDirectory: astroModuleLoader.readDirectory,
+    realpath: astroModuleLoader.realpath,
 
     getCompilationSettings: () => project.options,
     getCurrentDirectory: () => workspaceRoot,
@@ -127,7 +128,8 @@ async function createLanguageService(tsconfigPath: string, workspaceRoot: string
       return previousSnapshot;
     }
 
-    const snapshot = createDocumentSnapshot(filePath, docContext.createDocument);
+    const currentText = document ? document.getText() : null;
+    const snapshot = createDocumentSnapshot(filePath, currentText, docContext.createDocument);
     snapshotManager.set(filePath, snapshot);
     return snapshot;
   }
@@ -140,7 +142,7 @@ async function createLanguageService(tsconfigPath: string, workspaceRoot: string
       return doc;
     }
 
-    doc = createDocumentSnapshot(fileName, docContext.createDocument);
+    doc = createDocumentSnapshot(fileName, null, docContext.createDocument);
     snapshotManager.set(fileName, doc);
     return doc;
   }
diff --git a/tools/astro-languageserver/src/plugins/typescript/module-loader.ts b/tools/astro-languageserver/src/plugins/typescript/module-loader.ts
new file mode 100644
index 000000000..6bed70ac3
--- /dev/null
+++ b/tools/astro-languageserver/src/plugins/typescript/module-loader.ts
@@ -0,0 +1,132 @@
+import ts from 'typescript';
+import type { DocumentSnapshot } from './SnapshotManager';
+import {
+    isVirtualAstroFilePath,
+    ensureRealAstroFilePath,
+    getExtensionFromScriptKind
+} from './utils';
+import { createAstroSys } from './astro-sys';
+
+/**
+ * Caches resolved modules.
+ */
+class ModuleResolutionCache {
+    private cache = new Map<string, ts.ResolvedModule>();
+
+    /**
+     * Tries to get a cached module.
+     */
+    get(moduleName: string, containingFile: string): ts.ResolvedModule | undefined {
+        return this.cache.get(this.getKey(moduleName, containingFile));
+    }
+
+    /**
+     * Caches resolved module, if it is not undefined.
+     */
+    set(moduleName: string, containingFile: string, resolvedModule: ts.ResolvedModule | undefined) {
+        if (!resolvedModule) {
+            return;
+        }
+        this.cache.set(this.getKey(moduleName, containingFile), resolvedModule);
+    }
+
+    /**
+     * Deletes module from cache. Call this if a file was deleted.
+     * @param resolvedModuleName full path of the module
+     */
+    delete(resolvedModuleName: string): void {
+        this.cache.forEach((val, key) => {
+            if (val.resolvedFileName === resolvedModuleName) {
+                this.cache.delete(key);
+            }
+        });
+    }
+
+    private getKey(moduleName: string, containingFile: string) {
+        return containingFile + ':::' + ensureRealAstroFilePath(moduleName);
+    }
+}
+
+/**
+ * Creates a module loader specifically for `.astro` files.
+ *
+ * The typescript language service tries to look up other files that are referenced in the currently open astro file.
+ * For `.ts`/`.js` files this works, for `.astro` files it does not by default.
+ * Reason: The typescript language service does not know about the `.astro` file ending,
+ * so it assumes it's a normal typescript file and searches for files like `../Component.astro.ts`, which is wrong.
+ * In order to fix this, we need to wrap typescript's module resolution and reroute all `.astro.ts` file lookups to .astro.
+ *
+ * @param getSnapshot A function which returns a (in case of astro file fully preprocessed) typescript/javascript snapshot
+ * @param compilerOptions The typescript compiler options
+ */
+export function createAstroModuleLoader(
+    getSnapshot: (fileName: string) => DocumentSnapshot,
+    compilerOptions: ts.CompilerOptions
+) {
+    const astroSys = createAstroSys(getSnapshot);
+    const moduleCache = new ModuleResolutionCache();
+
+    return {
+        fileExists: astroSys.fileExists,
+        readFile: astroSys.readFile,
+        writeFile: astroSys.writeFile,
+        readDirectory: astroSys.readDirectory,
+        directoryExists: astroSys.directoryExists,
+        getDirectories: astroSys.getDirectories,
+        realpath: astroSys.realpath,
+        deleteFromModuleCache: (path: string) => moduleCache.delete(path),
+        resolveModuleNames
+    };
+
+    function resolveModuleNames(
+        moduleNames: string[],
+        containingFile: string
+    ): Array<ts.ResolvedModule | undefined> {
+        return moduleNames.map((moduleName) => {
+            const cachedModule = moduleCache.get(moduleName, containingFile);
+            if (cachedModule) {
+                return cachedModule;
+            }
+
+            const resolvedModule = resolveModuleName(moduleName, containingFile);
+            moduleCache.set(moduleName, containingFile, resolvedModule);
+            return resolvedModule;
+        });
+    }
+
+    function resolveModuleName(
+        name: string,
+        containingFile: string
+    ): ts.ResolvedModule | undefined {
+        // Delegate to the TS resolver first.
+        // If that does not bring up anything, try the Astro Module loader
+        // which is able to deal with .astro files.
+        const tsResolvedModule = ts.resolveModuleName(name, containingFile, compilerOptions, ts.sys)
+            .resolvedModule;
+        if (tsResolvedModule && !isVirtualAstroFilePath(tsResolvedModule.resolvedFileName)) {
+            return tsResolvedModule;
+        }
+
+        const astroResolvedModule = ts.resolveModuleName(
+            name,
+            containingFile,
+            compilerOptions,
+            astroSys
+        ).resolvedModule;
+        if (
+            !astroResolvedModule ||
+            !isVirtualAstroFilePath(astroResolvedModule.resolvedFileName)
+        ) {
+            return astroResolvedModule;
+        }
+
+        const resolvedFileName = ensureRealAstroFilePath(astroResolvedModule.resolvedFileName);
+        const snapshot = getSnapshot(resolvedFileName);
+
+        const resolvedastroModule: ts.ResolvedModuleFull = {
+            extension: getExtensionFromScriptKind(snapshot && snapshot.scriptKind),
+            resolvedFileName
+        };
+        return resolvedastroModule;
+    }
+}
diff --git a/tools/astro-languageserver/src/plugins/typescript/utils.ts b/tools/astro-languageserver/src/plugins/typescript/utils.ts
index 1f42e7d0a..3c43e56d5 100644
--- a/tools/astro-languageserver/src/plugins/typescript/utils.ts
+++ b/tools/astro-languageserver/src/plugins/typescript/utils.ts
@@ -111,6 +111,22 @@ export function getScriptKindFromFileName(fileName: string): ts.ScriptKind {
   }
 }
 
+export function getExtensionFromScriptKind(kind: ts.ScriptKind | undefined): ts.Extension {
+  switch (kind) {
+      case ts.ScriptKind.JSX:
+          return ts.Extension.Jsx;
+      case ts.ScriptKind.TS:
+          return ts.Extension.Ts;
+      case ts.ScriptKind.TSX:
+          return ts.Extension.Tsx;
+      case ts.ScriptKind.JSON:
+          return ts.Extension.Json;
+      case ts.ScriptKind.JS:
+      default:
+          return ts.Extension.Js;
+  }
+}
+
 export function isAstroFilePath(filePath: string) {
   return filePath.endsWith('.astro');
 }