ChatApp/apps/desktop/native/audio-loopback/src/lib.rs

// Native system-audio capture addon for the ChatApp Electron desktop
// client. Mirrors the proven Tauri implementation
// (apps/desktop/src-tauri/src/screen_audio.rs in the legacy repo).
//
// Why a native addon at all when Electron already exposes a 'loopback'
// audio source through setDisplayMediaRequestHandler? Because Chromium's
// loopback captures the entire OS mixer including our own renderer's
// playback — peers in a video call hear themselves echoed back when the
// sharer ticks "system audio". Windows ships an EXCLUDE_TARGET_PROCESS_TREE
// process-loopback mode that captures every render session except the
// targeted PID's tree. We pass our own PID so the LiveKit playback never
// re-enters the outgoing share.
//
// Wire format is fixed at 48kHz interleaved f32 stereo. The renderer-
// side AudioWorklet (lib/loopbackAudio.ts) assumes that layout and feeds
// samples into a MediaStreamDestination so LiveKit publishes a plain
// ScreenShareAudio track.
//
// Windows-only for v1. macOS/Linux stubs return a clear napi::Error so
// the renderer can fall through to the existing getUserMedia path.

#![allow(clippy::needless_return)]

use napi::bindgen_prelude::{Float32Array, Result};
use napi::threadsafe_function::{
    ErrorStrategy, ThreadsafeFunction, ThreadsafeFunctionCallMode,
};
use napi::JsFunction;
use napi_derive::napi;
use std::collections::HashMap;
use std::sync::atomic::{AtomicBool, AtomicU32, Ordering};
use std::sync::{Arc, Mutex, OnceLock};
use std::thread;

// Output format we always deliver to the frontend. Pinning a single fixed
// format means the AudioWorklet never has to renegotiate — it just
// assumes interleaved f32 stereo at 48kHz. WASAPI mix format is usually
// already this on Windows 10+, so the resample/upmix branch inside
// process-loopback's autoconvert is rarely hit.
const OUTPUT_SAMPLE_RATE: u32 = 48_000;
const OUTPUT_CHANNELS: u16 = 2;

static NEXT_ID: AtomicU32 = AtomicU32::new(1);
static SESSIONS: OnceLock<Mutex<HashMap<u32, Session>>> = OnceLock::new();

fn sessions() -> &'static Mutex<HashMap<u32, Session>> {
    SESSIONS.get_or_init(|| Mutex::new(HashMap::new()))
}

struct Session {
    stop: Arc<AtomicBool>,
    handle: Option<thread::JoinHandle<()>>,
}

/// Start a process-loopback capture that excludes the current process
/// tree. The supplied JS callback is invoked from a background thread
/// with one argument: a Float32Array of interleaved f32 stereo samples
/// at 48kHz. Returns a numeric capture id that must be passed to
/// `stopCapture` when the share ends.
///
/// Always excludes `std::process::id()` (whole-OS-mixer-minus-self).
/// For "include only this app" use `start_capture_for_pid` instead.
#[napi]
pub fn start_capture(callback: JsFunction) -> Result<u32> {
    spawn_capture_session(callback, std::process::id(), false)
}

/// Start a process-loopback capture that INCLUDES the target PID's
/// process tree (and only that tree) — the WASAPI
/// INCLUDE_TARGET_PROCESS_TREE mode. Used for window-shares where we
/// want only the picked app's audio (Discord parity).
#[napi]
pub fn start_capture_for_pid(pid: u32, callback: JsFunction) -> Result<u32> {
    spawn_capture_session(callback, pid, true)
}

/// Resolve the owning process id of a top-level window handle. The
/// renderer derives `hwnd` from desktopCapturer's `window:<HWND>:0`
/// source ids and we hand that to `start_capture_for_pid`.
#[napi]
pub fn resolve_window_pid(hwnd: u32) -> Result<u32> {
    #[cfg(target_os = "windows")]
    {
        // Minimal FFI to user32!GetWindowThreadProcessId — pulling in a
        // full windows crate just for one call would balloon build
        // times. The function returns the thread id (we ignore it) and
        // writes the process id through the pointer.
        #[allow(non_snake_case)]
        extern "system" {
            fn GetWindowThreadProcessId(hWnd: usize, lpdwProcessId: *mut u32) -> u32;
        }
        let mut pid: u32 = 0;
        // SAFETY: GetWindowThreadProcessId tolerates an invalid HWND
        // (returns 0 thread id and leaves *lpdwProcessId untouched). We
        // detect the failure case by checking for pid == 0 below.
        let thread_id = unsafe { GetWindowThreadProcessId(hwnd as usize, &mut pid) };
        if thread_id == 0 || pid == 0 {
            return Err(napi::Error::from_reason(format!(
                "GetWindowThreadProcessId({hwnd}) failed — window may have closed"
            )));
        }
        Ok(pid)
    }

    #[cfg(not(target_os = "windows"))]
    {
        let _ = hwnd;
        Err(napi::Error::from_reason(
            "resolve_window_pid only supported on Windows",
        ))
    }
}

