From 15bf0da82e064d2dcc907c892f43971683cec5e0 Mon Sep 17 00:00:00 2001
From: Aaron <aaronklinker1@gmail.com>
Date: Fri, 2 Aug 2024 16:43:40 -0500
Subject: [PATCH] fix(content-script-context): Fix initialization logic for
 Firefox (#895)

---
 .../__tests__/content-script-context.test.ts  | 59 +++++++++++++++++++
 .../content-scripts/content-script-context.ts | 16 +++--
 2 files changed, 70 insertions(+), 5 deletions(-)
 create mode 100644 packages/wxt/src/client/content-scripts/__tests__/content-script-context.test.ts

diff --git a/packages/wxt/src/client/content-scripts/__tests__/content-script-context.test.ts b/packages/wxt/src/client/content-scripts/__tests__/content-script-context.test.ts
new file mode 100644
index 00000000..6f7d14c1
--- /dev/null
+++ b/packages/wxt/src/client/content-scripts/__tests__/content-script-context.test.ts
@@ -0,0 +1,59 @@
+/** @vitest-environment happy-dom */
+
+import { describe, it, expect, vi, beforeEach } from 'vitest';
+import { ContentScriptContext } from '..';
+import { fakeBrowser } from '@webext-core/fake-browser';
+import { sleep } from '../../../core/utils/time';
+
+describe('Content Script Context', () => {
+  beforeEach(() => {
+    vi.useRealTimers();
+    fakeBrowser.runtime.id = 'anything';
+  });
+
+  it("should recognize when the content script has lost it's connection to the extension API", () => {
+    const ctx = new ContentScriptContext('test');
+    const onInvalidated = vi.fn();
+
+    ctx.onInvalidated(onInvalidated);
+    // @ts-expect-error
+    delete fakeBrowser.runtime.id;
+    const isValid = ctx.isValid;
+
+    expect(onInvalidated).toBeCalled();
+    expect(isValid).toBe(false);
+  });
+
+  it('should invalidate the current content script when a new context is created', async () => {
+    const name = 'test';
+    const onInvalidated = vi.fn();
+    const ctx = new ContentScriptContext(name);
+    ctx.onInvalidated(onInvalidated);
+
+    // Wait for events to run before next tick next tick
+    await sleep(0);
+
+    // Create a new context after first is initialized, and wait for it to initialize
+    new ContentScriptContext(name);
+    await sleep(0);
+
+    expect(onInvalidated).toBeCalled();
+    expect(ctx.isValid).toBe(false);
+  });
+
+  it('should not invalidate the current content script when a new context is created with a different name', async () => {
+    const onInvalidated = vi.fn();
+    const ctx = new ContentScriptContext('test1');
+    ctx.onInvalidated(onInvalidated);
+
+    // Wait for events to run before next tick next tick
+    await sleep(0);
+
+    // Create a new context after first is initialized, and wait for it to initialize
+    new ContentScriptContext('test2');
+    await sleep(0);
+
+    expect(onInvalidated).not.toBeCalled();
+    expect(ctx.isValid).toBe(true);
+  });
+});
diff --git a/packages/wxt/src/client/content-scripts/content-script-context.ts b/packages/wxt/src/client/content-scripts/content-script-context.ts
index 47d414e0..38cb9ed3 100644
--- a/packages/wxt/src/client/content-scripts/content-script-context.ts
+++ b/packages/wxt/src/client/content-scripts/content-script-context.ts
@@ -46,13 +46,13 @@ export class ContentScriptContext implements AbortController {
     public readonly options?: Omit<ContentScriptDefinition, 'main'>,
   ) {
     this.#abortController = new AbortController();
+
     if (this.#isTopFrame) {
+      this.#listenForNewerScripts({ ignoreFirstEvent: true });
       this.#stopOldScripts();
-    }
-    this.setTimeout(() => {
-      // Run on next tick so the listener it adds isn't triggered by stopOldScript
+    } else {
       this.#listenForNewerScripts();
-    });
+    }
   }
 
   get signal() {
@@ -224,12 +224,18 @@ export class ContentScriptContext implements AbortController {
     );
   }
 
-  #listenForNewerScripts() {
+  #listenForNewerScripts(options?: { ignoreFirstEvent?: boolean }) {
+    let isFirst = true;
+
     const cb = (event: MessageEvent) => {
       if (
         event.data?.type === ContentScriptContext.SCRIPT_STARTED_MESSAGE_TYPE &&
         event.data?.contentScriptName === this.contentScriptName
       ) {
+        const wasFirst = isFirst;
+        isFirst = false;
+        if (wasFirst && options?.ignoreFirstEvent) return;
+
         this.notifyInvalidated();
       }
     };