browser-cli/extension/src/classes/NativeConnection.ts

import { webExtApi as api } from '../browser-api';
/**
 * Native-messaging port lifecycle: connect/keepalive/reconnect plus the inbound
 * message router that hands commands to the CommandRegistry.
 */

import { getErrorMessage, getProfileAlias } from '../core';
import type { CommandRegistry } from './CommandRegistry';
import type { SessionCommands } from '../commands/session';
import type { ControlMessage, ResponseMessage, IncomingMessage, PageRequest, DispatchArgs, Serializable, RuntimePort } from '../types';

const NATIVE_HOST = "com.browsercli.host";
const DEBUG_LOG = false;

function debugLog(...args: Serializable[]) {
  if (DEBUG_LOG) console.log("[browser-cli]", ...args);
}

export class NativeConnection {
  private port: RuntimePort | null = null;
  private keepaliveEnabled = true;

  constructor(
    private readonly registry: CommandRegistry,
    private readonly session: SessionCommands,
  ) {}

  /** Registers all runtime listeners and opens the initial connection. */
  start() {
    api.runtime.onInstalled.addListener(() => this.connect());
    api.runtime.onStartup.addListener(() => this.connect());
    api.runtime.onSuspend.addListener(() => {
      this.disconnectPort({ sendBye: true });
    });
    api.windows.onCreated.addListener(() => {
      this.keepaliveEnabled = true;
      if (!this.port) this.connect();
    });
    api.windows.onRemoved.addListener(async () => {
      const windows = await api.windows.getAll({});
      if (windows.length > 0) return;

      this.keepaliveEnabled = false;
      this.disconnectPort({ sendBye: true });
    });

    // Reconnect poll — wakes the worker to re-establish the native port if it
    // dropped. 0.5 min is Chrome's minimum alarm period; lower values (e.g. 0.4)
    // are silently clamped and log a warning, so we set it explicitly.
    api.alarms.create("keepalive", { periodInMinutes: 0.5 });
    api.alarms.onAlarm.addListener((alarm) => {
      if (alarm.name === "keepalive") {
        if (!this.port && this.keepaliveEnabled) this.connect();
      }
    });
  }

  private sendControlMessage(targetPort: RuntimePort | null, message: ControlMessage) {
    if (!targetPort) return;
    try {
      targetPort.postMessage(message);
    } catch (e) {
      console.warn("[browser-cli] Failed to send control message:", e);
    }
  }

  private sendResponse(targetPort: RuntimePort | null, message: ResponseMessage) {
    if (!targetPort) return;
    try {
      targetPort.postMessage(message);
    } catch (e) {
      console.warn("[browser-cli] Failed to send response:", e);
    }
  }

  private disconnectPort({ sendBye = false }: { sendBye?: boolean } = {}) {
    const currentPort = this.port;
    if (!currentPort) return;

    if (sendBye) this.sendControlMessage(currentPort, { type: "bye" });

    if (this.port === currentPort) this.port = null;

    try {
      currentPort.disconnect();
    } catch (e) {
      console.warn("[browser-cli] Failed to disconnect native port:", e);
    }
  }

  private async connect() {
    if (this.port || !this.keepaliveEnabled) return;
    try {
      const nativePort = api.runtime.connectNative(NATIVE_HOST);
      this.port = nativePort;
      nativePort.onMessage.addListener((msg: IncomingMessage) => this.onMessage(msg));
      nativePort.onDisconnect.addListener(() => {
        if (this.port === nativePort) this.port = null;
        const err = api.runtime.lastError;
        if (err) console.warn("[browser-cli] Native host disconnected:", err.message);
      });
      // Send hello so native host knows which profile/alias this is
      const alias = await getProfileAlias();
      nativePort.postMessage({ type: "hello", alias });
      debugLog("Connected to native host as profile:", alias);
    } catch (e) {
      this.port = null;
      console.error("[browser-cli] Failed to connect:", e);
    }
  }

  private async onMessage(msg: IncomingMessage) {
    const { id, command, args } = msg;
    if (!id || !command) return;

    // Capture the port that delivered this message. dispatch() is awaited, and
    // during that await onDisconnect/rename can null out this.port — so we reply
    // on the captured port and bail if it (or this.port) is already gone.
    const replyPort = this.port;

    debugLog("←", command, args);

    let data: Serializable, error: string | undefined;
    try {
      const { __page, __background, ...commandArgs } = args || {};
      data = await this.registry.dispatch(command, commandArgs as DispatchArgs, {
        background: Boolean(__background),
        page: __page as PageRequest | undefined,
      });
    } catch (e) {
      error = getErrorMessage(e);
    }

    if (this.port !== replyPort) {
      console.warn("[browser-cli] Port changed before reply, dropping:", command);
      return;
    }

    if (error !== undefined) {
      debugLog("→ ERROR", command, error);
      this.sendResponse(replyPort, { id, success: false, error });
    } else {
      debugLog("→", command, data);
      this.sendResponse(replyPort, { id, success: true, data });
    }

    if (command === "clients.rename_profile" && error === undefined) {
      this.disconnectPort({ sendBye: true });
      this.keepaliveEnabled = true;
      await this.connect();
    }
  }
}