rea/src/browser/CdpConnection.ts

import WebSocket, { type RawData } from "ws";

import {
  AnalysisCancelledError,
  AnalysisError,
  BrowserObservationError,
  type BrowserObservationOperation,
} from "../domain/errors.js";

export interface CdpEvent {
  readonly method: string;
  readonly params: unknown;
  readonly sessionId?: string;
}

interface PendingCommand {
  readonly resolve: (value: unknown) => void;
  readonly reject: (error: AnalysisError) => void;
  readonly removeAbort: () => void;
}

/** Correlated JSON-RPC transport for one CDP target operation. */
export class CdpConnection {
  readonly #pending = new Map<number, PendingCommand>();
  readonly #listeners = new Set<(event: CdpEvent) => void>();
  readonly #disconnectListeners = new Set<() => void>();
  #nextId = 1;
  #closed = false;
  #protocolFailed = false;

  private constructor(
    private readonly socket: WebSocket,
    private readonly operation: BrowserObservationOperation,
  ) {
    socket.on("message", (data) => this.#receive(data));
    socket.on("close", () => this.#disconnect());
    socket.on("error", () => this.#disconnect());
  }

  /** Connect to one already-validated loopback CDP WebSocket. */
  static async connect(
    url: string,
    operation: BrowserObservationOperation,
    signal?: AbortSignal,
  ): Promise<CdpConnection> {
    if (signal?.aborted === true) throw new AnalysisCancelledError(operation);
    const socket = new WebSocket(url, {
      handshakeTimeout: 0,
      maxPayload: 0,
      perMessageDeflate: false,
    });
    await waitForOpen(socket, operation, signal);
    return new CdpConnection(socket, operation);
  }

  /** Subscribe to validated CDP event envelopes. */
  onEvent(listener: (event: CdpEvent) => void): () => void {
    this.#listeners.add(listener);
    return () => this.#listeners.delete(listener);
  }

  /** Subscribe to an unexpected transport loss while an operation is active. */
  onDisconnect(listener: () => void): () => void {
    if (this.#closed) {
      listener();
      return () => undefined;
    }
    this.#disconnectListeners.add(listener);
    return () => this.#disconnectListeners.delete(listener);
  }

  /** Execute one command, optionally within a flat target session. */
  async send(
    method: string,
    params: Readonly<Record<string, unknown>> = {},
    sessionId?: string,
    signal?: AbortSignal,
  ): Promise<unknown> {
    if (this.#protocolFailed)
      throw new BrowserObservationError(this.operation, "protocol_error");
    if (this.#closed || this.socket.readyState !== WebSocket.OPEN)
      throw new BrowserObservationError(this.operation, "disconnected");
    if (signal?.aborted === true)
      throw new AnalysisCancelledError(this.operation);
    const id = this.#nextId;
    this.#nextId += 1;
    return await new Promise((resolve, reject) => {
      const onAbort = (): void => {
        const pending = this.#pending.get(id);
        if (pending === undefined) return;
        pending.removeAbort();
        this.#pending.delete(id);
        reject(new AnalysisCancelledError(this.operation));
      };
      signal?.addEventListener("abort", onAbort, { once: true });
      this.#pending.set(id, {
        resolve,
        reject,
        removeAbort: () => signal?.removeEventListener("abort", onAbort),
      });
      this.socket.send(
        JSON.stringify({
          id,
          method,
          params,
          ...(sessionId === undefined ? {} : { sessionId }),
        }),
        (error) => {
          if (error === undefined || error === null) return;
          const pending = this.#pending.get(id);
          if (pending === undefined) return;
          this.#complete(id, pending);
          reject(
            new BrowserObservationError(this.operation, "disconnected", {
              cause: error,
            }),
          );
        },
      );
    });
  }

  /** Close only REA's socket; never close the browser or selected page. */
  async close(): Promise<void> {
    if (this.#closed) return;
    this.#closed = true;
    this.#failPending("disconnected");
    if (this.socket.readyState === WebSocket.CLOSED) return;
    await new Promise<void>((resolve) => {
      const timer = setTimeout(() => {
        this.socket.terminate();
        resolve();
      }, 1_000);
      this.socket.once("close", () => {
        clearTimeout(timer);
        resolve();
      });
      this.socket.close();
    });
  }

  #receive(data: RawData): void {
    let message: unknown;
    try {
      message = JSON.parse(rawText(data));
    } catch {
      this.#failPending("protocol_error");
      return;
    }
    if (!isRecord(message)) {
      this.#failPending("protocol_error");
      return;
    }
    if ("id" in message) {
      if (
        typeof message.id !== "number" ||
        !Number.isSafeInteger(message.id) ||
        message.id < 1
      ) {
        this.#failPending("protocol_error");
        return;
      }
      this.#receiveResponse(message, message.id);
      return;
    }
    if (
      typeof message.method !== "string" ||
      message.method.length === 0 ||
      ("params" in message && !isRecord(message.params)) ||
      ("sessionId" in message && typeof message.sessionId !== "string")
    ) {
      this.#failPending("protocol_error");
      return;
    }
    const event: CdpEvent = {
      method: message.method,
      params: "params" in message ? message.params : {},
      ...(typeof message.sessionId === "string"
        ? { sessionId: message.sessionId }
        : {}),
    };
    for (const listener of this.#listeners) listener(event);
  }

  #receiveResponse(message: Record<string, unknown>, id: number): void {
    const pending = this.#pending.get(id);
    if (pending === undefined) return;
    this.#complete(id, pending);
    if ("error" in message) {
      pending.reject(
        new BrowserObservationError(this.operation, "protocol_error"),
      );
      return;
    }
    pending.resolve("result" in message ? message.result : {});
  }

  #complete(id: number, pending: PendingCommand): void {
    pending.removeAbort();
    this.#pending.delete(id);
  }

  #failPending(reason: "disconnected" | "protocol_error"): void {
    if (reason === "disconnected") this.#closed = true;
    else this.#protocolFailed = true;
    for (const [id, pending] of this.#pending) {
      this.#complete(id, pending);
      pending.reject(new BrowserObservationError(this.operation, reason));
    }
  }

  #disconnect(): void {
    const wasClosed = this.#closed;
    this.#failPending("disconnected");
    if (wasClosed) return;
    for (const listener of this.#disconnectListeners) listener();
    this.#disconnectListeners.clear();
  }
}

const rawText = (data: RawData): string =>
  Array.isArray(data)
    ? Buffer.concat(data).toString("utf8")
    : Buffer.isBuffer(data)
      ? data.toString("utf8")
      : Buffer.from(data).toString("utf8");

const isRecord = (value: unknown): value is Record<string, unknown> =>
  typeof value === "object" && value !== null && !Array.isArray(value);

const waitForOpen = async (
  socket: WebSocket,
  operation: BrowserObservationOperation,
  signal?: AbortSignal,
): Promise<void> =>
  await new Promise((resolve, reject) => {
    const onOpen = (): void => {
      cleanup();
      resolve();
    };
    const onFailure = (cause: unknown): void => {
      cleanup();
      reject(
        signal?.aborted === true
          ? new AnalysisCancelledError(operation)
          : new BrowserObservationError(operation, "endpoint_unreachable", {
              cause,
            }),
      );
    };
    const onClose = (code: number, reason: Buffer): void => {
      onFailure(
        new Error(
          `CDP WebSocket closed before opening (${code}: ${reason.toString("utf8")})`,
        ),
      );
    };
    const onAbort = (): void => {
      terminateSocket(socket);
      onFailure(new AnalysisCancelledError(operation));
    };
    const cleanup = (): void => {
      socket.off("open", onOpen);
      socket.off("error", onFailure);
      socket.off("close", onClose);
      signal?.removeEventListener("abort", onAbort);
    };
    socket.once("open", onOpen);
    socket.once("error", onFailure);
    socket.once("close", onClose);
    signal?.addEventListener("abort", onAbort, { once: true });
    if (signal?.aborted === true) onAbort();
  });

const terminateSocket = (socket: WebSocket): void => {
  socket.once("error", () => undefined);
  socket.terminate();
};