chore: tracing over bidi+ (#10370)

Co-authored-by: Nikolay Vitkov <34244704+Lightning00Blade@users.noreply.github.com>

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

@@ -184,7 +184,25 @@ export class CDPPage extends Page {
       this.#timeoutSettings
     );
     this.#emulationManager = new EmulationManager(client);
-    this.#tracing = new Tracing(client);
+    this.#tracing = new Tracing({
+      read: opts => {
+        return this.#client.send('IO.read', opts);
+      },
+      close: opts => {
+        return this.#client.send('IO.close', opts);
+      },
+      start: opts => {
+        return client.send('Tracing.start', opts);
+      },
+      stop: async () => {
+        const deferred = Deferred.create();
+        this.#client.once('Tracing.tracingComplete', event => {
+          deferred.resolve(event);
+        });
+        await this.#client.send('Tracing.end');
+        return deferred.valueOrThrow() as Promise<Protocol.Tracing.TracingCompleteEvent>;
+      },
+    });
     this.#coverage = new Coverage(client);
     this.#screenshotTaskQueue = screenshotTaskQueue;
     this.#viewport = null;
@@ -1476,7 +1494,17 @@ export class CDPPage extends Page {
     }
 
     assert(result.stream, '`stream` is missing from `Page.printToPDF');
