refactor: update helpers

packages/puppeteer-core/src/common/Connection.ts

@@ -308,7 +308,6 @@ export class Connection extends EventEmitter {
     const object = JSON.parse(message);
     if (object.method === 'Target.attachedToTarget') {
       const sessionId = object.params.sessionId;
-      const parentSession = this.#sessions.get(object.sessionId);
       const session = new CDPSessionImpl(
         this,
         object.params.targetInfo.type,
@@ -317,6 +316,7 @@ export class Connection extends EventEmitter {
       );
       this.#sessions.set(sessionId, session);
       this.emit('sessionattached', session);
+      const parentSession = this.#sessions.get(object.sessionId);
       if (parentSession) {
         parentSession.emit('sessionattached', session);
       }
@@ -535,16 +535,21 @@ export class CDPSessionImpl extends CDPSession {
   }
 
   /**
+   * Sets the CDPTarget associated with the session instance.
+   *
    * @internal
    */
   _setTarget(target: CDPTarget): void {
     this.#target = target;
   }
 
   /**
+   * Gets the CDPTarget associated with the session instance.
+   *
    * @internal
    */
-  _target(): CDPTarget | undefined {
+  _target(): CDPTarget {
+    assert(this.#target, 'Target must exist');
     return this.#target;
   }
 

packages/puppeteer-core/src/common/FirefoxTargetManager.ts

@@ -108,13 +108,6 @@ export class FirefoxTargetManager
     this.setupAttachmentListeners(this.#connection);
   }
 
-  /**
-   * @internal
-   */
-  _tabTargetBySession(_session?: CDPSession): CDPSession | undefined {
-    return undefined;
-  }
-
   addTargetInterceptor(
     client: CDPSession | Connection,
     interceptor: TargetInterceptor

packages/puppeteer-core/src/common/Frame.ts

@@ -85,11 +85,16 @@ export class Frame extends BaseFrame {
     this.updateClient(client);
 
     this.on(FrameEmittedEvents.FrameSwappedByActivation, () => {
+      // Emulate loading process for swapped frames.
       this._onLoadingStarted();
       this._onLoadingStopped();
     });
   }
 
+  /**
+   * Updates the frame ID with the new ID. This happens when the main frame is
+   * replaced by a different frame.
+   */
   updateId(id: string): void {
     this._id = id;
   }

packages/puppeteer-core/src/common/FrameManager.ts

@@ -61,6 +61,8 @@ export const FrameManagerEmittedEvents = {
   ),
 };
 
+const TIME_FOR_WAITING_FOR_SWAP = 100; // ms.
+
 /**
  * A frame manager manages the frames for a given {@link Page | page}.
  *
@@ -114,41 +116,62 @@ export class FrameManager extends EventEmitter {
     this.#networkManager = new NetworkManager(client, ignoreHTTPSErrors, this);
     this.#timeoutSettings = timeoutSettings;
     this.setupEventListeners(this.#client);
-    client.once(CDPSessionEmittedEvents.Disconnected, async () => {
-      const mainFrame = this._frameTree.getMainFrame();
-      if (!mainFrame) {
-        return;
+    client.once(CDPSessionEmittedEvents.Disconnected, () => {
+      this.#onClientDisconnect().catch(debugError);
+    });
+  }
+
+  /**
+   * Called when the frame's client is disconnected. We don't know if the
+   * disconnect means that the frame is removed or if it will be replaced by a
+   * new frame. Therefore, we wait for a swap event.
+   */
+  async #onClientDisconnect() {
+    const mainFrame = this._frameTree.getMainFrame();
+    if (!mainFrame) {
+      return;
+    }
+    const swapped = Deferred.create<void>({
+      timeout: TIME_FOR_WAITING_FOR_SWAP,
+      message: 'Frame was not swapped',
+    });
+    mainFrame.once(FrameEmittedEvents.FrameSwappedByActivation, () => {
+      swapped.resolve();
+    });
+    try {
+      await swapped.valueOrThrow();
+      for (const child of mainFrame.childFrames()) {
+        this.#removeFramesRecursively(child);
       }
-      const swapped = Deferred.create<void>({
-        timeout: 100,
-        message: 'Frame was not swapped',
-      });
-      mainFrame.once(FrameEmittedEvents.FrameSwappedByActivation, () => {
-        swapped.resolve();
-      });
-      try {
-        await swapped.valueOrThrow();
-        for (const child of mainFrame.childFrames()) {
-          this.#removeFramesRecursively(child);
-        }
-      } catch (err) {
-        if (mainFrame) {
-          this.#removeFramesRecursively(mainFrame);
-        }
+    } catch (err) {
+      if (mainFrame) {
+        this.#removeFramesRecursively(mainFrame);
       }
-    });
+    }
   }
 
+  /**
+   * When the main frame is replaced by another main frame,
+   * we maintain the main frame object identity while updating
+   * its frame tree and ID.
+   */
   async swapFrameTree(client: CDPSession): Promise<void> {
     this.#onExecutionContextsCleared(this.#client);
 
     this.#client = client;
+    assert(
+      this.#client instanceof CDPSessionImpl,
+      'CDPSession is not an instance of CDPSessionImpl.'
+    );
     const frame = this._frameTree.getMainFrame();
     if (frame) {
       this._frameTree.removeFrame(frame);
-      frame.updateId((this.#client as CDPSessionImpl)._target()!._targetId);
+      frame.updateId(this.#client._target()._targetId);
     }
     this.setupEventListeners(client);
+    client.once(CDPSessionEmittedEvents.Disconnected, () => {
+      this.#onClientDisconnect().catch(debugError);
+    });
     await this.initialize(client);
     if (frame) {
       frame.emit(FrameEmittedEvents.FrameSwappedByActivation);

packages/puppeteer-core/src/common/LifecycleWatcher.ts

@@ -120,7 +120,7 @@ export class LifecycleWatcher {
       addEventListener(
         frame,
         FrameEmittedEvents.FrameSwappedByActivation,
-        this.#frameSwappedByActivation.bind(this)
+        this.#frameSwapped.bind(this)
       ),
       addEventListener(
         frame,
@@ -227,11 +227,6 @@ export class LifecycleWatcher {
     this.#checkLifecycleComplete();
   }
 
-  #frameSwappedByActivation(): void {
-    this.#swapped = true;
-    this.#checkLifecycleComplete();
-  }
-
   #checkLifecycleComplete(): void {
     // We expect navigation to commit.
     if (!checkLifecycle(this.#frame, this.#expectedLifecycle)) {

packages/puppeteer-core/src/common/Page.ts

@@ -313,8 +313,13 @@ export class CDPPage extends Page {
 
     this.#tabSession?.on('sessionswapped', async newSession => {
       this.#client = newSession;
-      this.#target = (this.#client as CDPSessionImpl)._target()!;
+      assert(
+        this.#client instanceof CDPSessionImpl,
+        'CDPSession is not instance of CDPSessionImpl'
+      );
+      this.#target = this.#client._target();
       assert(this.#target, 'Missing target on swap');
+      // TODO: swap the session for other members.
       await this.#frameManager.swapFrameTree(newSession);
       this.#setupEventListeners();
     });

packages/puppeteer-core/src/common/Target.ts

@@ -84,7 +84,9 @@ export class CDPTarget extends Target {
     this.#browserContext = browserContext;
     this._targetId = targetInfo.targetId;
     this.#sessionFactory = sessionFactory;
-    (this.#session as CDPSessionImpl | undefined)?._setTarget(this);
+    if (this.#session && this.#session instanceof CDPSessionImpl) {
+      this.#session._setTarget(this);
+    }
   }
 
   /**