From a1ddaa37e2b664994f8fddc33a3730c249c79984 Mon Sep 17 00:00:00 2001
From: JongHak Seo <unqocn@gmail.com>
Date: Sun, 8 Dec 2024 19:45:12 +0900
Subject: [PATCH] [Feat] Support content.css HMR (#799)

Co-authored-by: Jonghakseo <jonghak.seo@creatrip.com>
---
 chrome-extension/manifest.js                  |  1 +
 chrome-extension/public/refresh.js            | 72 +++++++++++++++++++
 chrome-extension/vite.config.mts              |  2 +-
 .../hmr/lib/plugins/watch-rebuild-plugin.ts   | 10 +--
 packages/hmr/lib/types.ts                     |  1 +
 5 files changed, 80 insertions(+), 6 deletions(-)
 create mode 100644 chrome-extension/public/refresh.js

diff --git a/chrome-extension/manifest.js b/chrome-extension/manifest.js
index 9e23a5244..e5113ea99 100755
--- a/chrome-extension/manifest.js
+++ b/chrome-extension/manifest.js
@@ -55,6 +55,7 @@ const manifest = deepmerge(
       },
       {
         matches: ['http://*/*', 'https://*/*', '<all_urls>'],
+        js: ['refresh.js'], // for public's HMR(refresh) support
         css: ['content.css'], // public folder
       },
     ],
diff --git a/chrome-extension/public/refresh.js b/chrome-extension/public/refresh.js
new file mode 100644
index 000000000..5d3dad3dd
--- /dev/null
+++ b/chrome-extension/public/refresh.js
@@ -0,0 +1,72 @@
+/* eslint-disable */
+(function () {
+  'use strict';
+  // This is the custom ID for HMR (chrome-extension/vite.config.mts)
+  const __HMR_ID = 'chrome-extension-hmr';
+
+  const LOCAL_RELOAD_SOCKET_PORT = 8081;
+  const LOCAL_RELOAD_SOCKET_URL = `ws://localhost:${LOCAL_RELOAD_SOCKET_PORT}`;
+
+  const DO_UPDATE = 'do_update';
+  const DONE_UPDATE = 'done_update';
+
+  class MessageInterpreter {
+    // eslint-disable-next-line @typescript-eslint/no-empty-function
+    constructor() {}
+
+    static send(message) {
+      return JSON.stringify(message);
+    }
+
+    static receive(serializedMessage) {
+      return JSON.parse(serializedMessage);
+    }
+  }
+
+  function initClient({ id, onUpdate }) {
+    const ws = new WebSocket(LOCAL_RELOAD_SOCKET_URL);
+
+    ws.onopen = () => {
+      ws.addEventListener('message', event => {
+        const message = MessageInterpreter.receive(String(event.data));
+
+        if (message.type === DO_UPDATE && message.id === id) {
+          onUpdate();
+          ws.send(MessageInterpreter.send({ type: DONE_UPDATE }));
+          return;
+        }
+      });
+    };
+  }
+
+  function addRefresh() {
+    let pendingReload = false;
+
+    initClient({
+      id: __HMR_ID,
+      onUpdate: () => {
+        // disable reload when tab is hidden
+        if (document.hidden) {
+          pendingReload = true;
+          return;
+        }
+        reload();
+      },
+    });
+
+    // reload
+    function reload() {
+      pendingReload = false;
+      window.location.reload();
+    }
+
+    // reload when tab is visible
+    function reloadWhenTabIsVisible() {
+      !document.hidden && pendingReload && reload();
+    }
+
+    document.addEventListener('visibilitychange', reloadWhenTabIsVisible);
+  }
+
+  addRefresh();
+})();
diff --git a/chrome-extension/vite.config.mts b/chrome-extension/vite.config.mts
index 9664b4f76..295dc339f 100644
--- a/chrome-extension/vite.config.mts
+++ b/chrome-extension/vite.config.mts
@@ -23,7 +23,7 @@ export default defineConfig({
     }) as PluginOption,
     watchPublicPlugin(),
     makeManifestPlugin({ outDir }),
-    isDev && watchRebuildPlugin({ reload: true }),
+    isDev && watchRebuildPlugin({ reload: true, id: 'chrome-extension-hmr' }),
   ],
   publicDir: resolve(rootDir, 'public'),
   build: {
diff --git a/packages/hmr/lib/plugins/watch-rebuild-plugin.ts b/packages/hmr/lib/plugins/watch-rebuild-plugin.ts
index e37da7a74..3b3b88a11 100644
--- a/packages/hmr/lib/plugins/watch-rebuild-plugin.ts
+++ b/packages/hmr/lib/plugins/watch-rebuild-plugin.ts
@@ -12,14 +12,14 @@ const refreshCode = fs.readFileSync(path.resolve(injectionsPath, 'refresh.js'),
 const reloadCode = fs.readFileSync(path.resolve(injectionsPath, 'reload.js'), 'utf-8');
 
 export function watchRebuildPlugin(config: PluginConfig): PluginOption {
+  const { refresh, reload, id: _id, onStart } = config;
+  const hmrCode = (refresh ? refreshCode : '') + (reload ? reloadCode : '');
+
   let ws: WebSocket | null = null;
 
-  const id = Math.random().toString(36);
+  const id = _id ?? Math.random().toString(36);
   let reconnectTries = 0;
 
-  const { refresh, reload } = config;
-  const hmrCode = (refresh ? refreshCode : '') + (reload ? reloadCode : '');
-
   function initializeWebSocket() {
     ws = new WebSocket(LOCAL_RELOAD_SOCKET_URL);
 
@@ -46,7 +46,7 @@ export function watchRebuildPlugin(config: PluginConfig): PluginOption {
   return {
     name: 'watch-rebuild',
     writeBundle() {
-      config.onStart?.();
+      onStart?.();
       if (!ws) {
         initializeWebSocket();
         return;
diff --git a/packages/hmr/lib/types.ts b/packages/hmr/lib/types.ts
index b93862ee7..d9c8e0556 100644
--- a/packages/hmr/lib/types.ts
+++ b/packages/hmr/lib/types.ts
@@ -16,4 +16,5 @@ export type PluginConfig = {
   onStart?: () => void;
   reload?: boolean;
   refresh?: boolean;
+  id?: string;
 };