-    return getReadableFromProtocolStream(this.#client, result.stream);
+    return getReadableFromProtocolStream(
+      {
+        read: opts => {
+          return this.#client.send('IO.read', opts);
+        },
+        close: opts => {
+          return this.#client.send('IO.close', opts);
+        },
+      },
+      result.stream
+    );
   }
 
   override async pdf(options: PDFOptions = {}): Promise<Buffer> {

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

@@ -13,12 +13,15 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-import {assert} from '../util/assert.js';
-import {Deferred} from '../util/Deferred.js';
-import {isErrorLike} from '../util/ErrorLike.js';
+import Protocol from 'devtools-protocol';
 
-import {CDPSession} from './Connection.js';
-import {getReadableAsBuffer, getReadableFromProtocolStream} from './util.js';
+import {assert} from '../util/assert.js';
+
+import {
+  getReadableAsBuffer,
+  getReadableFromProtocolStream,
+  ProtocolReadable,
+} from './util.js';
 
 /**
  * @public
@@ -29,6 +32,14 @@ export interface TracingOptions {
   categories?: string[];
 }
 
+/**
+ * @internal
+ */
+export interface TracingSource extends ProtocolReadable {
+  start(opts: Protocol.Tracing.StartRequest): Promise<void>;
+  stop(): Promise<Protocol.Tracing.TracingCompleteEvent>;
+}
+
 /**
  * The Tracing class exposes the tracing audit interface.
  * @remarks
@@ -46,15 +57,15 @@ export interface TracingOptions {
  * @public
  */
 export class Tracing {
-  #client: CDPSession;
+  #source: TracingSource;
   #recording = false;
   #path?: string;
 
   /**
    * @internal
    */
-  constructor(client: CDPSession) {
-    this.#client = client;
+  constructor(source: TracingSource) {
+    this.#source = source;
   }
 
   /**
@@ -102,7 +113,7 @@ export class Tracing {
 
     this.#path = path;
     this.#recording = true;
-    await this.#client.send('Tracing.start', {
+    await this.#source.start({
       transferMode: 'ReturnAsStream',
       traceConfig: {
         excludedCategories,
@@ -116,25 +127,13 @@ export class Tracing {
    * @returns Promise which resolves to buffer with trace data.
    */
   async stop(): Promise<Buffer | undefined> {
-    const contentDeferred = Deferred.create<Buffer | undefined>();
-    this.#client.once('Tracing.tracingComplete', async event => {
-      try {
-        const readable = await getReadableFromProtocolStream(
-          this.#client,
-          event.stream
-        );
-        const buffer = await getReadableAsBuffer(readable, this.#path);
-        contentDeferred.resolve(buffer ?? undefined);
-      } catch (error) {
-        if (isErrorLike(error)) {
-          contentDeferred.reject(error);
-        } else {
-          contentDeferred.reject(new Error(`Unknown error: ${error}`));
-        }
-      }
-    });
-    await this.#client.send('Tracing.end');
+    const result = await this.#source.stop();
+    const readable = await getReadableFromProtocolStream(
+      this.#source,
+      result.stream!
+    );
+    const buffer = await getReadableAsBuffer(readable, this.#path);
     this.#recording = false;
-    return contentDeferred.valueOrThrow();
+    return buffer ?? undefined;
   }
 }

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

@@ -35,7 +35,13 @@ import {debugError} from './utils.js';
  * @internal
  */
 export class Browser extends BrowserBase {
-  static readonly subscribeModules = ['browsingContext', 'network', 'log'];
+  static readonly subscribeModules = [
+    'browsingContext',
+    'network',
+    'log',
+    'cdp',
+  ];
+
   #browserName = '';
   #browserVersion = '';
 
@@ -55,11 +61,16 @@ export class Browser extends BrowserBase {
       browserName = result.capabilities.browserName ?? '';
       browserVersion = result.capabilities.browserVersion ?? '';
     } catch (err) {
+      // Chrome does not support session.new.
       debugError(err);
     }
 
     await opts.connection.send('session.subscribe', {
-      events: Browser.subscribeModules as Bidi.Message.EventNames[],
+      events: (browserName.toLocaleLowerCase().includes('firefox')
+        ? Browser.subscribeModules.filter(module => {
+            return !['cdp'].includes(module);
+          })
+        : Browser.subscribeModules) as Bidi.Message.EventNames[],
     });
 
     return new Browser({

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

@@ -59,6 +59,7 @@ export class BrowsingContext extends EventEmitter {
   #timeoutSettings: TimeoutSettings;
   #id: string;
   #url = 'about:blank';
+  #cdpSessionId?: string;
 
   constructor(
     connection: Connection,
@@ -95,6 +96,10 @@ export class BrowsingContext extends EventEmitter {
     return this.#id;
   }
 
+  get cdpSessionId(): string | undefined {
+    return this.#cdpSessionId;
+  }
+
   async goto(
     url: string,
     options: {
@@ -285,14 +290,17 @@ export class BrowsingContext extends EventEmitter {
     method: T,
     params: ProtocolMapping.Commands[T]['paramsType'][0] = {}
   ): Promise<ProtocolMapping.Commands[T]['returnType']> {
-    const session = await this.connection.send('cdp.getSession', {
-      context: this.#id,
-    });
-    const sessionId = session.result.cdpSession;
+    if (!this.#cdpSessionId) {
+      const session = await this.connection.send('cdp.getSession', {
+        context: this.#id,
+      });
+      const sessionId = session.result.cdpSession;
+      this.#cdpSessionId = sessionId;
+    }
     const result = await this.connection.send('cdp.sendCommand', {
       cdpMethod: method,
       cdpParams: params,
-      cdpSession: sessionId,
+      cdpSession: this.#cdpSessionId,
     });
     return result.result;
   }

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

@@ -210,6 +210,16 @@ export class Connection extends EventEmitter {
       // `log.entryAdded` specific context
     } else if ('source' in event.params && event.params.source.context) {
       context = this.#browsingContexts.get(event.params.source.context);
+    } else if (event.method === 'cdp.eventReceived') {
+      // TODO: this is not a good solution and we need to find a better one.
+      // Perhaps we need to have a dedicated CDP event emitter or emulate
+      // the CDPSession interface with BiDi?.
+      const cdpSessionId = event.params.cdpSession;
+      for (const context of this.#browsingContexts.values()) {
+        if (context.cdpSessionId === cdpSessionId) {
+          context?.emit(event.params.cdpMethod, event.params.cdpParams);
+        }
+      }
     }
     context?.emit(event.method, event.params);
   }

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

@@ -17,6 +17,7 @@
 import type {Readable} from 'stream';
 
 import * as Bidi from 'chromium-bidi/lib/cjs/protocol/protocol.js';
+import Protocol from 'devtools-protocol';
 
 import {
   Page as PageBase,
@@ -36,6 +37,7 @@ import {NetworkManagerEmittedEvents} from '../NetworkManager.js';
 import {PDFOptions} from '../PDFOptions.js';
 import {Viewport} from '../PuppeteerViewport.js';
 import {TimeoutSettings} from '../TimeoutSettings.js';
+import {Tracing} from '../Tracing.js';
 import {EvaluateFunc, HandleFor} from '../types.js';
 import {
   debugError,
@@ -117,6 +119,7 @@ export class Page extends PageBase {
       },
     ],
   ]);
+  #tracing: Tracing;
 
   constructor(browserContext: BrowserContext, info: {context: string}) {
     super();
@@ -151,12 +154,38 @@ export class Page extends PageBase {
           .sendCDPCommand('Accessibility.getFullAXTree');
       },
     });
+
+    this.#tracing = new Tracing({
+      read: opts => {
+        return this.mainFrame().context().sendCDPCommand('IO.read', opts);
+      },
+      close: opts => {
+        return this.mainFrame().context().sendCDPCommand('IO.close', opts);
+      },
+      start: opts => {
+        return this.mainFrame().context().sendCDPCommand('Tracing.start', opts);
+      },
+      stop: async () => {
+        const deferred = Deferred.create();
+        this.mainFrame()
+          .context()
+          .once('Tracing.tracingComplete', event => {
+            deferred.resolve(event);
+          });
+        await this.mainFrame().context().sendCDPCommand('Tracing.end');
+        return deferred.valueOrThrow() as Promise<Protocol.Tracing.TracingCompleteEvent>;
+      },
+    });
   }
 
   override get accessibility(): Accessibility {
     return this.#accessibility;
   }
 
+  override get tracing(): Tracing {
+    return this.#tracing;
+  }
+
   override browser(): Browser {
     return this.#browserContext.browser();
   }

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

@@ -574,11 +574,19 @@ export async function getReadableAsBuffer(
   }
 }
 
+/**
+ * @internal
+ */
+export interface ProtocolReadable {
+  read(opts: {handle: string; size: number}): Promise<Protocol.IO.ReadResponse>;
+  close(opts: {handle: string}): Promise<void>;
+}
+
 /**
  * @internal
  */
 export async function getReadableFromProtocolStream(
-  client: CDPSession,
+  source: ProtocolReadable,
   handle: string
 ): Promise<Readable> {
   // TODO: Once Node 18 becomes the lowest supported version, we can migrate to
@@ -597,11 +605,11 @@ export async function getReadableFromProtocolStream(
       }
 
       try {
-        const response = await client.send('IO.read', {handle, size});
+        const response = await source.read({handle, size});
         this.push(response.data, response.base64Encoded ? 'base64' : undefined);
         if (response.eof) {
           eof = true;
-          await client.send('IO.close', {handle});
+          await source.close({handle});
           this.push(null);
         }
       } catch (error) {

test/TestExpectations.json

@@ -269,12 +269,6 @@
     "parameters": ["webDriverBiDi"],
     "expectations": ["PASS"]
   },
-  {
-    "testIdPattern": "[browser.spec] Browser specs Browser.version should return version",
-    "platforms": ["darwin", "linux", "win32"],
-    "parameters": ["firefox", "webDriverBiDi"],
-    "expectations": ["PASS"]
-  },
   {
     "testIdPattern": "[chromiumonly.spec] *",
     "platforms": ["darwin", "linux", "win32"],
@@ -857,12 +851,24 @@
     "parameters": ["cdp", "firefox"],
     "expectations": ["FAIL", "SKIP"]
   },
+  {
+    "testIdPattern": "[tracing.spec] *",
+    "platforms": ["darwin", "linux", "win32"],
+    "parameters": ["chrome", "webDriverBiDi"],
+    "expectations": ["PASS"]
+  },
   {
     "testIdPattern": "[tracing.spec] *",
     "platforms": ["darwin", "linux", "win32"],
     "parameters": ["cdp", "firefox"],
     "expectations": ["SKIP"]
   },
+  {
+    "testIdPattern": "[tracing.spec] Tracing should throw if tracing on two pages",
+    "platforms": ["darwin", "linux", "win32"],
+    "parameters": ["webDriverBiDi"],
+    "expectations": ["FAIL"]
+  },
   {
     "testIdPattern": "[waittask.spec] waittask specs Frame.waitForTimeout waits for the given timeout before resolving",
     "platforms": ["darwin", "linux", "win32"],
@@ -899,6 +905,12 @@
     "parameters": ["chrome", "webDriverBiDi"],
     "expectations": ["FAIL"]
   },
+  {
+    "testIdPattern": "[browser.spec] Browser specs Browser.version should return version",
+    "platforms": ["darwin", "linux", "win32"],
+    "parameters": ["firefox", "webDriverBiDi"],
+    "expectations": ["PASS"]
+  },
   {
     "testIdPattern": "[browsercontext.spec] BrowserContext should fire target events",
     "platforms": ["darwin", "linux", "win32"],
@@ -2439,7 +2451,7 @@
     "testIdPattern": "[queryhandler.spec] Query handler tests P selectors should work for ARIA selectors in multiple isolated worlds",
     "platforms": ["darwin", "linux", "win32"],
     "parameters": ["cdp", "firefox"],
-    "expectations": ["FAIL"]
+    "expectations": ["FAIL", "TIMEOUT"]
   },
   {
     "testIdPattern": "[requestinterception-experimental.spec] request interception \"after each\" hook in \"request interception\"",

test/src/tracing.spec.ts

@@ -40,9 +40,9 @@ describe('Tracing', function () {
       fs.unlinkSync(outputFile);
     }
   });
+
   it('should output a trace', async () => {
     const {server, page} = testState;
-
     await page.tracing.start({screenshots: true, path: outputFile});
     await page.goto(server.PREFIX + '/grid.html');
     await page.tracing.stop();