refactor: unify tab target handling in page (#11115)

2024-06-14 14:02:48 +00:00 · 2023-10-16 09:53:03 +02:00 · 2023-10-16 09:53:03 +02:00 · bb9ef6ee8b
commit bb9ef6ee8b
parent 78c335e611
3 changed files with 174 additions and 121 deletions
--- a/packages/puppeteer-core/src/cdp/CDPSession.ts
+++ b/packages/puppeteer-core/src/cdp/CDPSession.ts
@ -68,7 +68,8 @@ export class CdpCDPSession extends CDPSession {

  override parentSession(): CDPSession | undefined {
    if (!this.#parentSessionId) {
-      return;
+      // To make it work in Firefox that does not have parent (tab) sessions.
+      return this;
    }
    const parent = this.#connection?.session(this.#parentSessionId);
    return parent ?? undefined;
--- a/packages/puppeteer-core/src/cdp/FirefoxTargetManager.ts
+++ b/packages/puppeteer-core/src/cdp/FirefoxTargetManager.ts
@ -22,6 +22,7 @@ import {EventEmitter} from '../common/EventEmitter.js';
 import {assert} from '../util/assert.js';
 import {Deferred} from '../util/Deferred.js';

+import type {CdpCDPSession} from './CDPSession.js';
 import type {Connection} from './Connection.js';
 import type {CdpTarget} from './Target.js';
 import {
@ -205,6 +206,7 @@ export class FirefoxTargetManager

    assert(target, `Target ${targetInfo.targetId} is missing`);

+    (session as CdpCDPSession)._setTarget(target);
    this.setupAttachmentListeners(session);

    this.#availableTargetsBySessionId.set(
--- a/packages/puppeteer-core/src/cdp/Page.ts
+++ b/packages/puppeteer-core/src/cdp/Page.ts
@ -80,6 +80,7 @@ import {CdpKeyboard, CdpMouse, CdpTouchscreen} from './Input.js';
 import {MAIN_WORLD} from './IsolatedWorlds.js';
 import type {Credentials, NetworkConditions} from './NetworkManager.js';
 import type {CdpTarget} from './Target.js';
+import type {TargetManager} from './TargetManager.js';
 import {TargetManagerEvent} from './TargetManager.js';
 import {Tracing} from './Tracing.js';
 import {WebWorker} from './WebWorker.js';
@ -111,9 +112,12 @@ export class CdpPage extends Page {
  }

  #closed = false;
-  #client: CDPSession;
-  #tabSession: CDPSession | undefined;
-  #target: CdpTarget;
+  readonly #targetManager: TargetManager;
+
+  #primaryTargetClient: CDPSession;
+  #primaryTarget: CdpTarget;
+  #tabTargetClient: CDPSession;
+  #tabTarget: CdpTarget;
  #keyboard: CdpKeyboard;
  #mouse: CdpMouse;
  #touchscreen: CdpTouchscreen;
@ -222,9 +226,13 @@ export class CdpPage extends Page {
    ignoreHTTPSErrors: boolean
  ) {
    super();
-    this.#client = client;
-    this.#tabSession = client.parentSession();
-    this.#target = target;
+    this.#primaryTargetClient = client;
+    this.#tabTargetClient = client.parentSession()!;
+    assert(this.#tabTargetClient, 'Tab target session is not defined.');
+    this.#tabTarget = (this.#tabTargetClient as CdpCDPSession)._target();
+    assert(this.#tabTarget, 'Tab target is not defined.');
+    this.#primaryTarget = target;
+    this.#targetManager = target._targetManager();
    this.#keyboard = new CdpKeyboard(client);
    this.#mouse = new CdpMouse(client, this.#keyboard);
    this.#touchscreen = new CdpTouchscreen(client, this.#keyboard);
@ -249,36 +257,65 @@ export class CdpPage extends Page {
      this.#frameManager.networkManager.on(eventName, handler as any);
    }

-    this.#setupPrimaryTargetListeners();
+    this.#tabTargetClient.on(
+      CDPSessionEvent.Swapped,
+      this.#onActivation.bind(this)
+    );

-    this.#tabSession?.on(CDPSessionEvent.Swapped, async newSession => {
-      this.#client = newSession;
-      assert(
-        this.#client instanceof CdpCDPSession,
-        'CDPSession is not instance of CDPSessionImpl'
-      );
-      this.#target = this.#client._target();
-      assert(this.#target, 'Missing target on swap');
-      this.#keyboard.updateClient(newSession);
-      this.#mouse.updateClient(newSession);
-      this.#touchscreen.updateClient(newSession);
-      this.#accessibility.updateClient(newSession);
-      this.#emulationManager.updateClient(newSession);
-      this.#tracing.updateClient(newSession);
-      this.#coverage.updateClient(newSession);
-      await this.#frameManager.swapFrameTree(newSession);
-      this.#setupPrimaryTargetListeners();
-    });
-    this.#tabSession?.on(CDPSessionEvent.Ready, session => {
-      assert(session instanceof CdpCDPSession);
-      if (session._target()._subtype() !== 'prerender') {
-        return;
-      }
-      this.#frameManager.registerSpeculativeSession(session).catch(debugError);
-      this.#emulationManager
-        .registerSpeculativeSession(session)
-        .catch(debugError);
-    });
+    this.#tabTargetClient.on(
+      CDPSessionEvent.Ready,
+      this.#onSecondaryTarget.bind(this)
+    );
+
+    this.#targetManager.on(
+      TargetManagerEvent.TargetGone,
+      this.#onDetachedFromTarget
+    );
+
+    this.#tabTarget._isClosedDeferred
+      .valueOrThrow()
+      .then(() => {
+        this.#targetManager.off(
+          TargetManagerEvent.TargetGone,
+          this.#onDetachedFromTarget
+        );
+
+        this.emit(PageEvent.Close, undefined);
+        this.#closed = true;
+      })
+      .catch(debugError);
+
+    this.#setupPrimaryTargetListeners();
+  }
+
+  async #onActivation(newSession: CDPSession): Promise<void> {
+    this.#primaryTargetClient = newSession;
+    assert(
+      this.#primaryTargetClient instanceof CdpCDPSession,
+      'CDPSession is not instance of CDPSessionImpl'
+    );
+    this.#primaryTarget = this.#primaryTargetClient._target();
+    assert(this.#primaryTarget, 'Missing target on swap');
+    this.#keyboard.updateClient(newSession);
+    this.#mouse.updateClient(newSession);
+    this.#touchscreen.updateClient(newSession);
+    this.#accessibility.updateClient(newSession);
+    this.#emulationManager.updateClient(newSession);
+    this.#tracing.updateClient(newSession);
+    this.#coverage.updateClient(newSession);
+    await this.#frameManager.swapFrameTree(newSession);
+    this.#setupPrimaryTargetListeners();
+  }
+
+  async #onSecondaryTarget(session: CDPSession): Promise<void> {
+    assert(session instanceof CdpCDPSession);
+    if (session._target()._subtype() !== 'prerender') {
+      return;
+    }
+    this.#frameManager.registerSpeculativeSession(session).catch(debugError);
+    this.#emulationManager
+      .registerSpeculativeSession(session)
+      .catch(debugError);
  }

  /**
@ -286,30 +323,15 @@ export class CdpPage extends Page {
   * during a navigation to a prerended page.
   */
  #setupPrimaryTargetListeners() {
