fix: content script is incorrectly invalidated when injected multiple times

packages/wxt/src/utils/content-script-context.ts

@@ -43,23 +43,19 @@ export class ContentScriptContext implements AbortController {
     'wxt:content-script-started',
   );
 
-  private isTopFrame = window.self === window.top;
+  private id: string;
   private abortController: AbortController;
   private locationWatcher = createLocationWatcher(this);
-  private receivedMessageIds = new Set<string>();
 
   constructor(
     private readonly contentScriptName: string,
     public readonly options?: Omit<ContentScriptDefinition, 'main'>,
   ) {
+    this.id = Math.random().toString(36).slice(2);
     this.abortController = new AbortController();
 
-    if (this.isTopFrame) {
-      this.listenForNewerScripts({ ignoreFirstEvent: true });
-      this.stopOldScripts();
-    } else {
-      this.listenForNewerScripts();
-    }
+    this.stopOldScripts();
+    this.listenForNewerScripts();
   }
 
   get signal() {
@@ -243,43 +239,56 @@ export class ContentScriptContext implements AbortController {
   }
 
   stopOldScripts() {
-    // Use postMessage so it get's sent to all the frames of the page.
+    document.dispatchEvent(
+      new CustomEvent(ContentScriptContext.SCRIPT_STARTED_MESSAGE_TYPE, {
+        detail: {
+          contentScriptName: this.contentScriptName,
+          messageId: this.id,
+        },
+      }),
+    );
+
+    // Send message using `window.postMessage` for backwards compatibility to invalidate old versions before WXT changed to `document.dispatchEvent`
+    // TODO: Remove this once WXT version using `document.dispatchEvent` has been released for a while
     window.postMessage(
       {
         type: ContentScriptContext.SCRIPT_STARTED_MESSAGE_TYPE,
         contentScriptName: this.contentScriptName,
-        messageId: Math.random().toString(36).slice(2),
+        messageId: this.id,
       },
       '*',
     );
   }
 
-  verifyScriptStartedEvent(event: MessageEvent) {
-    const isScriptStartedEvent =
-      event.data?.type === ContentScriptContext.SCRIPT_STARTED_MESSAGE_TYPE;
+  verifyScriptStartedEvent(event: CustomEvent) {
     const isSameContentScript =
-      event.data?.contentScriptName === this.contentScriptName;
-    const isNotDuplicate = !this.receivedMessageIds.has(event.data?.messageId);
-    return isScriptStartedEvent && isSameContentScript && isNotDuplicate;
+      event.detail?.contentScriptName === this.contentScriptName;
+    // Handle case where website dispatches the event again for some reason
+    const isFromSelf = event.detail?.messageId === this.id;
+    return isSameContentScript && !isFromSelf;
   }
 
-  listenForNewerScripts(options?: { ignoreFirstEvent?: boolean }) {
-    let isFirst = true;
-
-    const cb = (event: MessageEvent) => {
-      if (this.verifyScriptStartedEvent(event)) {
-        this.receivedMessageIds.add(event.data.messageId);
-
-        const wasFirst = isFirst;
-        isFirst = false;
-        if (wasFirst && options?.ignoreFirstEvent) return;
-
-        this.notifyInvalidated();
+  listenForNewerScripts() {
+    const cb: EventListener = (event) => {
+      if (
+        !(event instanceof CustomEvent) ||
+        !this.verifyScriptStartedEvent(event)
+      ) {
+        return;
       }
+      this.notifyInvalidated();
     };
 
-    addEventListener('message', cb);
-    this.onInvalidated(() => removeEventListener('message', cb));
+    document.addEventListener(
+      ContentScriptContext.SCRIPT_STARTED_MESSAGE_TYPE,
+      cb,
+    );
+    this.onInvalidated(() =>
+      document.removeEventListener(
+        ContentScriptContext.SCRIPT_STARTED_MESSAGE_TYPE,
+        cb,
+      ),
+    );
   }
 }