ChatApp/apps/desktop/electron/modules/audio-loopback.ts

// Native WASAPI process-loopback bridge. Wraps the
// @chatapp/audio-loopback-native napi-rs addon and forwards PCM chunks
// to the renderer over IPC.
//
// Why the addon over Chromium's built-in 'loopback' source: the OS
// process-loopback API supports EXCLUDE_TARGET_PROCESS_TREE, which lets
// us capture every render session *except* our own PID tree. That keeps
// LiveKit's call playback out of the outgoing share so peers don't hear
// themselves echoed back. Chromium's 'loopback' has no such filter.
//
// Windows-only. The addon's start_capture call rejects on macOS/Linux
// with a clear error string and the renderer falls through to the
// existing Chromium getUserMedia path (lib/screenAudio.ts).

import { app, BrowserWindow, ipcMain } from 'electron';
import { createRequire } from 'node:module';
import path from 'node:path';
import { fileURLToPath } from 'node:url';

import {
  CHANNELS,
  type AudioLoopbackChunk,
  type AudioLoopbackStartResult,
} from '../ipc-types';

interface NativeAudioLoopback {
  startCapture: (callback: (samples: Float32Array) => void) => number;
  /** INCLUDE_TARGET_PROCESS_TREE variant — capture only `pid`'s tree.
   *  Used for window-shares so we get just the picked app's audio. */
  startCaptureForPid?: (
    pid: number,
    callback: (samples: Float32Array) => void,
  ) => number;
  /** Look up the owning process id of a top-level window handle. */
  resolveWindowPid?: (hwnd: number) => number;
  stopCapture: (captureId: number) => void;
}

// __dirname in an ESM main bundle resolves to out/main after
// electron-vite builds. We need a CJS-style require to load the native
// addon — `import` would trigger ESM resolution which doesn't handle
// .node files cleanly across electron-vite's transform.
const __filenameSafe =
  typeof __filename === 'string' ? __filename : fileURLToPath(import.meta.url);
const __dirnameSafe = path.dirname(__filenameSafe);
const requireCjs = createRequire(__filenameSafe);

let cachedAddon: NativeAudioLoopback | null | undefined;

function resolveAddon(): NativeAudioLoopback | null {
  if (cachedAddon !== undefined) return cachedAddon;

  // Production: electron-builder copies the .node into
  // resources/native/audio-loopback.node (see extraResources in
  // package.json).
  // Dev: napi build emits the binary alongside the addon's package.json
  // at apps/desktop/native/audio-loopback/audio-loopback.<triple>.node,
  // and writes an index.js shim that auto-selects the right triple.
  // Loading the shim works in both layouts when the binary sits
  // adjacent to it; for the packaged single-file layout we require the
  // .node directly.
  const candidates: string[] = [];
  if (app.isPackaged && process.resourcesPath) {
    candidates.push(path.join(process.resourcesPath, 'native', 'audio-loopback.node'));
  }
  // Dev workspace layout — main bundle lives at out/main/main.js,
  // addon lives at native/audio-loopback/. Walk up two levels from the
  // bundle dir (out/main → out → apps/desktop) and into native/.
  candidates.push(
    path.join(__dirnameSafe, '..', '..', 'native', 'audio-loopback', 'index.js'),
    // Direct .node fallback in case the JS shim is missing (e.g. user
    // ran `cargo build` manually instead of `napi build`).
    path.join(
      __dirnameSafe,
      '..',
      '..',
      'native',
      'audio-loopback',
      `audio-loopback.${process.platform}-${process.arch}-msvc.node`,
    ),
  );

  for (const candidate of candidates) {
    try {
      const mod = requireCjs(candidate) as NativeAudioLoopback;
      if (
        typeof mod.startCapture === 'function' &&
        typeof mod.stopCapture === 'function'
      ) {
        cachedAddon = mod;
        return mod;
      }
    } catch {
      /* try next candidate */
    }
  }

  cachedAddon = null;
  return null;
}

/** Track active capture ids per BrowserWindow so we can tear them down
 *  if the renderer is destroyed mid-capture (renderer crash, window
 *  close during share). Without this the WASAPI thread would leak. */