-    this.#client.on(CDPSessionEvent.Ready, this.#onAttachedToTarget);
+    this.#primaryTargetClient.on(
+      CDPSessionEvent.Ready,
+      this.#onAttachedToTarget
+    );

    for (const [eventName, handler] of this.#sessionHandlers) {
      // TODO: Remove any.
-      this.#client.on(eventName, handler as any);
+      this.#primaryTargetClient.on(eventName, handler as any);
    }
-
-    this.#target
-      ._targetManager()
-      .on(TargetManagerEvent.TargetGone, this.#onDetachedFromTarget);
-
-    this.#target._isClosedDeferred
-      .valueOrThrow()
-      .then(() => {
-        this.#client.off(CDPSessionEvent.Ready, this.#onAttachedToTarget);
-
-        this.#target
-          ._targetManager()
-          .off(TargetManagerEvent.TargetGone, this.#onDetachedFromTarget);
-
-        this.emit(PageEvent.Close, undefined);
-        this.#closed = true;
-      })
-      .catch(debugError);
  }

  #onDetachedFromTarget = (target: CdpTarget) => {
@ -341,9 +363,9 @@ export class CdpPage extends Page {
  async #initialize(): Promise<void> {
    try {
      await Promise.all([
-        this.#frameManager.initialize(this.#client),
-        this.#client.send('Performance.enable'),
-        this.#client.send('Log.enable'),
+        this.#frameManager.initialize(this.#primaryTargetClient),
+        this.#primaryTargetClient.send('Performance.enable'),
+        this.#primaryTargetClient.send('Log.enable'),
      ]);
    } catch (err) {
      if (isErrorLike(err) && isTargetClosedError(err)) {
@ -377,7 +399,7 @@ export class CdpPage extends Page {
  }

  _client(): CDPSession {
-    return this.#client;
+    return this.#primaryTargetClient;
  }

  override isServiceWorkerBypassed(): boolean {
@ -404,9 +426,12 @@ export class CdpPage extends Page {
    this.#fileChooserDeferreds.add(deferred);
    let enablePromise: Promise<void> | undefined;
    if (needsEnable) {
-      enablePromise = this.#client.send('Page.setInterceptFileChooserDialog', {
-        enabled: true,
-      });
+      enablePromise = this.#primaryTargetClient.send(
+        'Page.setInterceptFileChooserDialog',
+        {
+          enabled: true,
+        }
+      );
    }
    try {
      const [result] = await Promise.all([
@ -425,15 +450,15 @@ export class CdpPage extends Page {
  }

  override target(): CdpTarget {
-    return this.#target;
+    return this.#primaryTarget;
  }

  override browser(): Browser {
-    return this.#target.browser();
+    return this.#primaryTarget.browser();
  }

  override browserContext(): BrowserContext {
-    return this.#target.browserContext();
+    return this.#primaryTarget.browserContext();
  }

  #onTargetCrashed(): void {
@ -444,7 +469,7 @@ export class CdpPage extends Page {
    const {level, text, args, source, url, lineNumber} = event.entry;
    if (args) {
      args.map(arg => {
-        return releaseObject(this.#client, arg);
+        return releaseObject(this.#primaryTargetClient, arg);
      });
    }
    if (source !== 'worker') {
@ -495,12 +520,17 @@ export class CdpPage extends Page {

  override async setBypassServiceWorker(bypass: boolean): Promise<void> {
    this.#serviceWorkerBypassed = bypass;
-    return await this.#client.send('Network.setBypassServiceWorker', {bypass});
+    return await this.#primaryTargetClient.send(
+      'Network.setBypassServiceWorker',
+      {bypass}
+    );
  }

  override async setDragInterception(enabled: boolean): Promise<void> {
    this.#userDragInterceptionEnabled = enabled;
-    return await this.#client.send('Input.setInterceptDrags', {enabled});
+    return await this.#primaryTargetClient.send('Input.setInterceptDrags', {
+      enabled,
+    });
  }

  override async setOfflineMode(enabled: boolean): Promise<void> {
@ -551,7 +581,7 @@ export class CdpPage extends Page {
    ...urls: string[]
  ): Promise<Protocol.Network.Cookie[]> {
    const originalCookies = (
-      await this.#client.send('Network.getCookies', {
+      await this.#primaryTargetClient.send('Network.getCookies', {
        urls: urls.length ? urls : [this.url()],
      })
    ).cookies;
@ -577,7 +607,7 @@ export class CdpPage extends Page {
      if (!cookie.url && pageURL.startsWith('http')) {
        item.url = pageURL;
      }
-      await this.#client.send('Network.deleteCookies', item);
+      await this.#primaryTargetClient.send('Network.deleteCookies', item);
    }
  }

@ -603,7 +633,9 @@ export class CdpPage extends Page {
    });
    await this.deleteCookie(...items);
    if (items.length) {
-      await this.#client.send('Network.setCookies', {cookies: items});
+      await this.#primaryTargetClient.send('Network.setCookies', {
+        cookies: items,
+      });
    }
  }

@ -636,8 +668,8 @@ export class CdpPage extends Page {
    this.#bindings.set(name, binding);

    const expression = pageBindingInitString('exposedFun', name);
-    await this.#client.send('Runtime.addBinding', {name});
-    const {identifier} = await this.#client.send(
+    await this.#primaryTargetClient.send('Runtime.addBinding', {name});
+    const {identifier} = await this.#primaryTargetClient.send(
      'Page.addScriptToEvaluateOnNewDocument',
      {
        source: expression,
@ -661,7 +693,7 @@ export class CdpPage extends Page {
      );
    }

-    await this.#client.send('Runtime.removeBinding', {name});
+    await this.#primaryTargetClient.send('Runtime.removeBinding', {name});
    await this.removeScriptToEvaluateOnNewDocument(exposedFun);

    await Promise.all(
@ -701,7 +733,9 @@ export class CdpPage extends Page {
  }

  override async metrics(): Promise<Metrics> {
-    const response = await this.#client.send('Performance.getMetrics');
+    const response = await this.#primaryTargetClient.send(
+      'Performance.getMetrics'
+    );
    return this.#buildMetricsObject(response.metrics);
  }

@ -753,7 +787,7 @@ export class CdpPage extends Page {
    }
    const context = this.#frameManager.getExecutionContextById(
      event.executionContextId,
-      this.#client
+      this.#primaryTargetClient
    );
    if (!context) {
      debugError(
@ -789,7 +823,7 @@ export class CdpPage extends Page {

    const context = this.#frameManager.executionContextById(
      event.executionContextId,
-      this.#client
+      this.#primaryTargetClient
    );
    if (!context) {
      return;
@ -843,7 +877,7 @@ export class CdpPage extends Page {
  #onDialog(event: Protocol.Page.JavascriptDialogOpeningEvent): void {
    const type = validateDialogType(event.type);
    const dialog = new CdpDialog(
-      this.#client,
+      this.#primaryTargetClient,
      type,
      event.message,
      event.defaultPrompt
@ -856,7 +890,7 @@ export class CdpPage extends Page {
  ): Promise<HTTPResponse | null> {
    const [result] = await Promise.all([
      this.waitForNavigation(options),
-      this.#client.send('Page.reload'),
+      this.#primaryTargetClient.send('Page.reload'),
    ]);

    return result;
@ -925,20 +959,24 @@ export class CdpPage extends Page {
    delta: number,
    options: WaitForOptions
  ): Promise<HTTPResponse | null> {
-    const history = await this.#client.send('Page.getNavigationHistory');
+    const history = await this.#primaryTargetClient.send(
+      'Page.getNavigationHistory'
+    );
    const entry = history.entries[history.currentIndex + delta];
    if (!entry) {
      return null;
    }
    const result = await Promise.all([
      this.waitForNavigation(options),
-      this.#client.send('Page.navigateToHistoryEntry', {entryId: entry.id}),
+      this.#primaryTargetClient.send('Page.navigateToHistoryEntry', {
+        entryId: entry.id,
+      }),
    ]);
    return result[0];
  }

  override async bringToFront(): Promise<void> {
-    await this.#client.send('Page.bringToFront');
+    await this.#primaryTargetClient.send('Page.bringToFront');
  }

  override async setJavaScriptEnabled(enabled: boolean): Promise<void> {
@ -946,7 +984,7 @@ export class CdpPage extends Page {
  }

  override async setBypassCSP(enabled: boolean): Promise<void> {
-    await this.#client.send('Page.setBypassCSP', {enabled});
+    await this.#primaryTargetClient.send('Page.setBypassCSP', {enabled});
  }

  override async emulateMediaType(type?: string): Promise<void> {
@ -1000,7 +1038,7 @@ export class CdpPage extends Page {
    ...args: Params
  ): Promise<NewDocumentScriptEvaluation> {
    const source = evaluationString(pageFunction, ...args);
-    const {identifier} = await this.#client.send(
+    const {identifier} = await this.#primaryTargetClient.send(
      'Page.addScriptToEvaluateOnNewDocument',
      {
        source,
@ -1013,9 +1051,12 @@ export class CdpPage extends Page {
  override async removeScriptToEvaluateOnNewDocument(
    identifier: string
  ): Promise<void> {
-    await this.#client.send('Page.removeScriptToEvaluateOnNewDocument', {
-      identifier,
-    });
+    await this.#primaryTargetClient.send(
+      'Page.removeScriptToEvaluateOnNewDocument',
+      {
+        identifier,
+      }
+    );
  }

  override async setCacheEnabled(enabled = true): Promise<void> {
@ -1066,17 +1107,20 @@ export class CdpPage extends Page {
    }

    // We need to do these spreads because Firefox doesn't allow unknown options.
-    const {data} = await this.#client.send('Page.captureScreenshot', {
-      format: type,
-      ...(optimizeForSpeed ? {optimizeForSpeed} : {}),
-      ...(quality !== undefined ? {quality: Math.round(quality)} : {}),
-      clip: clip && {
-        ...clip,
-        scale: clip.scale ?? 1,
-      },
-      ...(!fromSurface ? {fromSurface} : {}),
-      captureBeyondViewport,
-    });
+    const {data} = await this.#primaryTargetClient.send(
+      'Page.captureScreenshot',
+      {
+        format: type,
+        ...(optimizeForSpeed ? {optimizeForSpeed} : {}),
+        ...(quality !== undefined ? {quality: Math.round(quality)} : {}),
+        clip: clip && {
+          ...clip,
+          scale: clip.scale ?? 1,
+        },
+        ...(!fromSurface ? {fromSurface} : {}),
+        captureBeyondViewport,
+      }
+    );
    return data;
  }

@ -1101,23 +1145,26 @@ export class CdpPage extends Page {
      await this.#emulationManager.setTransparentBackgroundColor();
    }

-    const printCommandPromise = this.#client.send('Page.printToPDF', {
-      transferMode: 'ReturnAsStream',
-      landscape,
-      displayHeaderFooter,
-      headerTemplate,
-      footerTemplate,
-      printBackground,
-      scale,
-      paperWidth,
-      paperHeight,
-      marginTop: margin.top,
-      marginBottom: margin.bottom,
-      marginLeft: margin.left,
-      marginRight: margin.right,
-      pageRanges,
-      preferCSSPageSize,
-    });
+    const printCommandPromise = this.#primaryTargetClient.send(
+      'Page.printToPDF',
+      {
+        transferMode: 'ReturnAsStream',
+        landscape,
+        displayHeaderFooter,
+        headerTemplate,
+        footerTemplate,
+        printBackground,
+        scale,
+        paperWidth,
+        paperHeight,
+        marginTop: margin.top,
+        marginBottom: margin.bottom,
+        marginLeft: margin.left,
+        marginRight: margin.right,
+        pageRanges,
+        preferCSSPageSize,
+      }
+    );

    const result = await waitWithTimeout(
      printCommandPromise,
@ -1130,7 +1177,10 @@ export class CdpPage extends Page {
    }

    assert(result.stream, '`stream` is missing from `Page.printToPDF');
-    return await getReadableFromProtocolStream(this.#client, result.stream);
+    return await getReadableFromProtocolStream(
+      this.#primaryTargetClient,
+      result.stream
+    );
  }

  override async pdf(options: PDFOptions = {}): Promise<Buffer> {
@ -1144,19 +1194,19 @@ export class CdpPage extends Page {
  override async close(
    options: {runBeforeUnload?: boolean} = {runBeforeUnload: undefined}
  ): Promise<void> {
-    const connection = this.#client.connection();
+    const connection = this.#primaryTargetClient.connection();
    assert(
      connection,
      'Protocol error: Connection closed. Most likely the page has been closed.'
    );
    const runBeforeUnload = !!options.runBeforeUnload;
    if (runBeforeUnload) {
-      await this.#client.send('Page.close');
+      await this.#primaryTargetClient.send('Page.close');
    } else {
      await connection.send('Target.closeTarget', {
-        targetId: this.#target._targetId,
+        targetId: this.#primaryTarget._targetId,
      });
-      await this.#target._isClosedDeferred.valueOrThrow();
+      await this.#tabTarget._isClosedDeferred.valueOrThrow();
    }
  }