feat: generalize worker realm (#1779)

This PR generalizes DedicatedWorkerRealm to other worker realms.

Also (hopefully) fixes the flaky test by ensuring the worker is actively
doing something for the duration of the test.

Fixed: #1166,
#1830

src/bidiMapper/domains/context/BrowsingContextProcessor.ts

@@ -28,16 +28,22 @@ import {
 import {CdpErrorConstants} from '../../../utils/CdpErrorConstants.js';
 import {LogType, type LoggerFn} from '../../../utils/log.js';
 import type {NetworkStorage} from '../network/NetworkStorage.js';
-import {DedicatedWorkerRealm} from '../script/DedicatedWorkerRealm.js';
 import type {PreloadScriptStorage} from '../script/PreloadScriptStorage.js';
 import type {Realm} from '../script/Realm.js';
 import type {RealmStorage} from '../script/RealmStorage.js';
+import {WorkerRealm, type WorkerRealmType} from '../script/WorkerRealm.js';
 import type {EventManager} from '../session/EventManager.js';
 
 import {BrowsingContextImpl, serializeOrigin} from './BrowsingContextImpl.js';
 import type {BrowsingContextStorage} from './BrowsingContextStorage.js';
 import {CdpTarget} from './CdpTarget.js';
 
+const cdpToBidiTargetTypes = {
+  service_worker: 'service-worker',
+  shared_worker: 'shared-worker',
+  worker: 'dedicated-worker',
+} as const;
+
 export class BrowsingContextProcessor {
   readonly #browserCdpClient: CdpClient;
   readonly #cdpConnection: CdpConnection;
@@ -325,6 +331,9 @@ export class BrowsingContextProcessor {
     cdpClient.on('Target.targetInfoChanged', (params) => {
       this.#handleTargetInfoChangedEvent(params);
     });
+    cdpClient.on('Inspector.targetCrashed', () => {
+      this.#handleTargetCrashedEvent(cdpClient);
+    });
 
     cdpClient.on(
       'Page.frameAttached',
@@ -413,25 +422,33 @@ export class BrowsingContextProcessor {
         }
         return;
       }
+      case 'service_worker':
       case 'worker': {
-        const browsingContext =
-          parentSessionCdpClient.sessionId &&
-          this.#browsingContextStorage.findContextBySession(
-            parentSessionCdpClient.sessionId
-          );
+        const realm = this.#realmStorage.findRealm({
+          cdpSessionId: parentSessionCdpClient.sessionId,
+        });
         // If there is no browsing context, this worker is already terminated.
-        if (!browsingContext) {
+        if (!realm) {
           break;
         }
 
         const cdpTarget = this.#createCdpTarget(targetCdpClient, targetInfo);
         this.#handleWorkerTarget(
+          cdpToBidiTargetTypes[targetInfo.type],
           cdpTarget,
-          this.#realmStorage.getRealm({
-            browsingContextId: browsingContext.id,
-            type: 'window',
-            sandbox: undefined,
-          })
+          realm
+        );
+        return;
+      }
+      // In CDP, we only emit shared workers on the browser and not the set of
+      // frames that use the shared worker. If we change this in the future to
+      // behave like service workers (emits on both browser and frame targets),
+      // we can remove this block and merge service workers with the above one.
+      case 'shared_worker': {
+        const cdpTarget = this.#createCdpTarget(targetCdpClient, targetInfo);
+        this.#handleWorkerTarget(
+          cdpToBidiTargetTypes[targetInfo.type],
+          cdpTarget
         );
         return;
       }
@@ -467,18 +484,23 @@ export class BrowsingContextProcessor {
   }
 
   #workers = new Map<string, Realm>();
-  #handleWorkerTarget(cdpTarget: CdpTarget, ownerRealm: Realm) {
+  #handleWorkerTarget(
+    realmType: WorkerRealmType,
+    cdpTarget: CdpTarget,
+    ownerRealm?: Realm
+  ) {
     cdpTarget.cdpClient.on('Runtime.executionContextCreated', (params) => {
       const {uniqueId, id, origin} = params.context;
-      const workerRealm = new DedicatedWorkerRealm(
+      const workerRealm = new WorkerRealm(
         cdpTarget.cdpClient,
         this.#eventManager,
         id,
         this.#logger,
         serializeOrigin(origin),
-        ownerRealm,
+        ownerRealm ? [ownerRealm] : [],
         uniqueId,
-        this.#realmStorage
+        this.#realmStorage,
+        realmType
       );
       this.#workers.set(cdpTarget.cdpSessionId, workerRealm);
     });
@@ -516,4 +538,16 @@ export class BrowsingContextProcessor {
       context.onTargetInfoChanged(params);
     }
   }
+
+  #handleTargetCrashedEvent(cdpClient: CdpClient) {
+    // This is primarily used for service and shared workers. CDP tends to not
+    // signal they closed gracefully and instead says they crashed to signal
+    // they are closed.
+    const realms = this.#realmStorage.findRealms({
+      cdpSessionId: cdpClient.sessionId,
+    });
+    for (const realm of realms) {
+      realm.dispose();
+    }
+  }
 }

src/bidiMapper/domains/script/Realm.ts

@@ -234,19 +234,24 @@ export abstract class Realm {
     };
   }
 
