fix: HMR plugin causes full page reload

src/components/app.tsx

@@ -28,6 +28,10 @@
   showExample: boolean;
 };
 
+// Swapped out during dev when a linked vega-lite checkout is edited,
+// so specs recompile with the fresh code without a page reload
+let activeVegaLite: typeof vegaLite = vegaLite;
+
 const App: React.FC<Props> = (props) => {
   const appContext = useAppContext();
   const {state, setState} = appContext;
@@ -206,7 +210,7 @@
                 `The specification expects Vega-Lite ${parsed.version} but the editor uses v${vega.version}.`,
               );
             }
           } catch (e) {
             throw new Error('Could not parse $schema url.');
           }
         }
@@ -214,7 +218,7 @@
         validateVegaLite(vegaLiteSpec, currLogger);
 
         const compileResult =
-          editorString !== '{}' ? vegaLite.compile(vegaLiteSpec, options) : {spec: {}, normalized: {}};
+          editorString !== '{}' ? activeVegaLite.compile(vegaLiteSpec, options) : {spec: {}, normalized: {}};
         const normalizedSpec = compileResult.normalized;
 
         setState((s) => ({
@@ -257,7 +261,7 @@
             if (!satisfies(vega.version, `^${parsed.version.slice(1)}`)) {
               currLogger.warn(`The specification expects Vega ${parsed.version} but the editor uses v${vega.version}.`);
             }
           } catch (e) {
             throw new Error('Could not parse $schema url.');
           }
         }
@@ -290,6 +294,23 @@
     [setState],
   );
 