#[cfg(target_os = "windows")]
fn spawn_capture_session(
    callback: JsFunction,
    pid: u32,
    include_tree: bool,
) -> Result<u32> {
    // ErrorStrategy::Fatal — the JS callback signature is `(samples)`
    // not `(err, samples)`, so we don't want napi to inject an
    // error slot. If anything goes wrong on the Rust side we tear
    // the session down and stop calling the callback.
    let tsfn: ThreadsafeFunction<Float32Array, ErrorStrategy::Fatal> =
        callback.create_threadsafe_function(0, |ctx| Ok(vec![ctx.value]))?;

    let capture_id = NEXT_ID.fetch_add(1, Ordering::Relaxed);
    let stop = Arc::new(AtomicBool::new(false));
    let stop_clone = Arc::clone(&stop);

    let handle = thread::Builder::new()
        .name(format!("audio-loopback-{capture_id}"))
        .spawn(move || {
            if let Err(err) = windows_loopback::capture_loop(
                capture_id,
                tsfn,
                stop_clone,
                pid,
                include_tree,
            ) {
                eprintln!("audio-loopback {capture_id}: {err}");
            }
        })
        .map_err(|e| {
            napi::Error::from_reason(format!("failed to spawn audio thread: {e}"))
        })?;

    sessions().lock().unwrap().insert(
        capture_id,
        Session {
            stop,
            handle: Some(handle),
        },
    );
    Ok(capture_id)
}

#[cfg(not(target_os = "windows"))]
fn spawn_capture_session(
    callback: JsFunction,
    _pid: u32,
    _include_tree: bool,
) -> Result<u32> {
    let _ = callback;
    Err(napi::Error::from_reason(
        "system audio capture only supported on Windows",
    ))
}

/// Tear down the capture for the given id. Safe to call on a missing id
/// (no-op) so the JS side doesn't have to track whether the stop has
/// already been issued by the screen-share teardown path.
#[napi]
pub fn stop_capture(capture_id: u32) -> Result<()> {
    let session = sessions().lock().unwrap().remove(&capture_id);
    let Some(mut session) = session else {
        return Ok(());
    };
    session.stop.store(true, Ordering::Relaxed);
    if let Some(handle) = session.handle.take() {
        // Best-effort join — the capture loop polls `stop` every event
        // cycle (≤100ms) so this usually returns promptly. If a WASAPI
        // call is wedged we'd rather drop the handle than hang the JS
        // teardown path.
        let _ = handle.join();
    }
    Ok(())
}

// ---------------------------------------------------------------------------
// Windows loopback implementation
// ---------------------------------------------------------------------------

#[cfg(target_os = "windows")]
mod windows_loopback {
    use super::*;
    use wasapi::{initialize_mta, AudioClient, Direction, SampleType, ShareMode, WaveFormat};

    // 200ms request buffer in 100ns units. Process-loopback clients
    // ignore the period for shared-mode but the API still requires a
    // non-zero value — picking 200ms keeps wakeups infrequent enough
    // that we don't spin the capture thread on idle audio.
    const REQUESTED_BUFFER_HNS: i64 = 2_000_000;