const activeCaptures = new Map<number, Set<number>>();

function registerWindowCleanup(win: BrowserWindow, addon: NativeAudioLoopback): void {
  const wcId = win.webContents.id;
  if (activeCaptures.has(wcId)) return;
  activeCaptures.set(wcId, new Set());
  const cleanup = (): void => {
    const ids = activeCaptures.get(wcId);
    if (!ids) return;
    for (const id of ids) {
      try {
        addon.stopCapture(id);
      } catch {
        /* best effort */
      }
    }
    activeCaptures.delete(wcId);
  };
  win.webContents.once('destroyed', cleanup);
  win.once('closed', cleanup);
}

export function register(mainWindow: BrowserWindow): void {
  ipcMain.handle(
    CHANNELS.AUDIO_LOOPBACK_START,
    async (event): Promise<AudioLoopbackStartResult> => {
      const addon = resolveAddon();
      if (!addon) {
        throw new Error('audio-loopback native addon unavailable on this platform');
      }
      const wc = event.sender;
      // captureId is assigned synchronously by addon.startCapture below,
      // but the chunk callback needs to reference it — we forward-declare
      // via a closure-shared holder. The first chunk can only fire after
      // startCapture returns (the worker thread spawn happens inside it).
      let assignedId = 0;
      const cb = makeChunkCallback(wc, () => assignedId);
      assignedId = addon.startCapture(cb);
      registerWindowCleanup(mainWindow, addon);
      const set = activeCaptures.get(mainWindow.webContents.id);
      if (set) set.add(assignedId);
      return { captureId: assignedId };
    },
  );

  ipcMain.handle(
    CHANNELS.AUDIO_LOOPBACK_START_FOR_WINDOW,
    async (
      event,
      args: { hwnd: number },
    ): Promise<AudioLoopbackStartResult> => {
      const addon = resolveAddon();
      if (!addon) {
        throw new Error('audio-loopback native addon unavailable on this platform');
      }
      if (!addon.startCaptureForPid || !addon.resolveWindowPid) {
        throw new Error(
          'audio-loopback native addon is too old: missing startCaptureForPid / resolveWindowPid (rebuild with `pnpm build:native`)',
        );
      }
      if (!args || typeof args.hwnd !== 'number' || !Number.isFinite(args.hwnd)) {
        throw new Error('AUDIO_LOOPBACK_START_FOR_WINDOW: hwnd must be a finite number');
      }
      const pid = addon.resolveWindowPid(args.hwnd);
      const wc = event.sender;
      let assignedId = 0;
      const cb = makeChunkCallback(wc, () => assignedId);
      assignedId = addon.startCaptureForPid(pid, cb);
      registerWindowCleanup(mainWindow, addon);
      const set = activeCaptures.get(mainWindow.webContents.id);
      if (set) set.add(assignedId);
      return { captureId: assignedId };
    },
  );

  ipcMain.handle(
    CHANNELS.AUDIO_LOOPBACK_STOP,
    async (_event, captureId: number): Promise<void> => {
      const addon = resolveAddon();
      if (!addon) return;
      try {
        addon.stopCapture(captureId);
      } finally {
        const set = activeCaptures.get(mainWindow.webContents.id);
        if (set) set.delete(captureId);
      }
    },
  );
}

/** Shared chunk-forwarding callback. Both start variants produce the
 *  same wire format (interleaved f32 stereo @ 48kHz) so the renderer
 *  doesn't need to know which start path was used — the captureId
 *  routes the chunks. */
function makeChunkCallback(
  wc: Electron.WebContents,
  getCaptureId: () => number,
): (samples: Float32Array) => void {
  return (samples: Float32Array): void => {
    // The addon delivers each chunk on its WASAPI capture thread —
    // marshal to the renderer's webContents from the main loop. If the
    // webContents has been destroyed (window closed during a share)
    // silently drop; the cleanup hook will stop the capture.
    if (wc.isDestroyed()) return;
    const captureId = getCaptureId();
    const payload: AudioLoopbackChunk = { captureId, samples };
    wc.send(CHANNELS.AUDIO_LOOPBACK_CHUNK, payload);
  };
}