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.
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";

/**
 * 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, so denoise
 * survives EC's own mid-call reconnect — the root cause of the A7
 * "mic dead after reconnect" bug.
 */
export class LotusDenoiseProcessor
  implements TrackProcessor<Track.Kind.Audio, AudioProcessorOptions>
{
  public readonly name = "lotus-denoise";
  public processedTrack?: MediaStreamTrack;

  private source?: MediaStreamAudioSourceNode;
  private nodes: AudioWorkletNode[] = [];

  public constructor(private readonly config: LotusDenoiseConfig) {}

  public async init(opts: AudioProcessorOptions): Promise<void> {
    await this.build(opts.audioContext, opts.track);
  }

  public async restart(opts: AudioProcessorOptions): Promise<void> {
    this.teardownGraph();
    await this.build(opts.audioContext, opts.track);
  }

  public async destroy(): Promise<void> {
    this.teardownGraph();
  }

  private async build(
    ctx: AudioContext,
    track: MediaStreamTrack,
  ): Promise<void> {
    const base = this.config.assetBase;
    const proc = PROCESSORS[this.config.model];

    // Register worklet modules (idempotent) and load the wasm binary.
    await ctx.audioWorklet.addModule(base + proc.script);
    if (this.config.gate) await ctx.audioWorklet.addModule(base + GATE_SCRIPT);
    const wasmBinary = await fetch(base + proc.wasm).then((r) => {
      if (!r.ok) throw new Error(`denoise wasm ${r.status}`);
      return r.arrayBuffer();
    });

    const source = ctx.createMediaStreamSource(new MediaStream([track]));
    const dest = ctx.createMediaStreamDestination();
    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;
      this.nodes.push(gate);
    }

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

    this.source = source;
    this.processedTrack = dest.stream.getAudioTracks()[0];
    logger.info(`[lotus] denoise processor active (${this.config.model})`);
  }

  private teardownGraph(): void {
    for (const node of this.nodes) {
      try {
        node.port.postMessage("destroy");
      } catch {
        /* ignore */
      }
      try {
        node.disconnect();
      } catch {
        /* ignore */
      }
    }
    this.nodes = [];
    try {
      this.source?.disconnect();
    } catch {
      /* ignore */
    }
    this.source = undefined;
    this.processedTrack?.stop();
    this.processedTrack = undefined;
  }
}