element-call/src/lotus/lotusDenoise.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 LocalAudioTrack,
  ParticipantEvent,
  type Room as LivekitRoom,
  Track,
} from "livekit-client";
import { logger } from "matrix-js-sdk/lib/logger";

import { type CallViewModel } from "../state/CallViewModel/CallViewModel";
import { lotusFlag, lotusParam } from "./lotusWidget";
import {
  type LotusDenoiseConfig,
  LotusDenoiseProcessor,
} from "./lotusDenoiseProcessor";

/**
 * Apply Lotus ML noise suppression to the local mic track as a first-class
 * Element Call audio TrackProcessor (#1), replacing the host's build-time
 * `getUserMedia` monkeypatch.
 *
 * Opt-in via `lotusDenoiseSource=1` (+ `lotusModel`, `lotusGate`,
 * `lotusGateThreshold`, `lotusDenoiseBase`). Because the processor is attached
 * to the mic publication and re-attached on every (re)publish, denoise
 * survives EC's mid-call reconnect — fixing the A7 "mic dead after reconnect"
 * bug. Additive: no-op without the flag.
 */
/**
 * Resolve the denoise asset base, forcing it to be SAME-ORIGIN. The base is fed
 * to `audioWorklet.addModule()`, which executes the target as code in the
 * worklet scope (processing the live mic) — so a cross-origin base from a
 * crafted call link would be arbitrary code execution. Any non-same-origin or
 * malformed value falls back to the bundled "./denoise/".
 */
function safeAssetBase(raw: string | null): string {
  // Resolve to an ABSOLUTE same-origin href against the document. Absolute is
  // required because native dynamic `import()` (DTLN/DeepFilterNet) resolves a
  // relative specifier against the JS chunk's URL, not the document — so a
  // relative "./denoise/" would 404. addModule()/fetch() work with absolute
  // too, so this keeps all three asset-load paths consistent.
  const fallback = new URL("./denoise/", window.location.href).href;
  if (!raw) return fallback;
  try {
    const u = new URL(raw, window.location.href);
    if (u.origin !== window.location.origin) return fallback;
    return u.href.endsWith("/") ? u.href : `${u.href}/`;
  } catch {
    return fallback;
  }
}

export function startLotusDenoise(vm: CallViewModel): () => void {
  // Gate on a DEDICATED flag — NOT the existing `lotusDenoise=ml` that the
  // host's build-time getUserMedia shim already uses. Otherwise simply shipping
  // this fork (while the host still injects its shim and still sets
  // lotusDenoise=ml) would denoise twice. The host opts into the in-source
  // engine with `lotusDenoiseSource=1` AND stops injecting the shim at the same
  // time. Default off ⇒ the fork is inert and behaviour is unchanged.
  if (!lotusFlag("lotusDenoiseSource")) return () => undefined;

  const requested = lotusParam("lotusModel");
  const model: LotusDenoiseConfig["model"] =
    requested === "speex" ||
    requested === "dtln" ||
    requested === "deepfilternet"
      ? requested
      : "rnnoise";

  const rawThreshold = lotusParam("lotusGateThreshold");
  const rawFloor = lotusParam("lotusDenoiseFloor");
  const config: LotusDenoiseConfig = {
    model,
    assetBase: safeAssetBase(lotusParam("lotusDenoiseBase")),
    gate: lotusFlag("lotusGate"),
    // Default -45 (matches the reference shim); accept an explicit 0 (don't
    // coerce it away via `|| default`).
    gateThreshold:
      rawThreshold !== null && Number.isFinite(Number(rawThreshold))
        ? Number(rawThreshold)
        : -45,
    // Dry/wet attenuation floor. Default 0.15 (~-16 dB) tames the
    // over-suppression "underwater"/pumping artifact; host can tune via
    // `lotusDenoiseFloor` (0 = full suppression, no floor).
    floor:
      rawFloor !== null && Number.isFinite(Number(rawFloor))
        ? Math.min(0.5, Math.max(0, Number(rawFloor)))
        : 0.15,
  };

  const micOf = (room: LivekitRoom): LocalAudioTrack | undefined =>
    room.localParticipant.getTrackPublication(Track.Source.Microphone)
      ?.track as LocalAudioTrack | undefined;

  // [lotus] `mic.getProcessor()` only becomes set once `setProcessor()`
  // resolves — i.e. after the whole wasm/model load. LiveKit fires
  // `LocalTrackPublished` once per local track (mic, then camera on join with
  // video), so two calls to `apply()` can both observe `!mic.getProcessor()`
  // and race to construct a second `LotusDenoiseProcessor` (a second
  // AudioContext + model load) before the first has attached. Track an
  // in-flight setProcessor per room and skip `apply()` while one is pending.
  const pendingProcessors = new Map<LivekitRoom, LotusDenoiseProcessor>();

  const apply = (room: LivekitRoom): void => {
    const mic = micOf(room);
    if (!mic || mic.getProcessor() || pendingProcessors.has(room)) return;
    const processor = new LotusDenoiseProcessor(config);
    pendingProcessors.set(room, processor);
    void mic
      .setProcessor(processor)
      .catch((e) => logger.warn("[lotus] denoise setProcessor failed", e))
      .finally(() => {
        // Only clear if we're still the pending entry (a teardown that ran
        // while this was in flight may have already replaced/removed it).
        if (pendingProcessors.get(room) === processor)
          pendingProcessors.delete(room);
      });
  };

  const roomListeners = new Map<LivekitRoom, () => void>();
  let rooms: LivekitRoom[] = [];

  // Drive activation off the LOCAL participant's connection(s), not
  // `livekitRoomItems$` — that stream excludes the local participant and only
  // surfaces rooms with ≥1 remote member, so it wouldn't denoise you while
  // you're alone and is a fragile coupling to a remote-render concern.
  const sub = vm.allConnections$.subscribe((data) => {
    const next = data.getConnections().map((c) => c.livekitRoom);
    rooms = next;
    for (const [room, off] of roomListeners) {
      if (!next.includes(room)) {
        off();
        roomListeners.delete(room);
      }
    }
    for (const room of next) {
      if (!roomListeners.has(room)) {
        // Re-attach on every (re)publish — this is what makes denoise survive
        // reconnects (A7), unlike the old getUserMedia patch.
        const onPublished = (): void => apply(room);
        room.localParticipant.on(
          ParticipantEvent.LocalTrackPublished,
          onPublished,
        );
        roomListeners.set(room, () =>
          room.localParticipant.off(
            ParticipantEvent.LocalTrackPublished,
            onPublished,
          ),
        );
        apply(room);
      }
    }
  });

  return () => {
    sub.unsubscribe();
    for (const off of roomListeners.values()) off();
    roomListeners.clear();
    for (const room of rooms) {
      const mic = micOf(room);
      if (mic?.getProcessor()) void mic.stopProcessor();
      else {
        // [lotus] A setProcessor() call may still be in flight (mid wasm/model
        // load) when teardown runs, in which case `mic.getProcessor()` is
        // still undefined and `stopProcessor()` above is a no-op. Destroy the
        // pending processor directly so its AudioContext/graph don't leak.
        const pending = pendingProcessors.get(room);
        if (pending) void pending.destroy().catch(() => undefined);
      }
    }
    pendingProcessors.clear();
  };
}