    pub fn capture_loop(
        capture_id: u32,
        tsfn: ThreadsafeFunction<Float32Array, ErrorStrategy::Fatal>,
        stop: Arc<AtomicBool>,
        pid: u32,
        include_tree: bool,
    ) -> std::result::Result<(), String> {
        // COM must be initialised on every thread that touches WASAPI.
        // MTA is the right model for a background capture thread —
        // STA would require message pumping we don't want to add.
        initialize_mta()
            .ok()
            .map_err(|e| format!("initialize_mta: {e:?}"))?;

        // Process-loopback. Two modes — selected by `include_tree`:
        //   false → EXCLUDE_TARGET_PROCESS_TREE: every render session on
        //           the box *except* the given PID's tree. Default for
        //           full-screen shares so LiveKit playback stays out of
        //           the outgoing audio (we pass our own PID).
        //   true  → INCLUDE_TARGET_PROCESS_TREE: only the given PID's
        //           tree. Used for window-shares so the captured audio
        //           is exactly the picked app (Discord parity).
        // The `include_tree` flag is the load-bearing arg to
        // `new_application_loopback_client`; do not flip without
        // re-reading the wasapi crate's docs.
        let mut audio_client =
            AudioClient::new_application_loopback_client(pid, include_tree)
                .map_err(|e| format!("new_application_loopback_client: {e:?}"))?;

        // Process-loopback only accepts caller-specified formats —
        // `get_mixformat` is documented as broken on this client. We
        // pin the wire format we already deliver downstream: 48kHz,
        // 32-bit float, stereo. `autoconvert=true` (the trailing `true`
        // arg to initialize_client) lets WASAPI mix arbitrary session
        // formats into ours so games at 44.1k or mono notification
        // sounds don't blow up the capture.
        let wave_format = WaveFormat::new(
            32,
            32,
            &SampleType::Float,
            OUTPUT_SAMPLE_RATE as usize,
            OUTPUT_CHANNELS as usize,
            None,
        );

        audio_client
            .initialize_client(
                &wave_format,
                REQUESTED_BUFFER_HNS,
                &Direction::Capture,
                &ShareMode::Shared,
                true,
            )
            .map_err(|e| format!("initialize_client (process-loopback): {e:?}"))?;

        let h_event = audio_client
            .set_get_eventhandle()
            .map_err(|e| format!("set_get_eventhandle: {e:?}"))?;
        let capture_client = audio_client
            .get_audiocaptureclient()
            .map_err(|e| format!("get_audiocaptureclient: {e:?}"))?;

        audio_client
            .start_stream()
            .map_err(|e| format!("start_stream: {e:?}"))?;

        let block_align = wave_format.get_blockalign() as usize;

        while !stop.load(Ordering::Relaxed) {
            // 100ms timeout lets the loop check the stop flag even when
            // every excluded session is silent — there's nothing to
            // render so the event handle never fires.
            if h_event.wait_for_event(100).is_err() {
                continue;
            }

            // Drain all packets available since the last wake — there
            // can be several queued if we were preempted.
            loop {
                if stop.load(Ordering::Relaxed) {
                    break;
                }
                let frames_available = match capture_client.get_next_nbr_frames() {
                    Ok(Some(n)) if n > 0 => n,
                    Ok(_) => break,
                    Err(e) => {
                        eprintln!(
                            "audio-loopback {capture_id}: get_next_nbr_frames: {e:?}"
                        );
                        break;
                    }
                };

                let bytes_needed = frames_available as usize * block_align;
                let mut raw = vec![0u8; bytes_needed];
                if let Err(e) = capture_client.read_from_device(&mut raw) {
                    eprintln!(
                        "audio-loopback {capture_id}: read_from_device: {e:?}"
                    );
                    break;
                }

                // Reinterpret bytes as f32 little-endian samples. The
                // buffer is already 48kHz f32 stereo because process-
                // loopback autoconverted to our requested format. We
                // copy out into a Vec<f32> so napi can hand ownership
                // of a JS-owned ArrayBuffer to the renderer.
                let mut samples = Vec::<f32>::with_capacity(raw.len() / 4);
                let mut idx = 0;
                while idx + 4 <= raw.len() {
                    let bytes = [raw[idx], raw[idx + 1], raw[idx + 2], raw[idx + 3]];
                    samples.push(f32::from_le_bytes(bytes));
                    idx += 4;
                }
                let _ = (OUTPUT_SAMPLE_RATE, OUTPUT_CHANNELS);

                let payload = Float32Array::new(samples);
                // NonBlocking: never block the WASAPI capture thread on
                // a slow JS event loop — at 48kHz stereo a stalled
                // renderer would otherwise back-pressure the WASAPI
                // event handle and underrun every other consumer.
                let status = tsfn.call(payload, ThreadsafeFunctionCallMode::NonBlocking);
                if status != napi::Status::Ok {
                    // Renderer went away or the threadsafe function was
                    // released — stop cleanly.
                    stop.store(true, Ordering::Relaxed);
                    break;
                }
            }
        }

        let _ = audio_client.stop_stream();
        Ok(())
    }
}