chore: extract BiDi context to allow emitting only to it (#9742)

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

@@ -103,13 +103,11 @@ export class ChromeTargetManager extends EventEmitter implements TargetManager {
     this.#connection.on('sessiondetached', this.#onSessionDetached);
     this.#setupAttachmentListeners(this.#connection);
 
-    // TODO: remove `as any` once the protocol definitions are updated with the
-    // next Chromium roll.
     this.#connection
       .send('Target.setDiscoverTargets', {
         discover: true,
         filter: [{type: 'tab', exclude: true}, {}],
-      } as any)
+      })
       .then(this.#storeExistingTargetsForInit)
       .catch(debugError);
   }

packages/puppeteer-core/src/common/bidi/BidiOverCDP.ts

@@ -1,3 +1,19 @@
+/**
+ * Copyright 2023 Google Inc. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
 import * as BidiMapper from 'chromium-bidi/lib/cjs/bidiMapper/bidiMapper.js';
 import * as Bidi from 'chromium-bidi/lib/cjs/protocol/protocol.js';
 import type {ProtocolMapping} from 'devtools-protocol/types/protocol-mapping.js';

packages/puppeteer-core/src/common/bidi/BrowserContext.ts

@@ -18,6 +18,7 @@ import {BrowserContext as BrowserContextBase} from '../../api/BrowserContext.js'
 import {Page as PageBase} from '../../api/Page.js';
 
 import {Connection} from './Connection.js';
+import {Context} from './Context.js';
 import {Page} from './Page.js';
 
 /**
@@ -32,10 +33,11 @@ export class BrowserContext extends BrowserContextBase {
   }
 
   override async newPage(): Promise<PageBase> {
-    const response = await this.#connection.send('browsingContext.create', {
+    const {result} = await this.#connection.send('browsingContext.create', {
       type: 'tab',
     });
-    return new Page(this.#connection, response.result.context);
+    const context = this.#connection.context(result.context) as Context;
+    return new Page(context);
   }
 
   override async close(): Promise<void> {}

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

@@ -22,6 +22,8 @@ import {debug} from '../Debug.js';
 import {ProtocolError} from '../Errors.js';
 import {EventEmitter} from '../EventEmitter.js';
 
+import {Context} from './Context.js';
+
 const debugProtocolSend = debug('puppeteer:webDriverBiDi:SEND ►');
 const debugProtocolReceive = debug('puppeteer:webDriverBiDi:RECV ◀');
 
@@ -78,6 +80,7 @@ export class Connection extends EventEmitter {
   #lastId = 0;
   #closed = false;
   #callbacks: Map<number, ConnectionCallback> = new Map();
+  #contexts: Map<string, Context> = new Map();
 
   constructor(transport: ConnectionTransport, delay = 0) {
     super();
@@ -92,6 +95,10 @@ export class Connection extends EventEmitter {
     return this.#closed;
   }
 
+  context(contextId: string): Context | null {
+    return this.#contexts.get(contextId) || null;
+  }
+
   send<T extends keyof Commands>(
     method: T,
     params: Commands[T]['params']
@@ -126,7 +133,8 @@ export class Connection extends EventEmitter {
     debugProtocolReceive(message);
     const object = JSON.parse(message) as
       | Bidi.Message.CommandResponse
-      | Bidi.EventResponse<string, unknown>;
+      | Bidi.Message.EventMessage;
+
     if ('id' in object) {
       const callback = this.#callbacks.get(object.id);
       // Callbacks could be all rejected if someone has called `.dispose()`.
@@ -137,10 +145,21 @@ export class Connection extends EventEmitter {
             createProtocolError(callback.error, callback.method, object)
           );
         } else {
+          if (callback.method === 'browsingContext.create') {
+            this.#contexts.set(
+              object.result.context,
+              new Context(this, object.result.context)
+            );
+          }
           callback.resolve(object);
         }
       }
     } else {
+      if ('source' in object.params && !!object.params.source.context) {
+        const context = this.#contexts.get(object.params.source.context);
+        context?.emit(object.method, object.params);
+      }
+
       this.emit(object.method, object.params);
     }
   }

packages/puppeteer-core/src/common/bidi/Context.ts

@@ -0,0 +1,140 @@
+/**
+ * Copyright 2023 Google Inc. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import * as Bidi from 'chromium-bidi/lib/cjs/protocol/protocol.js';
+
+import {stringifyFunction} from '../../util/Function.js';
+import {EventEmitter} from '../EventEmitter.js';
+import {EvaluateFunc, HandleFor} from '../types.js';
+import {isString} from '../util.js';
+
+import {Connection} from './Connection.js';
+import {JSHandle} from './JSHandle.js';
+import {BidiSerializer} from './Serializer.js';
+
+/**
+ * @internal
+ */
+export class Context extends EventEmitter {
+  #connection: Connection;
+  _contextId: string;
+
+  constructor(connection: Connection, contextId: string) {
+    super();
+    this.#connection = connection;
+    this._contextId = contextId;
+  }
+
+  get connection(): Connection {
+    return this.#connection;
+  }
+
+  get id(): string {
+    return this._contextId;
+  }
+
+  async evaluateHandle<
+    Params extends unknown[],
+    Func extends EvaluateFunc<Params> = EvaluateFunc<Params>
+  >(
+    pageFunction: Func | string,
+    ...args: Params
+  ): Promise<HandleFor<Awaited<ReturnType<Func>>>> {
+    return this.#evaluate(false, pageFunction, ...args);
+  }
+
+  async evaluate<
+    Params extends unknown[],
+    Func extends EvaluateFunc<Params> = EvaluateFunc<Params>
+  >(
+    pageFunction: Func | string,
+    ...args: Params
+  ): Promise<Awaited<ReturnType<Func>>> {
+    return this.#evaluate(true, pageFunction, ...args);
+  }
+
+  async #evaluate<
+    Params extends unknown[],
+    Func extends EvaluateFunc<Params> = EvaluateFunc<Params>
+  >(
+    returnByValue: true,
+    pageFunction: Func | string,
+    ...args: Params
+  ): Promise<Awaited<ReturnType<Func>>>;
+  async #evaluate<
+    Params extends unknown[],
+    Func extends EvaluateFunc<Params> = EvaluateFunc<Params>
+  >(
+    returnByValue: false,
+    pageFunction: Func | string,
+    ...args: Params
+  ): Promise<HandleFor<Awaited<ReturnType<Func>>>>;
+  async #evaluate<
+    Params extends unknown[],
+    Func extends EvaluateFunc<Params> = EvaluateFunc<Params>
+  >(
+    returnByValue: boolean,
+    pageFunction: Func | string,
+    ...args: Params
+  ): Promise<HandleFor<Awaited<ReturnType<Func>>> | Awaited<ReturnType<Func>>> {
+    let responsePromise;
+    const resultOwnership = returnByValue ? 'none' : 'root';
+    if (isString(pageFunction)) {
+      responsePromise = this.#connection.send('script.evaluate', {
+        expression: pageFunction,
+        target: {context: this._contextId},
+        resultOwnership,
+        awaitPromise: true,
+      });
+    } else {
+      responsePromise = this.#connection.send('script.callFunction', {
+        functionDeclaration: stringifyFunction(pageFunction),
+        arguments: await Promise.all(
+          args.map(arg => {
+            return BidiSerializer.serialize(arg, this);
+          })
+        ),
+        target: {context: this._contextId},
+        resultOwnership,
+        awaitPromise: true,
+      });
+    }
+
+    const {result} = await responsePromise;
+
+    if ('type' in result && result.type === 'exception') {
+      throw new Error(result.exceptionDetails.text);
+    }
+
+    return returnByValue
+      ? BidiSerializer.deserialize(result.result)
+      : getBidiHandle(this, result.result);
+  }
+}
+
+/**
+ * @internal
+ */
+export function getBidiHandle(
+  context: Context,
+  result: Bidi.CommonDataTypes.RemoteValue
+): JSHandle {
+  if ((result.type === 'node' || result.type === 'window') && context) {
+    // TODO: Implement ElementHandle
+    return new JSHandle(context, result);
+  }
+  return new JSHandle(context, result);
+}

