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 `lotusDenoise=ml` (+ `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 {
  const fallback = "./denoise/";
  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 {
  if (lotusParam("lotusDenoise") !== "ml") return () => undefined;

  const requestedModel = lotusParam("lotusModel");
  // Only RNNoise/Speex are implemented in-source so far. Don't silently
  // degrade an unsupported model (e.g. the host's dtln/deepfilternet) to
  // rnnoise — log it so the mismatch is visible.
  let model: LotusDenoiseConfig["model"] = "rnnoise";
  if (requestedModel === "speex") model = "speex";
  else if (requestedModel && requestedModel !== "rnnoise")
    logger.warn(
      `[lotus] denoise model "${requestedModel}" not implemented in-source; using rnnoise`,
    );

  const rawThreshold = lotusParam("lotusGateThreshold");
  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,
  };

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

  const apply = (room: LivekitRoom): void => {
    const mic = micOf(room);
    if (mic && !mic.getProcessor()) {
      void mic
        .setProcessor(new LotusDenoiseProcessor(config))
        .catch((e) => logger.warn("[lotus] denoise setProcessor failed", e));
    }
  };

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

  const sub = vm.livekitRoomItems$.subscribe((items) => {
    const next = items.map((i) => i.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();
    }
  };
}