element-call/src/lotus/lotusDenoiseProcessor.ts

/*
Copyright 2026 Lotus Guild

SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-Element-Commercial
Please see LICENSE in the repository root for full details.
*/

import {
  type AudioProcessorOptions,
  type Track,
  type TrackProcessor,
} from "livekit-client";
import { logger } from "matrix-js-sdk/lib/logger";

export type LotusDenoiseModel = "rnnoise" | "speex";

export interface LotusDenoiseConfig {
  model: LotusDenoiseModel;
  /** Base URL the worklet scripts + wasm are served from (e.g. "./denoise/"). */
  assetBase: string;
  gate: boolean;
  gateThreshold: number;
}

// Flat sapphi worklets: each registers a processor under these names when its
// script module is added. Same assets the Lotus host already ships under
// public/element-call/denoise/, so no new EC dependency is needed.
//
// ⚠️ CONTRACT: this is an undeclared, cross-repo asset dependency. The embedded
// fork build does NOT bundle these — they are copied in by cinny's
// vite.config.js `lotusDenoise()` plugin from `@sapphi-red/web-noise-suppressor`.
// The worklet/wasm version must match what this processor expects (sapphi
// rnnoise/speex, 48kHz). If the fork ever wants to own them, add
// `@sapphi-red/web-noise-suppressor` as a dep + a copy step here. Until then,
// an integration smoke-check should assert GET .../denoise/rnnoise.wasm == 200.
const PROCESSORS: Record<
  LotusDenoiseModel,
  { name: string; script: string; wasm: string }
> = {
  rnnoise: {
    name: "@sapphi-red/web-noise-suppressor/rnnoise",
    script: "rnnoiseWorklet.js",
    wasm: "rnnoise.wasm",
  },
  speex: {
    name: "@sapphi-red/web-noise-suppressor/speex",
    script: "speexWorklet.js",
    wasm: "speex.wasm",
  },
};
const GATE_NAME = "@sapphi-red/web-noise-suppressor/noiseGate";
const GATE_SCRIPT = "noiseGateWorklet.js";

// The sapphi rnnoise/speex worklets are 48kHz-fullband and don't resample, so
// the graph MUST run at 48kHz regardless of the hardware default.
const SAMPLE_RATE = 48_000;

// Cache fetched wasm per URL so a reconnect/device-switch (which rebuilds the
// graph) doesn't re-download it.
const wasmCache = new Map<string, Promise<ArrayBuffer>>();
function loadWasm(url: string): Promise<ArrayBuffer> {
  let p = wasmCache.get(url);
  if (!p) {
    p = fetch(url).then((r) => {
      if (!r.ok) throw new Error(`denoise wasm ${r.status}`);
      return r.arrayBuffer();
    });
    wasmCache.set(url, p);
  }
  return p;
}

interface Graph {
  source: MediaStreamAudioSourceNode;
  nodes: AudioWorkletNode[];
  track: MediaStreamTrack;
}

/**
 * A LiveKit audio TrackProcessor that runs Lotus ML noise suppression
 * (RNNoise/Speex) on the local microphone track, as a first-class stage in
 * Element Call's publish pipeline.
 *
 * This replaces the host's `getUserMedia` monkeypatch: because it's a real
 * LiveKit processor, EC re-applies it on every (re)publish/restart, so denoise
 * survives EC's own mid-call reconnect — the root cause of the A7 "mic dead
 * after reconnect" bug.
 *
 * It owns a dedicated 48kHz AudioContext (the worklets require it, and LiveKit
 * does NOT pass an audioContext to restart()), reused across restarts and
 * closed on destroy. restart() never throws and never leaves a stopped track
 * on the sender: on failure it degrades to the raw (un-denoised) track rather
 * than silence.
 */
export class LotusDenoiseProcessor
  implements TrackProcessor<Track.Kind.Audio, AudioProcessorOptions>
{
  public readonly name = "lotus-denoise";
  public processedTrack?: MediaStreamTrack;

  private ctx?: AudioContext;
  private graph?: Graph;

  public constructor(private readonly config: LotusDenoiseConfig) {}

  public async init(_opts: AudioProcessorOptions): Promise<void> {
    // Own context at the required rate — do NOT use opts.audioContext (it's the
    // room's hardware-rate context, and is undefined on restart()).
    this.ctx = new AudioContext({ sampleRate: SAMPLE_RATE });
    const graph = await this.buildGraph(_opts.track);
    this.graph = graph;
    this.processedTrack = graph.track;
  }

  public async restart(opts: AudioProcessorOptions): Promise<void> {
    if (!this.ctx || this.ctx.state === "closed")
      this.ctx = new AudioContext({ sampleRate: SAMPLE_RATE });
    try {
      const next = await this.buildGraph(opts.track);
      this.disposeGraph(this.graph);
      this.graph = next;
      this.processedTrack = next.track;
    } catch (e) {
      // Never go silent on the A7/device-switch path: fall back to raw audio.
      logger.warn("[lotus] denoise restart failed; using raw mic", e);
      this.disposeGraph(this.graph);
      this.graph = undefined;
      this.processedTrack = opts.track;
    }
  }

  public async destroy(): Promise<void> {
    this.disposeGraph(this.graph);
    this.graph = undefined;
    this.processedTrack = undefined;
    if (this.ctx && this.ctx.state !== "closed")
      await this.ctx.close().catch(() => undefined);
    this.ctx = undefined;
  }

  private async buildGraph(track: MediaStreamTrack): Promise<Graph> {
    const ctx = this.ctx!;
    const base = this.config.assetBase;
    const proc = PROCESSORS[this.config.model];

    await ctx.audioWorklet.addModule(base + proc.script);
    if (this.config.gate) await ctx.audioWorklet.addModule(base + GATE_SCRIPT);
    const wasmBinary = await loadWasm(base + proc.wasm);

    const source = ctx.createMediaStreamSource(new MediaStream([track]));
    const dest = ctx.createMediaStreamDestination();
    const nodes: AudioWorkletNode[] = [];
    let head: AudioNode = source;

    if (this.config.gate) {
      const gate = new AudioWorkletNode(ctx, GATE_NAME, {
        processorOptions: {
          openThreshold: this.config.gateThreshold,
          closeThreshold: this.config.gateThreshold - 5,
          holdMs: 150,
          maxChannels: 1,
        },
      });
      head.connect(gate);
      head = gate;
      nodes.push(gate);
    }

    const ml = new AudioWorkletNode(ctx, proc.name, {
      channelCount: 1,
      numberOfInputs: 1,
      numberOfOutputs: 1,
      processorOptions: { maxChannels: 1, wasmBinary: wasmBinary.slice(0) },
    });
    head.connect(ml);
    ml.connect(dest);
    nodes.push(ml);

    logger.info(`[lotus] denoise processor active (${this.config.model})`);
    return { source, nodes, track: dest.stream.getAudioTracks()[0] };
  }

  private disposeGraph(graph: Graph | undefined): void {
    if (!graph) return;
    for (const node of graph.nodes) {
      try {
        node.port.postMessage("destroy");
      } catch {
        /* ignore */
      }
      try {
        node.disconnect();
      } catch {
        /* ignore */
      }
    }
    try {
      graph.source.disconnect();
    } catch {
      /* ignore */
    }
    graph.track.stop();
  }
}