packages/puppeteer-core/src/common/bidi/JSHandle.ts

@@ -21,7 +21,7 @@ import {JSHandle as BaseJSHandle} from '../../api/JSHandle.js';
 import {EvaluateFuncWith, HandleFor, HandleOr} from '../../common/types.js';
 
 import {Connection} from './Connection.js';
-import {Page} from './Page.js';
+import {Context} from './Context.js';
 import {BidiSerializer} from './Serializer.js';
 import {releaseReference} from './utils.js';
 
@@ -30,17 +30,17 @@ export class JSHandle<T = unknown> extends BaseJSHandle<T> {
   #context;
   #remoteValue;
 
-  constructor(context: Page, remoteValue: Bidi.CommonDataTypes.RemoteValue) {
+  constructor(context: Context, remoteValue: Bidi.CommonDataTypes.RemoteValue) {
     super();
     this.#context = context;
     this.#remoteValue = remoteValue;
   }
 
-  context(): Page {
+  context(): Context {
     return this.#context;
   }
 
-  get connecton(): Connection {
+  get connection(): Connection {
     return this.#context.connection;
   }
 
@@ -122,7 +122,7 @@ export class JSHandle<T = unknown> extends BaseJSHandle<T> {
     }
     this.#disposed = true;
     if ('handle' in this.#remoteValue) {
-      await releaseReference(this.connecton, this.#remoteValue);
+      await releaseReference(this.#context, this.#remoteValue);
     }
   }
 
@@ -153,7 +153,7 @@ export class JSHandle<T = unknown> extends BaseJSHandle<T> {
     return 'handle' in this.#remoteValue ? this.#remoteValue.handle : undefined;
   }
 
-  bidiObject(): Bidi.CommonDataTypes.RemoteValue {
+  remoteValue(): Bidi.CommonDataTypes.RemoteValue {
     return this.#remoteValue;
   }
 }

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

@@ -17,49 +17,46 @@
 import * as Bidi from 'chromium-bidi/lib/cjs/protocol/protocol.js';
 
 import {Page as PageBase, PageEmittedEvents} from '../../api/Page.js';
-import {stringifyFunction} from '../../util/Function.js';
 import {ConsoleMessage, ConsoleMessageLocation} from '../ConsoleMessage.js';
-import type {EvaluateFunc, HandleFor} from '../types.js';
-import {isString} from '../util.js';
+import {EvaluateFunc, HandleFor} from '../types.js';
 
 import {Connection} from './Connection.js';
-import {JSHandle} from './JSHandle.js';
+import {Context, getBidiHandle} from './Context.js';
 import {BidiSerializer} from './Serializer.js';
 
 /**
  * @internal
  */
 export class Page extends PageBase {
-  #connection: Connection;
+  #context: Context;
   #subscribedEvents = [
     'log.entryAdded',
   ] as Bidi.Session.SubscribeParameters['events'];
-  _contextId: string;
+  #boundOnLogEntryAdded = this.#onLogEntryAdded.bind(this);
 
-  constructor(connection: Connection, contextId: string) {
+  constructor(context: Context) {
     super();
-    this.#connection = connection;
-    this._contextId = contextId;
+    this.#context = context;
 
     // TODO: Investigate an implementation similar to CDPSession
     this.connection.send('session.subscribe', {
       events: this.#subscribedEvents,
-      contexts: [this._contextId],
+      contexts: [this.contextId],
     });
 
-    this.connection.on('log.entryAdded', this.#onLogEntryAdded.bind(this));
+    this.#context.on('log.entryAdded', this.#boundOnLogEntryAdded);
   }
 
   #onLogEntryAdded(event: Bidi.Log.LogEntry): void {
     if (isConsoleLogEntry(event)) {
       const args = event.args.map(arg => {
-        return getBidiHandle(this, arg);
+        return getBidiHandle(this.#context, arg);
       });
 
       const text = args
         .reduce((value, arg) => {
           const parsedValue = arg.isPrimitiveValue
-            ? BidiSerializer.deserialize(arg.bidiObject())
+            ? BidiSerializer.deserialize(arg.remoteValue())
             : arg.toString();
           return `${value} ${parsedValue}`;
         }, '')
@@ -88,20 +85,24 @@ export class Page extends PageBase {
   }
 
   override async close(): Promise<void> {
-    await this.#connection.send('browsingContext.close', {
-      context: this._contextId,
-    });
-
-    this.connection.send('session.unsubscribe', {
+    await this.connection.send('session.unsubscribe', {
       events: this.#subscribedEvents,
-      contexts: [this._contextId],
+      contexts: [this.contextId],
     });
 
-    this.connection.off('log.entryAdded', this.#onLogEntryAdded.bind(this));
+    await this.connection.send('browsingContext.close', {
+      context: this.contextId,
+    });
+
+    this.#context.off('log.entryAdded', this.#boundOnLogEntryAdded);
   }
 
   get connection(): Connection {
-    return this.#connection;
+    return this.#context.connection;
+  }
+
+  get contextId(): string {
+    return this.#context.id;
   }
 
   override async evaluateHandle<
@@ -111,7 +112,7 @@ export class Page extends PageBase {
     pageFunction: Func | string,
     ...args: Params
   ): Promise<HandleFor<Awaited<ReturnType<Func>>>> {
-    return this.#evaluate(false, pageFunction, ...args);
+    return this.#context.evaluateHandle(pageFunction, ...args);
   }
 
   override async evaluate<
@@ -121,83 +122,8 @@ export class Page extends PageBase {
     pageFunction: Func | string,
     ...args: Params
   ): Promise<Awaited<ReturnType<Func>>> {
-    return this.#evaluate(true, pageFunction, ...args);
+    return this.#context.evaluate(pageFunction, ...args);
   }
-
-  async #evaluate<
-    Params extends unknown[],
-    Func extends EvaluateFunc<Params> = EvaluateFunc<Params>
-  >(
-    returnByValue: true,
-    pageFunction: Func | string,
-    ...args: Params
-  ): Promise<Awaited<ReturnType<Func>>>;
-  async #evaluate<
-    Params extends unknown[],
-    Func extends EvaluateFunc<Params> = EvaluateFunc<Params>
-  >(
-    returnByValue: false,
-    pageFunction: Func | string,
-    ...args: Params
-  ): Promise<HandleFor<Awaited<ReturnType<Func>>>>;
-  async #evaluate<
-    Params extends unknown[],
-    Func extends EvaluateFunc<Params> = EvaluateFunc<Params>
-  >(
-    returnByValue: boolean,
-    pageFunction: Func | string,
-    ...args: Params
-  ): Promise<HandleFor<Awaited<ReturnType<Func>>> | Awaited<ReturnType<Func>>> {
-    let responsePromise;
-    const resultOwnership = returnByValue ? 'none' : 'root';
-    if (isString(pageFunction)) {
-      responsePromise = this.#connection.send('script.evaluate', {
-        expression: pageFunction,
-        target: {context: this._contextId},
-        resultOwnership,
-        awaitPromise: true,
-      });
-    } else {
-      responsePromise = this.#connection.send('script.callFunction', {
-        functionDeclaration: stringifyFunction(pageFunction),
-        arguments: await Promise.all(
-          args.map(arg => {
-            return BidiSerializer.serialize(arg, this);
-          })
-        ),
-        target: {context: this._contextId},
-        resultOwnership,
-        awaitPromise: true,
-      });
-    }
-
-    const {result} = await responsePromise;
-
-    if ('type' in result && result.type === 'exception') {
-      throw new Error(result.exceptionDetails.text);
-    }
-
-    return returnByValue
-      ? BidiSerializer.deserialize(result.result)
-      : getBidiHandle(this, result.result);
-  }
-}
-
-/**
- * @internal
- */
-export function getBidiHandle(
-  context: Page,
-  result: Bidi.CommonDataTypes.RemoteValue
-): JSHandle {
-  if (
-    (result.type === 'node' || result.type === 'window') &&
-    context._contextId
-  ) {
-    // TODO: Implement ElementHandle
-    return new JSHandle(context, result);
-  }
-  return new JSHandle(context, result);
 }
 
 function isConsoleLogEntry(

packages/puppeteer-core/src/common/bidi/Serializer.ts

@@ -1,9 +1,25 @@
+/**
+ * Copyright 2023 Google Inc. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
 import * as Bidi from 'chromium-bidi/lib/cjs/protocol/protocol.js';
 
 import {debugError, isDate, isPlainObject, isRegExp} from '../util.js';
 
+import {Context} from './Context.js';
 import {JSHandle} from './JSHandle.js';
-import {Page} from './Page.js';
 
 /**
  * @internal
@@ -130,7 +146,7 @@ export class BidiSerializer {
 
   static serialize(
     arg: unknown,
-    context: Page
+    context: Context
   ): Bidi.CommonDataTypes.LocalOrRemoteValue {
     // TODO: See use case of LazyArgs
     const objectHandle = arg && arg instanceof JSHandle ? arg : null;
@@ -143,7 +159,7 @@ export class BidiSerializer {
       if (objectHandle.disposed) {
         throw new Error('JSHandle is disposed!');
       }
-      return objectHandle.bidiObject();
+      return objectHandle.remoteValue();
     }
 
     return BidiSerializer.serializeRemoveValue(arg);

packages/puppeteer-core/src/common/bidi/utils.ts

@@ -18,7 +18,7 @@ import * as Bidi from 'chromium-bidi/lib/cjs/protocol/protocol.js';
 
 import {debug} from '../Debug.js';
 
-import {Connection} from './Connection.js';
+import {Context} from './Context.js';
 
 /**
  * @internal
@@ -28,15 +28,15 @@ export const debugError = debug('puppeteer:error');
  * @internal
  */
 export async function releaseReference(
-  client: Connection,
+  client: Context,
   remoteReference: Bidi.CommonDataTypes.RemoteReference
 ): Promise<void> {
   if (!remoteReference.handle) {
     return;
   }
-  await client
+  await client.connection
     .send('script.disown', {
-      target: {realm: '', context: ''}, // TODO: Populate
+      target: {context: client._contextId},
       handles: [remoteReference.handle],
     })
     .catch((error: any) => {