fix: Use HMR API for dev server communication

Author: aklinker1

packages/wxt/src/sandbox/dev-server-websocket.ts

@@ -1,32 +1,28 @@
 import definition from 'virtual:user-background-entrypoint';
 import { initPlugins } from 'virtual:wxt-plugins';
-import { getDevServerWebSocket } from '../sandbox/dev-server-websocket';
 import { logger } from '../sandbox/utils/logger';
 import { browser } from 'wxt/browser';
 import { keepServiceWorkerAlive } from './utils/keep-service-worker-alive';
 import { reloadContentScript } from './utils/reload-content-scripts';
 
-if (import.meta.env.COMMAND === 'serve') {
-  try {
-    const ws = getDevServerWebSocket();
-    ws.addWxtEventListener('wxt:reload-extension', () => {
-      browser.runtime.reload();
-    });
-    ws.addWxtEventListener('wxt:reload-content-script', (event) => {
-      reloadContentScript(event.detail);
-    });
+if (import.meta.hot) {
+  import.meta.hot.on('wxt:reload-extension', () => browser.runtime.reload());
+  import.meta.hot.on('wxt:reload-content-script', (event) =>
+    reloadContentScript(event.detail),
+  );
 
-    if (import.meta.env.MANIFEST_VERSION === 3) {
-      // Tell the server the background script is loaded and ready to go
-      ws.addEventListener('open', () =>
-        ws.sendCustom('wxt:background-initialized'),
-      );
+  if (import.meta.env.MANIFEST_VERSION === 3) {
+    let backgroundInitialized = false;
+    // Tell the server the background script is loaded and ready to go
+    import.meta.hot.on('vite:ws:connect', () => {
+      if (backgroundInitialized) return;
 
-      // Web Socket will disconnect if the service worker is killed
-      keepServiceWorkerAlive();
-    }
-  } catch (err) {
-    logger.error('Failed to setup web socket connection with dev server', err);
+      import.meta.hot?.send('wxt:background-initialized');
+      backgroundInitialized = true;
+    });
+
+    // Web Socket will disconnect if the service worker is killed
+    keepServiceWorkerAlive();
   }
 
   browser.commands.onCommand.addListener((command) => {

packages/wxt/src/virtual/background-entrypoint.ts

@@ -1,14 +1,6 @@
-import { logger } from '../sandbox/utils/logger';
-import { getDevServerWebSocket } from '../sandbox/dev-server-websocket';
-
-if (import.meta.env.COMMAND === 'serve') {
-  try {
-    const ws = getDevServerWebSocket();
-    ws.addWxtEventListener('wxt:reload-page', (event) => {
-      // "popup.html" === "/popup.html".substring(1)
-      if (event.detail === location.pathname.substring(1)) location.reload();
-    });
-  } catch (err) {
-    logger.error('Failed to setup web socket connection with dev server', err);
-  }
+if (import.meta.hot) {
+  import.meta.hot.on('wxt:reload-page', (event) => {
+    // "popup.html" === "/popup.html".substring(1)
+    if (event.detail === location.pathname.substring(1)) location.reload();
+  });
 }

packages/wxt/src/virtual/reload-html.ts

@@ -1,7 +1,6 @@
 import { browser } from 'wxt/browser';
 import { logger } from '../../sandbox/utils/logger';
 import { MatchPattern } from 'wxt/sandbox';
-import type { ReloadContentScriptPayload } from '../../sandbox/dev-server-websocket';
 
 export function reloadContentScript(payload: ReloadContentScriptPayload) {
   const manifest = browser.runtime.getManifest();
@@ -97,3 +96,12 @@ export async function reloadContentScriptMv2(
 ) {
   throw Error('TODO: reloadContentScriptMv2');
 }
+
+export interface ReloadContentScriptPayload {
+  registration?: 'manifest' | 'runtime';
+  contentScript: {
+    matches: string[];
+    js?: string[];
+    css?: string[];
+  };
+}

packages/wxt/src/virtual/utils/reload-content-scripts.ts

@@ -18,7 +18,7 @@ export default defineConfig({
       ignored: '**/dist/**',
     },
   },
-  plugins: [RandomSeed()],
+  plugins: [RandomSeed() as any],
   resolve: {
     alias: {
       'wxt/testing': path.resolve('src/testing'),