-  protected initialize() {
-    for (const browsingContext of this.associatedBrowsingContexts) {
-      this.#eventManager.registerEvent(
-        {
-          type: 'event',
-          method: ChromiumBidi.Script.EventNames.RealmCreated,
-          params: this.realmInfo,
-        },
-        browsingContext.id
-      );
+  #registerEvent(event: ChromiumBidi.Event) {
+    if (this.associatedBrowsingContexts.length === 0) {
+      this.#eventManager.registerEvent(event, null);
+    } else {
+      for (const browsingContext of this.associatedBrowsingContexts) {
+        this.#eventManager.registerEvent(event, browsingContext.id);
+      }
     }
   }
 
+  protected initialize() {
+    this.#registerEvent({
+      type: 'event',
+      method: ChromiumBidi.Script.EventNames.RealmCreated,
+      params: this.realmInfo,
+    });
+  }
+
   /**
    * Serializes a given CDP object into BiDi, keeping references in the
    * target's `globalThis`.
@@ -707,17 +712,12 @@ export abstract class Realm {
   }
 
   dispose(): void {
-    for (const browsingContext of this.associatedBrowsingContexts) {
-      this.#eventManager.registerEvent(
-        {
-          type: 'event',
-          method: ChromiumBidi.Script.EventNames.RealmDestroyed,
-          params: {
-            realm: this.realmId,
-          },
-        },
-        browsingContext.id
-      );
-    }
+    this.#registerEvent({
+      type: 'event',
+      method: ChromiumBidi.Script.EventNames.RealmDestroyed,
+      params: {
+        realm: this.realmId,
+      },
+    });
   }
 }

src/bidiMapper/domains/script/DedicatedWorkerRealm.ts → src/bidiMapper/domains/script/WorkerRealm.ts

@@ -26,18 +26,27 @@ import type {EventManager} from '../session/EventManager.js';
 import {Realm} from './Realm.js';
 import type {RealmStorage} from './RealmStorage.js';
 
-export class DedicatedWorkerRealm extends Realm {
-  readonly #ownerRealm: Realm;
+export type WorkerRealmType = Pick<
+  | Script.SharedWorkerRealmInfo
+  | Script.DedicatedWorkerRealmInfo
+  | Script.ServiceWorkerRealmInfo,
+  'type'
+>['type'];
+
+export class WorkerRealm extends Realm {
+  readonly #realmType: WorkerRealmType;
+  readonly #ownerRealms: Realm[];
 
   constructor(
     cdpClient: CdpClient,
     eventManager: EventManager,
     executionContextId: Protocol.Runtime.ExecutionContextId,
     logger: LoggerFn | undefined,
     origin: string,
-    ownerRealm: Realm,
+    ownerRealms: Realm[],
     realmId: Script.Realm,
-    realmStorage: RealmStorage
+    realmStorage: RealmStorage,
+    realmType: WorkerRealmType
   ) {
     super(
       cdpClient,
@@ -49,17 +58,20 @@ export class DedicatedWorkerRealm extends Realm {
       realmStorage
     );
 
-    this.#ownerRealm = ownerRealm;
+    this.#ownerRealms = ownerRealms;
+    this.#realmType = realmType;
 
     this.initialize();
   }
 
   override get associatedBrowsingContexts(): BrowsingContextImpl[] {
-    return this.#ownerRealm.associatedBrowsingContexts;
+    return this.#ownerRealms.flatMap(
+      (realm) => realm.associatedBrowsingContexts
+    );
   }
 
-  override get realmType(): 'dedicated-worker' {
-    return 'dedicated-worker';
+  override get realmType(): WorkerRealmType {
+    return this.#realmType;
   }
 
   override get source(): Script.Source {
@@ -71,11 +83,28 @@ export class DedicatedWorkerRealm extends Realm {
     };
   }
 
-  override get realmInfo(): Script.DedicatedWorkerRealmInfo {
-    return {
-      ...this.baseInfo,
-      type: this.realmType,
-      owners: [this.#ownerRealm.realmId],
-    };
+  override get realmInfo(): Script.RealmInfo {
+    const owners = this.#ownerRealms.map((realm) => realm.realmId);
+    const {realmType} = this;
+    switch (realmType) {
+      case 'dedicated-worker': {
+        const owner = owners[0];
+        if (owner === undefined || owners.length !== 1) {
+          throw new Error('Dedicated worker must have exactly one owner');
+        }
+        return {
+          ...this.baseInfo,
+          type: realmType,
+          owners: [owner],
+        };
+      }
+      case 'service-worker':
+      case 'shared-worker': {
+        return {
+          ...this.baseInfo,
+          type: realmType,
+        };
+      }
+    }
   }
 }

tests/script/test_evaluate.py

@@ -399,11 +399,10 @@ async def test_scriptEvaluate_realm(websocket, context_id):
     } == result
 
 
-@pytest.mark.timeout(20)
 @pytest.mark.asyncio
 async def test_scriptEvaluate_dedicated_worker(websocket, context_id, html,
                                                snapshot):
-    worker_url = 'data:application/javascript,'
+    worker_url = 'data:text/javascript,setTimeout(() => {}, 20000)'
     url = html(f"<script>window.w = new Worker('{worker_url}');</script>")
 
     await subscribe(websocket, ["script.realmCreated"])