From 3e9855c5a57b3c72745407906d081c0e57f042e9 Mon Sep 17 00:00:00 2001
From: Maksim Sadym <69349599+sadym-chromium@users.noreply.github.com>
Date: Thu, 4 Jul 2024 10:16:48 +0200
Subject: [PATCH] fix: restore OOPiF state (#2381)

Addressing #2362.

NodeJS runner nor chromedriver don't support re-connecting, so tested via puppeteer.
---
 src/bidiMapper/modules/cdp/CdpTarget.ts       | 21 +++++---
 .../modules/context/BrowsingContextImpl.ts    | 49 ++++++++++++++++---
 2 files changed, 56 insertions(+), 14 deletions(-)

diff --git a/src/bidiMapper/modules/cdp/CdpTarget.ts b/src/bidiMapper/modules/cdp/CdpTarget.ts
index 684717f7f3..b6f2bf3e78 100644
--- a/src/bidiMapper/modules/cdp/CdpTarget.ts
+++ b/src/bidiMapper/modules/cdp/CdpTarget.ts
@@ -191,12 +191,21 @@ export class CdpTarget {
 
   #restoreFrameTreeState(frameTree: Protocol.Page.FrameTree) {
     const frame = frameTree.frame;
-    if (
-      this.#browsingContextStorage.findContext(frame.id) === undefined &&
-      frame.parentId !== undefined
-    ) {
-      // Can restore only not yet known nested frames. The top-level frame is created
-      // when the target is attached.
+    const maybeContext = this.#browsingContextStorage.findContext(frame.id);
+    if (maybeContext !== undefined) {
+      // Restoring parent of already known browsing context. This means the target is
+      // OOPiF and the BiDi session was connected to already existing browser instance.
+      if (
+        maybeContext.parentId === null &&
+        frame.parentId !== null &&
+        frame.parentId !== undefined
+      ) {
+        maybeContext.parentId = frame.parentId;
+      }
+    }
+    if (maybeContext === undefined && frame.parentId !== undefined) {
+      // Restore not yet known nested frames. The top-level frame is created when the
+      // target is attached.
       const parentBrowsingContext = this.#browsingContextStorage.getContext(
         frame.parentId
       );
diff --git a/src/bidiMapper/modules/context/BrowsingContextImpl.ts b/src/bidiMapper/modules/context/BrowsingContextImpl.ts
index de48c87df9..b4f8e938ef 100644
--- a/src/bidiMapper/modules/context/BrowsingContextImpl.ts
+++ b/src/bidiMapper/modules/context/BrowsingContextImpl.ts
@@ -54,7 +54,7 @@ export class BrowsingContextImpl {
    * The ID of the parent browsing context.
    * If null, this is a top-level context.
    */
-  readonly #parentId: BrowsingContext.BrowsingContext | null;
+  #parentId: BrowsingContext.BrowsingContext | null = null;
 
   /** Direct children browsing contexts. */
   readonly #children = new Set<BrowsingContext.BrowsingContext>();
@@ -154,13 +154,30 @@ export class BrowsingContextImpl {
       context.parent!.addChild(context.id);
     }
 
-    eventManager.registerEvent(
-      {
-        type: 'event',
-        method: ChromiumBidi.BrowsingContext.EventNames.ContextCreated,
-        params: context.serializeToBidiValue(),
-      },
-      context.id
+    // Hold on the `contextCreated` event until the target is unblocked. This is required,
+    // as the parent of the context can be set later in case of reconnecting to an
+    // existing browser instance + OOPiF.
+    eventManager.registerPromiseEvent(
+      context.targetUnblockedOrThrow().then(
+        () => {
+          return {
+            kind: 'success',
+            value: {
+              type: 'event',
+              method: ChromiumBidi.BrowsingContext.EventNames.ContextCreated,
+              params: context.serializeToBidiValue(),
+            },
+          };
+        },
+        (error) => {
+          return {
+            kind: 'error',
+            error,
+          };
+        }
+      ),
+      context.id,
+      ChromiumBidi.BrowsingContext.EventNames.ContextCreated
     );
 
     return context;
@@ -226,6 +243,22 @@ export class BrowsingContextImpl {
     return this.#parentId;
   }
 
+  /** Sets the parent context ID and updates parent's children. */
+  set parentId(parentId: BrowsingContext.BrowsingContext | null) {
+    if (this.#parentId !== null) {
+      this.#logger?.(LogType.debugError, 'Parent context already set');
+      // Cannot do anything except logging, as throwing will stop event processing. So
+      // just return,
+      return;
+    }
+
+    this.#parentId = parentId;
+
+    if (!this.isTopLevelContext()) {
+      this.parent!.addChild(this.id);
+    }
+  }
+
   /** Returns the parent context. */
   get parent(): BrowsingContextImpl | null {
     if (this.parentId === null) {