+  // For dev, when a linked vega-lite checkout is edited, swap in the fresh
+  // module and re-parse the current spec so the change shows up immediately.
+  useEffect(() => {
+    if (!import.meta.hot) {
+      return;
+    }
+    const onVegaPackageHmr = (event: Event) => {
+      const {packageName, module} = (event as CustomEvent<{packageName: string; module: typeof vegaLite}>).detail;
+      if (packageName === 'vega-lite' && module) {
+        activeVegaLite = module;
+        setState((s) => ({...s, parse: true}));
+      }
+    };
+    window.addEventListener('vega-package-hmr', onVegaPackageHmr);
+    return () => window.removeEventListener('vega-package-hmr', onVegaPackageHmr);
+  }, [setState]);
+
   useEffect(() => {
     const handleMessage = (evt: MessageEvent) => {
       const data = evt.data as MessageData;

src/index.tsx

@@ -20,6 +20,21 @@ if (typeof window !== 'undefined') {
   console.log('Vega versions set on window.VEGA_DEBUG');
 }
 
+if (import.meta.hot) {
+  // Pointing to at the freshest linked vega source
+  window.addEventListener('vega-package-hmr', (event: Event) => {
+    const {packageName, module} = (event as CustomEvent).detail;
+    const w: typeof globalThis = window;
+    if (packageName === 'vega-lite' && module) {
+      w.VEGA_DEBUG.vegaLite = module;
+      w.VEGA_DEBUG.VEGA_LITE_VERSION = module.version;
+    } else if (packageName === 'vega' && module) {
+      w.VEGA_DEBUG.vega = module;
+      w.VEGA_DEBUG.VEGA_VERSION = module.version;
+    }
+  });
+}
+
 try {
   console.log('Setting up Monaco editor...');
   setupMonaco();

tsconfig.json

@@ -17,6 +17,7 @@
     ],
     "skipLibCheck": true,
     "types": [
+      "vite/client",
       "react",
       "react-dom",
       "node",

vite.config.ts

@@ -2,157 +2,83 @@
 import {defineConfig} from 'vitest/config';
 import react from '@vitejs/plugin-react';
 import {resolve} from 'path';
-import {watch as fsWatch, readdirSync, lstatSync, existsSync} from 'fs';
+import {readdirSync, lstatSync, existsSync, realpathSync} from 'fs';
 import {execSync} from 'child_process';
 import {playwright} from '@vitest/browser-playwright';
 
 const commitHash = execSync('git rev-parse HEAD', {encoding: 'utf8'}).trim();
 
-const vegaUtils = {
-  isVegaPackage: (packageName: string): boolean => packageName === 'vega' || packageName.startsWith('vega-'),
+const nodeModulesPath = resolve(__dirname, 'node_modules');
 
-  isVegaModule: (url: string): boolean => url?.includes('vega-') || url?.includes('vega'),
+const isVegaPackage = (packageName: string): boolean => packageName === 'vega' || packageName.startsWith('vega-');
 
-  getNodeModulesPath: () => resolve(process.cwd(), 'node_modules'),
+const getVegaPackageNames = (): string[] => readdirSync(nodeModulesPath).filter(isVegaPackage);
 
-  getVegaPackageNames: (nodeModulesPath: string): string[] => {
-    const items = readdirSync(nodeModulesPath);
-    return items.filter((item) => vegaUtils.isVegaPackage(item));
-  },
+// Vega packages symlinked into node modules mapped
+// to their specific source entry point so edits are picked up without rebuilding.
+const getLinkedVegaPackages = (): Map<string, {entry: string; srcDir: string}> => {
+  const linked = new Map<string, {entry: string; srcDir: string}>();
+  for (const packageName of getVegaPackageNames()) {
+    const packagePath = resolve(nodeModulesPath, packageName);
+    if (!lstatSync(packagePath).isSymbolicLink()) continue;
 
-  findEntryPoint: (realPath: string): string | null => {
-    const entryPoints = [
+    const realPath = realpathSync(packagePath);
+    const entry = [
       resolve(realPath, 'index.js'),
       resolve(realPath, 'src', 'index.js'),
       resolve(realPath, 'src', 'index.ts'),
-    ];
-    return entryPoints.find(existsSync) || null;
-  },
-
-  resolvePackagePath: (packageName: string, nodeModulesPath: string) => {
-    const packagePath = resolve(nodeModulesPath, packageName);
-    if (!existsSync(packagePath)) return null;
-
-    const stats = lstatSync(packagePath);
-    if (stats.isSymbolicLink()) {
-      const realPath = resolve(packagePath);
-      return {
-        realPath,
-        entryPoint: vegaUtils.findEntryPoint(realPath),
-        srcPath: resolve(realPath, 'src'),
-      };
-    }
-    return null;
-  },
+    ].find(existsSync);
+    const srcDir = resolve(realPath, 'src');
 
-  handleVegaFileChange: (server: any) => {
-    const moduleGraph = server.moduleGraph;
-    const modules = Array.from(moduleGraph.urlToModuleMap.values());
-    const modulesToUpdate: any[] = [];
-
-    modules.forEach((module: any) => {
-      const isVegaModule = vegaUtils.isVegaModule(module.url);
-      const importsVega =
-        module.importedModules &&
-        Array.from(module.importedModules).some((importedModule: any) =>
-          vegaUtils.isVegaModule(importedModule.id || importedModule.url),
-        );
-
-      if (isVegaModule || importsVega) {
-        modulesToUpdate.push(module);
-        moduleGraph.invalidateModule(module);
-      }
-    });
-
-    if (modulesToUpdate.length > 0) {
-      server.ws.send({
-        type: 'update',
-        updates: modulesToUpdate.map((module) => ({
-          type: 'js-update' as const,
-          path: module.url,
-          acceptedPath: module.url,
-          timestamp: Date.now(),
-        })),
-      });
-
-      server.ws.send({
-        type: 'custom',
-        event: 'vega-package-updating',
-        data: {timestamp: Date.now()},
-      });
+    if (entry && existsSync(srcDir)) {
+      linked.set(packageName, {entry, srcDir});
     }
-  },
+  }
+  return linked;
 };
 
 function createVegaHMRPlugin() {
+  const linkedPackages = getLinkedVegaPackages();
   return {
     name: 'vega-packages-hmr',
     enforce: 'pre' as const,
 
-    configureServer(server: any) {
-      const nodeModulesPath = resolve(__dirname, 'node_modules');
-      const vegaPackageNames = vegaUtils.getVegaPackageNames(nodeModulesPath);
-
-      const {vegaPackagePaths, vegaPackageAliases} = vegaPackageNames.reduce(
-        (acc, packageName) => {
-          const resolved = vegaUtils.resolvePackagePath(packageName, nodeModulesPath);
-          if (resolved) {
-            if (existsSync(resolved.srcPath)) {
-              acc.vegaPackagePaths.push(resolved.srcPath);
-            }
-            if (resolved.entryPoint) {
-              acc.vegaPackageAliases[packageName] = resolved.entryPoint;
-            }
-          }
-          return acc;
-        },
-        {vegaPackagePaths: [] as string[], vegaPackageAliases: {} as Record<string, string>},
-      );
-
-      server.vegaPackageAliases = vegaPackageAliases;
-
-      const watchers = vegaPackagePaths.map((srcPath) =>
-        fsWatch(srcPath, {recursive: true}, (_, filename) => {
-          if (filename?.match(/\.(ts|js)$/)) {
-            vegaUtils.handleVegaFileChange(server);
-          }
-        }),
-      );
-
-      server.httpServer?.on('close', () => {
-        watchers.forEach((watcher) => watcher.close());
-      });
+    resolveId(id: string) {
+      return linkedPackages.get(id)?.entry ?? null;
     },
 
-    resolveId(id: string) {
-      if (vegaUtils.isVegaPackage(id)) {
-        const nodeModulesPath = resolve(__dirname, 'node_modules');
-        const resolved = vegaUtils.resolvePackagePath(id, nodeModulesPath);
-        return resolved?.entryPoint || null;
+    configureServer(server: {watcher: {add: (path: string) => void}}) {
+      for (const {srcDir} of linkedPackages.values()) {
+        server.watcher.add(srcDir);
       }
-      return null;
     },
 
-    async transform(code: string, id: string) {
-      if (id.includes('/vega') && id.includes('/src/')) {
-        return {
-          code: `
+    // Make each linked package's entry module accept
+    // updates for the whole package tree, so propagation stops there
+    transform(code: string, id: string) {
+      for (const [packageName, {entry}] of linkedPackages) {
+        if (id === entry) {
+          return {
+            code: `${code}
 if (import.meta.hot) {
-    import.meta.hot.accept();
+  import.meta.hot.accept((newModule) => {
+    window.dispatchEvent(
+      new CustomEvent('vega-package-hmr', {detail: {packageName: ${JSON.stringify(packageName)}, module: newModule}}),
+    );
+  });
 }
-${code}`,
-          map: null,
-        };
+`,
+            map: null,
+          };
+        }
       }
+      return null;
     },
   };
 }
 
 export default defineConfig({
-  plugins: [
-    react(),
-    createVegaHMRPlugin(),
-  ],
+  plugins: [react(), createVegaHMRPlugin()],
 
   define: {
     'process.env.VITE_COMMIT_HASH': JSON.stringify(commitHash),
@@ -177,18 +103,14 @@ export default defineConfig({
   server: {
     port: 1234,
     open: true,
-    watch: {
-      ignored: ['!**/node_modules/vega-lite/**'],
-      followSymlinks: true,
-    },
     fs: {
       allow: ['..', resolve(__dirname, '../..')],
     },
   },
 
   optimizeDeps: {
     include: [],
-    exclude: vegaUtils.getVegaPackageNames(vegaUtils.getNodeModulesPath()),
+    exclude: getVegaPackageNames(),
   },
 
   publicDir: 'public',