fix: Use HMR API for dev server communication

packages/wxt-demo/src/entrypoints/background.ts

@@ -1,5 +1,5 @@
 export default defineBackground({
-  // type: 'module',
+  type: 'module',
 
   main() {
     console.log(browser.runtime.id);

packages/wxt/build.config.ts

@@ -32,10 +32,13 @@ export default defineBuildConfig([
       ...virtualEntrypointModuleNames.map((name) => `virtual:user-${name}`),
       'virtual:wxt-plugins',
       'virtual:app-config',
+      'wxt/background-client',
       'wxt/browser',
       'wxt/sandbox',
       'wxt/client',
       'wxt/testing',
+      'http://localhost:3000/@id/wxt/background-client',
+      'http://localhost:3000/@vite/client',
     ],
   })),
 ]);

packages/wxt/package.json

@@ -67,6 +67,10 @@
     "./modules": {
       "types": "./dist/modules.d.ts",
       "default": "./dist/modules.mjs"
+    },
+    "./background-client": {
+      "types": "./dist/background-client.d.ts",
+      "default": "./dist/background-client.mjs"
     }
   },
   "scripts": {

packages/wxt/src/background-client.ts

@@ -0,0 +1,141 @@
+import { browser } from 'wxt/browser';
+import { logger } from './sandbox/utils/logger';
+import { MatchPattern } from 'wxt/sandbox';
+
+console.log('HOT', import.meta.hot);
+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) {
+    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;
+
+      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) => {
+    if (command === 'wxt:reload-extension') {
+      browser.runtime.reload();
+    }
+  });
+
+  /**
+   * https://developer.chrome.com/blog/longer-esw-lifetimes/
+   */
+  function keepServiceWorkerAlive() {
+    setInterval(async () => {
+      // Calling an async browser API resets the service worker's timeout
+      await browser.runtime.getPlatformInfo();
+    }, 5e3);
+  }
+
+  function reloadContentScript(payload: ReloadContentScriptPayload) {
+    const manifest = browser.runtime.getManifest();
+    if (manifest.manifest_version == 2) {
+      void reloadContentScriptMv2(payload);
+    } else {
+      void reloadContentScriptMv3(payload);
+    }
+  }
+
+  async function reloadContentScriptMv3({
+    registration,
+    contentScript,
+  }: ReloadContentScriptPayload) {
+    if (registration === 'runtime') {
+      await reloadRuntimeContentScriptMv3(contentScript);
+    } else {
+      await reloadManifestContentScriptMv3(contentScript);
+    }
+  }
+
+  type ContentScript = ReloadContentScriptPayload['contentScript'];
+
+  async function reloadManifestContentScriptMv3(contentScript: ContentScript) {
+    const id = `wxt:${contentScript.js![0]}`;
+    logger.log('Reloading content script:', contentScript);
+    const registered = await browser.scripting.getRegisteredContentScripts();
+    logger.debug('Existing scripts:', registered);
+
+    const existing = registered.find((cs) => cs.id === id);
+
+    if (existing) {
+      logger.debug('Updating content script', existing);
+      await browser.scripting.updateContentScripts([{ ...contentScript, id }]);
+    } else {
+      logger.debug('Registering new content script...');
+      await browser.scripting.registerContentScripts([
+        { ...contentScript, id },
+      ]);
+    }
+
+    await reloadTabsForContentScript(contentScript);
+  }
+
+  async function reloadRuntimeContentScriptMv3(contentScript: ContentScript) {
+    logger.log('Reloading content script:', contentScript);
+    const registered = await browser.scripting.getRegisteredContentScripts();
+    logger.debug('Existing scripts:', registered);
+
+    const matches = registered.filter((cs) => {
+      const hasJs = contentScript.js?.find((js) => cs.js?.includes(js));
+      const hasCss = contentScript.css?.find((css) => cs.css?.includes(css));
+      return hasJs || hasCss;
+    });
+
+    if (matches.length === 0) {
+      logger.log(
+        'Content script is not registered yet, nothing to reload',
+        contentScript,
+      );
+      return;
+    }
+
+    await browser.scripting.updateContentScripts(matches);
+    await reloadTabsForContentScript(contentScript);
+  }
+
+  async function reloadTabsForContentScript(contentScript: ContentScript) {
+    const allTabs = await browser.tabs.query({});
+    const matchPatterns = contentScript.matches.map(
+      (match) => new MatchPattern(match),
+    );
+    const matchingTabs = allTabs.filter((tab) => {
+      const url = tab.url;
+      if (!url) return false;
+      return !!matchPatterns.find((pattern) => pattern.includes(url));
+    });
+    await Promise.all(
+      matchingTabs.map(async (tab) => {
+        try {
+          await browser.tabs.reload(tab.id);
+        } catch (err) {
+          logger.warn('Failed to reload tab:', err);
+        }
+      }),
+    );
+  }
+
+  async function reloadContentScriptMv2(_payload: ReloadContentScriptPayload) {
+    throw Error('TODO: reloadContentScriptMv2');
+  }
+
+  interface ReloadContentScriptPayload {
+    registration?: 'manifest' | 'runtime';
+    contentScript: {
+      matches: string[];
+      js?: string[];
+      css?: string[];
+    };
+  }
+}

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

@@ -1,40 +1,8 @@
 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.env.MANIFEST_VERSION === 3) {
-      // Tell the server the background script is loaded and ready to go
-      ws.addEventListener('open', () =>
-        ws.sendCustom('wxt:background-initialized'),
-      );
-
-      // 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);
-  }
-
-  browser.commands.onCommand.addListener((command) => {
-    if (command === 'wxt:reload-extension') {
-      browser.runtime.reload();
-    }
-  });
-}
+/* @vite-ignore */
+import 'http://localhost:3000/@id/wxt/background-client';
 
 let result;
 

packages/wxt/src/virtual/reload-html.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();
+  });
 }