From e6f7c9adafde151cfaf729f1a320aed7fcf84424 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E8=8D=89=E9=9E=8B=E6=B2=A1=E5=8F=B7?= <308487730@qq.com>
Date: Sat, 1 Apr 2023 18:20:32 +0800
Subject: [PATCH] =?UTF-8?q?feat(plugin-vite):=20hot-restart=20=F0=9F=94=A5?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 packages/plugin/vite/src/Config.ts           |  5 +++
 packages/plugin/vite/src/ViteConfig.ts       |  6 +--
 packages/plugin/vite/src/util/plugins.ts     | 43 +++++++++++++++++++-
 packages/plugin/vite/test/ViteConfig_spec.ts |  2 +-
 4 files changed, 51 insertions(+), 5 deletions(-)

diff --git a/packages/plugin/vite/src/Config.ts b/packages/plugin/vite/src/Config.ts
index 2f1ffd3601..e85b418caa 100644
--- a/packages/plugin/vite/src/Config.ts
+++ b/packages/plugin/vite/src/Config.ts
@@ -9,6 +9,11 @@ export interface VitePluginBuildConfig {
    * Vite config file path.
    */
   config?: string;
+  /**
+   * By default, when any entry in `build` is rebuilt it will restart the Electron App.
+   * If you want to customize this behavior, you can pass a function and control it with the `args.restart` provided by the function.
+   */
+  restart?: false | ((args: { restart: () => void }) => void | Promise<void>);
 }
 
 export interface VitePluginRendererConfig {
diff --git a/packages/plugin/vite/src/ViteConfig.ts b/packages/plugin/vite/src/ViteConfig.ts
index c03449054c..45590156ce 100644
--- a/packages/plugin/vite/src/ViteConfig.ts
+++ b/packages/plugin/vite/src/ViteConfig.ts
@@ -4,7 +4,7 @@ import debug from 'debug';
 import { ConfigEnv, loadConfigFromFile, mergeConfig, UserConfig } from 'vite';
 
 import { VitePluginConfig } from './Config';
-import { externalBuiltins } from './util/plugins';
+import { externalBuiltins, hotRestart } from './util/plugins';
 
 const d = debug('electron-forge:plugin:vite:viteconfig');
 
@@ -68,7 +68,7 @@ export default class ViteConfigGenerator {
     const plugins = [externalBuiltins()];
     const configs = this.pluginConfig.build
       .filter(({ entry, config }) => entry || config)
-      .map<Promise<UserConfig>>(async ({ entry, config }) => {
+      .map<Promise<UserConfig>>(async ({ entry, config, restart }) => {
         const defaultConfig: UserConfig = {
           // Ensure that each build config loads the .env file correctly.
           mode: this.mode,
@@ -89,7 +89,7 @@ export default class ViteConfigGenerator {
           },
           clearScreen: false,
           define,
-          plugins,
+          plugins: [...plugins, hotRestart(restart)],
         };
         if (config) {
           const loadResult = await this.resolveConfig(config);
diff --git a/packages/plugin/vite/src/util/plugins.ts b/packages/plugin/vite/src/util/plugins.ts
index fb5df3db1b..5ad459cb93 100644
--- a/packages/plugin/vite/src/util/plugins.ts
+++ b/packages/plugin/vite/src/util/plugins.ts
@@ -1,6 +1,8 @@
 import { builtinModules } from 'node:module';
 
-import type { Plugin } from 'vite';
+import { VitePluginBuildConfig } from '../Config';
+
+import type { Plugin, ResolvedConfig } from 'vite';
 
 /**
  * `electron` and Node.js built-in modules should always be externalize.
@@ -33,3 +35,42 @@ export function externalBuiltins() {
     },
   };
 }
+
+/**
+ * Hot restart App during development for DX.
+ */
+export function hotRestart(options: VitePluginBuildConfig['restart']) {
+  let config: ResolvedConfig;
+  const restart = () => {
+    // https://github.com/electron/forge/blob/v6.0.5/packages/api/core/src/api/start.ts#L204-L211
+    process.stdin.emit('data', 'rs');
+  };
+  // Avoid first start, it's stated by forge.
+  let isFirstStart = false;
+
+  return <Plugin>{
+    name: '@electron-forge/plugin-vite:hot-restart',
+    configResolved(_config) {
+      config = _config;
+    },
+    closeBundle() {
+      if (config.mode === 'production') {
+        // https://github.com/electron/forge/blob/98b621dcace753c1bd33aeb301c64d03335abfdc/packages/plugin/vite/src/ViteConfig.ts#L36-L41
+        return;
+      }
+      if (options === false) {
+        return;
+      }
+      if (!isFirstStart) {
+        isFirstStart = true;
+        return;
+      }
+      if (typeof options === 'function') {
+        // Leave it to the user to decide whether to restart.
+        options({ restart });
+      } else {
+        restart();
+      }
+    },
+  };
+}
diff --git a/packages/plugin/vite/test/ViteConfig_spec.ts b/packages/plugin/vite/test/ViteConfig_spec.ts
index 3699dab4f4..4209105730 100644
--- a/packages/plugin/vite/test/ViteConfig_spec.ts
+++ b/packages/plugin/vite/test/ViteConfig_spec.ts
@@ -42,7 +42,7 @@ describe('ViteConfigGenerator', () => {
       clearScreen: false,
       define: {},
       // shims
-      plugins: [buildConfig.plugins?.[0]],
+      plugins: buildConfig.plugins,
     } as UserConfig);
   });