From 8b8770c004ba842496e0ca4845642fe82a211051 Mon Sep 17 00:00:00 2001
From: jrandolf <101637635+jrandolf@users.noreply.github.com>
Date: Mon, 7 Aug 2023 14:44:12 +0200
Subject: [PATCH] fix: destroy puppeteer utility on context destruction
 (#10672)

---
 .../puppeteer-core/src/common/bidi/Browser.ts |  1 +
 .../puppeteer-core/src/common/bidi/Realm.ts   | 17 ++++++++
 .../puppeteer-core/src/common/bidi/Sandbox.ts |  9 ++--
 test/TestExpectations.json                    | 42 +++++++++++++++++++
 4 files changed, 63 insertions(+), 6 deletions(-)

diff --git a/packages/puppeteer-core/src/common/bidi/Browser.ts b/packages/puppeteer-core/src/common/bidi/Browser.ts
index e869ce524ef98..b3c4bcec7a6a5 100644
--- a/packages/puppeteer-core/src/common/bidi/Browser.ts
+++ b/packages/puppeteer-core/src/common/bidi/Browser.ts
@@ -54,6 +54,7 @@ export class Browser extends BrowserBase {
     'browsingContext',
     'network',
     'log',
+    'script',
   ];
   static readonly subscribeCdpEvents: Bidi.Cdp.EventNames[] = [
     // Coverage
diff --git a/packages/puppeteer-core/src/common/bidi/Realm.ts b/packages/puppeteer-core/src/common/bidi/Realm.ts
index 8c42d552511dc..707e5e9dda2dd 100644
--- a/packages/puppeteer-core/src/common/bidi/Realm.ts
+++ b/packages/puppeteer-core/src/common/bidi/Realm.ts
@@ -7,6 +7,7 @@ import {scriptInjector} from '../ScriptInjector.js';
 import {EvaluateFunc, HandleFor} from '../types.js';
 import {
   PuppeteerURL,
+  debugError,
   getSourcePuppeteerURLIfAvailable,
   isString,
 } from '../util.js';
@@ -46,6 +47,22 @@ export class Realm extends EventEmitter {
 
   setFrame(frame: Frame): void {
     this.#frame = frame;
+
+    // TODO(jrandolf): We should try to find a less brute-force way of doing
+    // this.
+    this.connection.on(
+      Bidi.ChromiumBidi.Script.EventNames.RealmDestroyed,
+      async () => {
+        const promise = this.internalPuppeteerUtil;
+        this.internalPuppeteerUtil = undefined;
+        try {
+          const util = await promise;
+          await util?.dispose();
+        } catch (error) {
+          debugError(error);
+        }
+      }
+    );
   }
 
   protected internalPuppeteerUtil?: Promise<JSHandle<PuppeteerUtil>>;
diff --git a/packages/puppeteer-core/src/common/bidi/Sandbox.ts b/packages/puppeteer-core/src/common/bidi/Sandbox.ts
index 3e8f871e9cdba..e83def6d23e68 100644
--- a/packages/puppeteer-core/src/common/bidi/Sandbox.ts
+++ b/packages/puppeteer-core/src/common/bidi/Sandbox.ts
@@ -60,7 +60,6 @@ export interface SandboxChart {
  * @internal
  */
 export class Sandbox implements RealmBase {
-  #document?: ElementHandle<Document>;
   #realm: Realm;
 
   #timeoutSettings: TimeoutSettings;
@@ -76,13 +75,11 @@ export class Sandbox implements RealmBase {
   }
 
   async document(): Promise<ElementHandle<Document>> {
-    if (this.#document) {
-      return this.#document;
-    }
-    this.#document = await this.#realm.evaluateHandle(() => {
+    // TODO(jrandolf): We should try to cache this because we need to dispose
+    // this when it's unused.
+    return await this.#realm.evaluateHandle(() => {
       return document;
     });
-    return this.#document;
   }
 
   async $<Selector extends string>(
diff --git a/test/TestExpectations.json b/test/TestExpectations.json
index 5daca7ca20c87..1d9034a1dfe69 100644
--- a/test/TestExpectations.json
+++ b/test/TestExpectations.json
@@ -341,12 +341,48 @@
     "parameters": ["webDriverBiDi"],
     "expectations": ["PASS"]
   },
+  {
+    "testIdPattern": "[browser.spec] Browser specs Browser.target should return browser target",
+    "platforms": ["darwin", "linux", "win32"],
+    "parameters": ["webDriverBiDi"],
+    "expectations": ["PASS"]
+  },
+  {
+    "testIdPattern": "[browsercontext.spec] BrowserContext should close all belonging targets once closing context",
+    "platforms": ["darwin", "linux", "win32"],
+    "parameters": ["webDriverBiDi"],
+    "expectations": ["PASS"]
+  },
+  {
+    "testIdPattern": "[browsercontext.spec] BrowserContext should create new incognito context",
+    "platforms": ["darwin", "linux", "win32"],
+    "parameters": ["webDriverBiDi"],
+    "expectations": ["PASS"]
+  },
   {
     "testIdPattern": "[browsercontext.spec] BrowserContext should have default context",
     "platforms": ["darwin", "linux", "win32"],
     "parameters": ["webDriverBiDi"],
     "expectations": ["PASS"]
   },
+  {
+    "testIdPattern": "[browsercontext.spec] BrowserContext should timeout waiting for a non-existent target",
+    "platforms": ["darwin", "linux", "win32"],
+    "parameters": ["webDriverBiDi"],
+    "expectations": ["PASS"]
+  },
+  {
+    "testIdPattern": "[browsercontext.spec] BrowserContext should wait for a target",
+    "platforms": ["darwin", "linux", "win32"],
+    "parameters": ["webDriverBiDi"],
+    "expectations": ["PASS"]
+  },
+  {
+    "testIdPattern": "[browsercontext.spec] BrowserContext window.open should use parent tab context",
+    "platforms": ["darwin", "linux", "win32"],
+    "parameters": ["webDriverBiDi"],
+    "expectations": ["PASS"]
+  },
   {
     "testIdPattern": "[CDPSession.spec] Target.createCDPSession *",
     "platforms": ["darwin", "linux", "win32"],
@@ -3899,6 +3935,12 @@
     "parameters": ["cdp", "chrome", "new-headless"],
     "expectations": ["FAIL", "PASS"]
   },
+  {
+    "testIdPattern": "[click.spec] Page.click should click the button after navigation",
+    "platforms": ["darwin"],
+    "parameters": ["chrome", "headless", "webDriverBiDi"],
+    "expectations": ["PASS"]
+  },
   {
     "testIdPattern": "[click.spec] Page.click should click the button with fixed position inside an iframe",
     "platforms": ["darwin", "linux", "win32"],