webext-typed-messages/src/index.ts

Better WebExtension sendMessage API

/**
 * Definitions for your custom runtime messages.
 *
 * Should be a map of message types to message handler functions.
 *
 * @example
 *
 * ```ts
 * declare module "webext-typed-messages" {
 *   interface RuntimeMessages {
 *     hello: (name: string) => string;
 *   }
 * }
 * ```
 */
export interface RuntimeMessages {}

/**
 * Definitions for your custom tab messages.
 *
 * Should be a map of message types to message handler functions.
 *
 * @example
 *
 * ```ts
 * declare module "webext-typed-messages" {
 *   interface TabMessages {
 *     add: (a: number, b: number) => Promise<number>;
 *   }
 * }
 * ```
 */
export interface TabMessages {}

/**
 * Parameters for a message handler.
 *
 * @template T - Message handler map
 * @template M - Message type
 */
type messageParams<T extends {}, M extends keyof T> = T[M] extends (
  ...params: infer P
) => any
  ? P
  : never;

/**
 * Return type of a message handler.
 *
 * @template T - Message handler map
 * @template M - Message type
 */
type messageResult<T extends {}, M extends keyof T> = T[M] extends (
  ...params: any[]
) => infer R
  ? Awaited<R>
  : undefined;

/**
 * Sends a runtime message to your registered handlers.
 *
 * Call {@link addRuntimeMessageListener} to register your handlers.
 *
 * @see {@link chrome.runtime.sendMessage}
 */
export function sendRuntimeMessage<M extends keyof RuntimeMessages>(
  type: M,
  ...params: messageParams<RuntimeMessages, M>
): Promise<messageResult<RuntimeMessages, M> | undefined> {
  return new Promise((resolve, reject) => {
    chrome.runtime.sendMessage({ type, params }, (response) => {
      if (chrome.runtime.lastError) {
        reject(new Error(chrome.runtime.lastError.message));
      } else {
        resolve(response);
      }
    });
  });
}

/**
 * Sends a message to the content script in the specified tab.
 *
 * The content script should call {@link addTabMessageListener} to receive the message.
 *
 * @see {@link chrome.tabs.sendMessage}
 */
export function sendTabMessage<M extends keyof TabMessages>(
  tabId: number,
  type: M,
  ...params: messageParams<TabMessages, M>
): Promise<messageResult<TabMessages, M> | undefined> {
  return new Promise((resolve, reject) => {
    chrome.tabs.sendMessage(tabId, { type, params }, (response) => {
      if (chrome.runtime.lastError) {
        reject(new Error(chrome.runtime.lastError.message));
      } else {
        resolve(response);
      }
    });
  });
}

/**
 * Object of functions for handling messages of each type.
 *
 * @template T - Message handler map
 */
type MessageListenerMap<T extends {}> = {
  [M in keyof T]: (
    this: void,
    ...params: messageParams<T, M>
  ) => messageResult<T, M> | Promise<messageResult<T, M>>;
};

/**
 * Registers listeners for runtime messages.
 *
 * Should be called in the background script.
 *
 * @see {@link chrome.runtime.onMessage}
 */
export function addRuntimeMessageListener(
  listenerMap: MessageListenerMap<RuntimeMessages>
) {
  const messageListener = getMessageListener(listenerMap);
  chrome.runtime.onMessage.addListener(messageListener);
  return function removeMessageListener() {
    chrome.runtime.onMessage.removeListener(messageListener);
  };
}

/**
 * Registers listeners for tab messages.
 *
 * Should be called in the content script.
 *
 * @see {@link chrome.runtime.onMessage}
 */
export function addTabMessageListener(
  listenerMap: MessageListenerMap<TabMessages>
) {
  const messageListener = getMessageListener(listenerMap);
  chrome.runtime.onMessage.addListener(messageListener);
  return function removeMessageListener() {
    chrome.runtime.onMessage.removeListener(messageListener);
  };
}

function getMessageListener<T extends {}>(listenerMap: MessageListenerMap<T>) {
  return function messageListener(
    message: any,
    sender: chrome.runtime.MessageSender,
    sendResponse: (response: any) => void
  ) {
    handleMessage(listenerMap, message, sendResponse);
    return true;
  };
}

async function handleMessage(
  listenerMap: Record<string, (...params: any[]) => any>,
  message: any,
  sendResponse: (response: any) => void
) {
  const listener = listenerMap[message.type];
  try {
    const result = await listener?.(...message.params);
    sendResponse(result);
  } catch (error) {
    sendResponse(undefined);
    throw error;
  }
}