chore: change puppeteer/devtools connection adapter

Author: szuend

src/DevToolsConnectionAdapter.ts

@@ -4,38 +4,101 @@
  * 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';
 
 /**
- * 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('sessionattached', this.#startForwardingCdpEvents.bind(this));
+    session.on('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;
+  // puppeteer types around '*' event listeners are broken
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  #handleEvent(
+    sessionId: string,
+    type: string | symbol | number,
+    event: any,
+  ): void {
+    if (
+      typeof type === 'string' &&
+      type !== 'sessionattached' &&
+      type !== 'sessiondetached'
+    ) {
+      this.#observers.forEach(observer =>
+        observer.onEvent({
+          method: type as devtools.Event,
+          sessionId,
+          params: event,
+        }),
+      );
+    }
   }
 }