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,
  type TrackPublication,
} from "livekit-client";
import { logger } from "matrix-js-sdk/lib/logger";

import { type CallViewModel } from "../state/CallViewModel/CallViewModel";
import {
  LotusWidgetActions,
  lotusFlag,
  lotusParam,
  lotusSendToHost,
} from "./lotusWidget";
import {
  type LotusDenoiseConfig,
  type LotusDenoiseModel,
  LotusDenoiseProcessor,
  prepareDenoiseAssets,
  releasePreparedDenoiseAssets,
} 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,
  };

  // [lotus #7] Start fetching wasm / creating the AudioContext / addModule-ing
  // the worklet (and loading the DFN model) NOW, before any mic track exists.
  // `setProcessor()` holds LiveKit's trackChangeLock while awaiting `init()`,
  // so anything still loading there freezes mute/unmute/device-switch.
  void prepareDenoiseAssets(config);

  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>();
  let stopped = false;

  const attach = async (
    room: LivekitRoom,
    mic: LocalAudioTrack,
    model: LotusDenoiseModel,
  ): Promise<void> => {
    const processor = new LotusDenoiseProcessor({ ...config, model });
    pendingProcessors.set(room, processor);
    // [lotus #9] Seed the mute state so a processor attached while already
    // muted starts suspended rather than burning inference on silence.
    processor.setMicMuted(mic.isMuted);
    try {
      await mic.setProcessor(processor);
    } 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 apply = (room: LivekitRoom): void => {
    const mic = micOf(room);
    if (!mic || mic.getProcessor() || pendingProcessors.has(room)) return;
    void (async () => {
      let model = config.model;
      try {
        try {
          await attach(room, mic, model);
        } catch (e) {
          logger.warn("[lotus] denoise setProcessor failed", e);
          // [lotus #8] A failed init leaves LiveKit with no processor (it only
          // assigns after init resolves), so retry once with the smallest,
          // most portable tier before giving up.
          if (model === "rnnoise" || stopped) throw e;
          model = "rnnoise";
          await attach(room, mic, model);
        }
        // [lotus #8] Tell the host what is ACTUALLY running so its toggle
        // reflects reality (including the fallback model).
        lotusSendToHost(LotusWidgetActions.DenoiseState, {
          active: true,
          model,
        });
      } catch (e) {
        logger.warn("[lotus] denoise unavailable; publishing raw mic", e);
        lotusSendToHost(LotusWidgetActions.DenoiseState, {
          active: false,
          error: e instanceof Error ? e.message : String(e),
          model,
        });
      }
    })();
  };

  // [lotus #9] Suspend/resume the processor's own AudioContext with the mic's
  // mute state (LiveKit keeps the muted track flowing silent frames, so the
  // worklet would otherwise run inference the whole time the user is muted).
  const onMuteChange =
    (room: LivekitRoom, muted: boolean) =>
    (pub: TrackPublication): void => {
      if (pub.source !== Track.Source.Microphone) return;
      const p = pendingProcessors.get(room) ?? micOf(room)?.getProcessor();
      if (p instanceof LotusDenoiseProcessor) p.setMicMuted(muted);
    };

  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);
        const onMuted = onMuteChange(room, true);
        const onUnmuted = onMuteChange(room, false);
        room.localParticipant.on(
          ParticipantEvent.LocalTrackPublished,
          onPublished,
        );
        room.localParticipant.on(ParticipantEvent.TrackMuted, onMuted);
        room.localParticipant.on(ParticipantEvent.TrackUnmuted, onUnmuted);
        roomListeners.set(room, () => {
          room.localParticipant.off(
            ParticipantEvent.LocalTrackPublished,
            onPublished,
          );
          room.localParticipant.off(ParticipantEvent.TrackMuted, onMuted);
          room.localParticipant.off(ParticipantEvent.TrackUnmuted, onUnmuted);
        });
        apply(room);
      }
    }
  });

  return () => {
    stopped = true;
    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();
    // [lotus #7] Close any prepared-but-never-claimed context (e.g. no mic).
    void releasePreparedDenoiseAssets();
  };
}