chore: change puppeteer/devtools connection adapter (#602)

Rather then adapting on the transport layer, we adapt the puppeteer
`CDPSession`/`Connection` and make them look like a `CDPConnection`. The
class assumes that callers create a dedicated `CDPSession` to be used.

The `PuppeteerDevToolsConnection` installs a generic `'*'` event
listener on all child sessions to funnel all CDP events into DevTools
via `CDPConnectionObservers`. As pointed out in the code comment, we
don't need to recursively listen to `'sessionattached'` events on child
sessions: Nested `sessionattached` events are reported on all parent
sessions.

While not strictly necessary, `PuppeteerDevToolsConnection` also
uninstalls the CDP event listener when a session gets detached (modulo
the root puppeteer `CDPSession`).

Author: szuend

src/DevToolsConnectionAdapter.ts

@@ -4,38 +4,106 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import {ConnectionTransport as DevToolsConnectionTransport} from '../node_modules/chrome-devtools-frontend/mcp/mcp.js';
+import type {CDPConnection as devtools} from '../node_modules/chrome-devtools-frontend/mcp/mcp.js';
 
-import {type ConnectionTransport} from './third_party/index.js';
+import type * as puppeteer from './third_party/index.js';
+import {CDPSessionEvent} from './third_party/index.js';
 
 /**
- * Allows a puppeteer {@link ConnectionTransport} to act like a DevTools {@link Connection}.
+ * This class makes a puppeteer connection look like DevTools CDPConnection.
+ *
+ * Since we connect "root" DevTools targets to specific pages, we scope everything to a puppeteer CDP session.
+ *
+ * We don't have to recursively listen for 'sessionattached' as the "root" CDP session sees all child session attached
+ * events, regardless how deeply nested they are.
  */
-export class DevToolsConnectionAdapter extends DevToolsConnectionTransport {
-  #transport: ConnectionTransport | null;
-  #onDisconnect: ((arg0: string) => void) | null = null;
-
-  constructor(transport: ConnectionTransport) {
-    super();
-    this.#transport = transport;
-    this.#transport.onclose = () => this.#onDisconnect?.('');
-    this.#transport.onmessage = msg => this.onMessage?.(msg);
+export class PuppeteerDevToolsConnection implements devtools.CDPConnection {
+  readonly #connection: puppeteer.Connection;
+  readonly #observers = new Set<devtools.CDPConnectionObserver>();
+  readonly #sessionEventHandlers = new Map<
+    string,
+    puppeteer.Handler<unknown>
+  >();
+
+  constructor(session: puppeteer.CDPSession) {
+    this.#connection = session.connection()!;
+
+    session.on(
+      CDPSessionEvent.SessionAttached,
+      this.#startForwardingCdpEvents.bind(this),
+    );
+    session.on(
+      CDPSessionEvent.SessionDetached,
+      this.#stopForwardingCdpEvents.bind(this),
+    );
+
+    this.#startForwardingCdpEvents(session);
+  }
+
+  send<T extends devtools.Command>(
+    method: T,
+    params: devtools.CommandParams<T>,
+    sessionId: string | undefined,
+  ): Promise<{result: devtools.CommandResult<T>} | {error: devtools.CDPError}> {
+    if (sessionId === undefined) {
+      throw new Error(
+        'Attempting to send on the root session. This must not happen',
+      );
+    }
+    const session = this.#connection.session(sessionId);
+    if (!session) {
+      throw new Error('Unknown session ' + sessionId);
+    }
+    // Rolled protocol version between puppeteer and DevTools doesn't necessarily match
+    /* eslint-disable @typescript-eslint/no-explicit-any */
+    return session
+      .send(method as any, params)
+      .then(result => ({result}))
+      .catch(error => ({error})) as any;
+    /* eslint-enable @typescript-eslint/no-explicit-any */
+  }
+
+  observe(observer: devtools.CDPConnectionObserver): void {
+    this.#observers.add(observer);
   }
 
-  override setOnMessage(onMessage: (arg0: object | string) => void): void {
-    this.onMessage = onMessage;
+  unobserve(observer: devtools.CDPConnectionObserver): void {
+    this.#observers.delete(observer);
   }
 
-  override setOnDisconnect(onDisconnect: (arg0: string) => void): void {
-    this.#onDisconnect = onDisconnect;
+  #startForwardingCdpEvents(session: puppeteer.CDPSession): void {
+    const handler = this.#handleEvent.bind(
+      this,
+      session.id(),
+    ) as puppeteer.Handler<unknown>;
+    this.#sessionEventHandlers.set(session.id(), handler);
+    session.on('*', handler);
   }
 
-  override sendRawMessage(message: string): void {
-    this.#transport?.send(message);
+  #stopForwardingCdpEvents(session: puppeteer.CDPSession): void {
+    const handler = this.#sessionEventHandlers.get(session.id());
+    if (handler) {
+      session.off('*', handler);
+    }
   }
 
-  override async disconnect(): Promise<void> {
-    this.#transport?.close();
-    this.#transport = null;
+  #handleEvent(
+    sessionId: string,
+    type: string | symbol | number,
+    event: any, // eslint-disable-line @typescript-eslint/no-explicit-any
+  ): void {
+    if (
+      typeof type === 'string' &&
+      type !== CDPSessionEvent.SessionAttached &&
+      type !== CDPSessionEvent.SessionDetached
+    ) {
+      this.#observers.forEach(observer =>
+        observer.onEvent({
+          method: type as devtools.Event,
+          sessionId,
+          params: event,
+        }),
+      );
+    }
   }
 }

src/third_party/index.ts

@@ -20,7 +20,11 @@ export {
   type TextContent,
 } from '@modelcontextprotocol/sdk/types.js';
 export {z as zod} from 'zod';
-export {Locator, PredefinedNetworkConditions} from 'puppeteer-core';
+export {
+  Locator,
+  PredefinedNetworkConditions,
+  CDPSessionEvent,
+} from 'puppeteer-core';
 export {default as puppeteer} from 'puppeteer-core';
 export type * from 'puppeteer-core';
 export type {CdpPage} from 'puppeteer-core/internal/cdp/Page.js';