fix: improve Vite plugin hot update handling with file invalidation

Author: HichemTab-tech

src/ts-transformer.ts

@@ -12,7 +12,16 @@ export const refreshProject = () => {
     });
 }
 
-refreshProject();
+export const invalidateOneFile = (file: string) => {
+
+    const sourceFileToRefresh = project.getSourceFile(file);
+    if (sourceFileToRefresh) {
+        sourceFileToRefresh.refreshFromFileSystemSync();
+        return true;
+    } else {
+        return false;
+    }
+}
 
 export function transform(code: string, filePath: string): string {
 
@@ -89,4 +98,6 @@ export function transform(code: string, filePath: string): string {
     }
 
     return sourceFile.getFullText();
-}
+}
+
+refreshProject();

src/vite-plugin.ts

@@ -1,5 +1,5 @@
 import { PluginOption } from "vite";
-import {fileToTypes, refreshProject, transform, typeToFile} from "./ts-transformer.js";
+import {fileToTypes, invalidateOneFile, transform, typeToFile} from "./ts-transformer.js";
 
 // noinspection JSUnusedGlobalSymbols
 /**
@@ -28,40 +28,61 @@ export default function TsRuntimePickerVitePlugin(): PluginOption {
             return null;
         },
 
-        handleHotUpdate: ({server, file}) => {
-            const affectedFiles = [];
+        handleHotUpdate: ({server, file, timestamp}) => {
+            const isKnownTypeDefinitionFile = [...typeToFile.values()].includes(file);
 
-            // Loop over all files that used types
+            if (!isKnownTypeDefinitionFile) {
+                // This file change doesn't affect any of our transformations, so we can ignore it.
+                return;
+            }
+
+            console.log(`[ts-runtime-picker] Change detected in a tracked type definition: ${file}`);
+
+
+            // Surgically update the ts-morph project
+            if (!invalidateOneFile(file)) return;
+
+            const dependentModulesToUpdate = [];
+
+            // Find all files (`userFile`) that import a type from the file that just changed (`file`).
             for (const [userFile, usedTypes] of fileToTypes.entries()) {
                 for (const typeName of usedTypes) {
                     if (typeToFile.get(typeName) === file) {
-                        affectedFiles.push(userFile);
-                        break;
+                        // This `userFile` depends on the changed type. We need to update it.
+                        const moduleNode = server.moduleGraph.getModuleById(userFile);
+                        if (moduleNode) {
+                            dependentModulesToUpdate.push(moduleNode);
+                            // Once we know the file is affected, we don't need to check its other types.
+                            break;
+                        }
                     }
                 }
             }
 
-            if (affectedFiles.length) {
-                refreshProject();
+            if (dependentModulesToUpdate.length > 0) {
+                console.log(`[ts-runtime-picker] Found ${dependentModulesToUpdate.length} module(s) to update.`);
 
-                for (const file of affectedFiles) {
-                    const mod = server.moduleGraph.getModuleById(file);
-                    if (mod) {
-                        server.moduleGraph.invalidateModule(mod);
-                        server.ws.send({
-                            type: 'update',
-                            updates: [
-                                {
-                                    type: 'js-update',
-                                    path: mod.url, // relative to root
-                                    acceptedPath: mod.url,
-                                    timestamp: Date.now()
-                                }
-                            ]
-                        });
-                    }
+                //Invalidate modules and notify the client
+                for (const mod of dependentModulesToUpdate) {
+                    server.moduleGraph.invalidateModule(mod);
                 }
+
+                // This sends the HMR update to the client (browser), telling it which
+                // modules have changed so it can re-request them.
+                server.ws.send({
+                    type: 'update',
+                    updates: dependentModulesToUpdate.map((mod) => ({
+                        type: 'js-update',
+                        path: mod.url,
+                        acceptedPath: mod.url,
+                        timestamp: timestamp
+                    }))
+                });
             }
+
+            // We return the list of modules we've handled so Vite's core HMR logic
+            // doesn't try to process them again.
+            return dependentModulesToUpdate;
         }
     } as PluginOption;
 }