element-call/src/lotus/lotusDenoise.test.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 { afterEach, beforeEach, describe, expect, test, vi } from "vitest";
import { ParticipantEvent, Track } from "livekit-client";

import { type CallViewModel } from "../state/CallViewModel/CallViewModel";
import {
  applySetDenoise,
  browserNsFor,
  type DenoiseTarget,
  startLotusDenoise,
} from "./lotusDenoise";
import { LotusWidgetActions } from "./lotusActions";

// Track constructed instances so tests can assert exactly one processor was
// built across racing apply() calls, and can assert the pending one is
// destroyed on early teardown. `vi.hoisted` is required because `vi.mock`
// factories are hoisted above this file's other top-level statements.
const instances = vi.hoisted(
  () =>
    [] as {
      config: { model: string };
      destroy: ReturnType<typeof vi.fn>;
      setMicMuted: ReturnType<typeof vi.fn>;
    }[],
);
// Ordered log of the heavy-lifting calls, to assert prepare runs before
// setProcessor (#7).
const calls = vi.hoisted(() => [] as string[]);
vi.mock("./lotusDenoiseProcessor", () => ({
  LotusDenoiseProcessor: class {
    public destroy = vi.fn().mockResolvedValue(undefined);
    public setMicMuted = vi.fn();
    public constructor(public config: { model: string }) {
      instances.push(this);
    }
  },
  prepareDenoiseAssets: vi.fn(async () => {
    calls.push("prepare");
    await Promise.resolve();
  }),
  releasePreparedDenoiseAssets: vi.fn().mockResolvedValue(undefined),
}));
const hostSend = vi.hoisted(() => vi.fn().mockResolvedValue({}));
// Host → widget actions: tests call `sendAction()` to deliver one.
const actionHandlers = vi.hoisted(
  () => new Map<string, (ev: CustomEvent) => void>(),
);
vi.mock("../widget", () => ({
  widget: {
    api: { transport: { send: hostSend, reply: vi.fn() } },
    lazyActions: {
      on: (name: string, cb: (ev: CustomEvent) => void): void => {
        actionHandlers.set(name, cb);
      },
      off: (name: string): void => {
        actionHandlers.delete(name);
      },
    },
  },
}));
const sendAction = (name: string, data: unknown): void =>
  actionHandlers.get(name)?.(new CustomEvent(name, { detail: { data } }));

/** A promise plus externally-callable resolve, for controlling ordering. */
function deferred<T>(): {
  promise: Promise<T>;
  resolve: (v: T) => void;
} {
  let resolve!: (v: T) => void;
  const promise = new Promise<T>((r) => {
    resolve = r;
  });
  return { promise, resolve };
}

function makeRoomAndVm(mic: {
  getProcessor: () => unknown;
  setProcessor: ReturnType<typeof vi.fn>;
  isMuted?: boolean;
}): {
  vm: CallViewModel;
  room: { localParticipant: Record<string, unknown> };
  firePublished: () => void;
  fire: (event: string, arg?: unknown) => void;
} {
  const handlers = new Map<string, (arg?: unknown) => void>();
  const localParticipant = {
    getTrackPublication: (
      source: Track.Source,
    ): { track: typeof mic } | undefined =>
      source === Track.Source.Microphone ? { track: mic } : undefined,
    on: (
      event: string,
      cb: (arg?: unknown) => void,
    ): Map<string, (arg?: unknown) => void> => handlers.set(event, cb),
    off: (event: string): boolean => handlers.delete(event),
  };
  const room = { localParticipant };
  const vm = {
    allConnections$: {
      subscribe: (
        cb: (data: {
          getConnections: () => { livekitRoom: unknown }[];
        }) => void,
      ) => {
        cb({ getConnections: () => [{ livekitRoom: room }] });
        return { unsubscribe: (): void => undefined };
      },
    },
  } as unknown as CallViewModel;
  return {
    vm,
    room,
    firePublished: () => handlers.get(ParticipantEvent.LocalTrackPublished)?.(),
    fire: (event, arg) => handlers.get(event)?.(arg),
  };
}

/** Let the async apply() chain (attach → retry → host notify) settle. */
async function flush(): Promise<void> {
  for (let i = 0; i < 8; i++) await Promise.resolve();
}

beforeEach(() => {
  instances.length = 0;
  calls.length = 0;
  hostSend.mockClear();
  actionHandlers.clear();
  // `lotusParam`/`lotusFlag` cache the URL params at first read; seed the hash
  // before startLotusDenoise() so the dedicated `lotusDenoiseSource` flag reads on.
  window.location.hash = "#/room?lotusDenoiseSource=1";
});

afterEach(() => {
  vi.resetModules();
});

describe("startLotusDenoise", () => {
  test("racing LocalTrackPublished events only construct one processor", async () => {
    const setProcessorDeferred = deferred<void>();
    let attached: unknown;
    const mic = {
      getProcessor: (): unknown => attached,
      setProcessor: vi.fn(async (p: unknown) => {
        calls.push("setProcessor");
        await setProcessorDeferred.promise;
        attached = p;
      }),
    };
    const { vm, firePublished } = makeRoomAndVm(mic);

    startLotusDenoise(vm);
    // [#7] The heavy assets were kicked off before setProcessor() was called.
    expect(calls).toEqual(["prepare", "setProcessor"]);
    // Simulate a second LocalTrackPublished (e.g. camera) firing before the
    // first setProcessor() has resolved.
    firePublished();

    expect(mic.setProcessor).toHaveBeenCalledTimes(1);
    expect(instances).toHaveLength(1);

    setProcessorDeferred.resolve();
    await Promise.resolve();
    await Promise.resolve();

    // Once attached, a further publish should be a no-op (mic.getProcessor()
    // is now set).
    firePublished();
    expect(mic.setProcessor).toHaveBeenCalledTimes(1);
  });

  test("destroys the in-flight processor if torn down before setProcessor resolves", async () => {
    const setProcessorDeferred = deferred<void>();
    const mic = {
      getProcessor: (): unknown => undefined,
      setProcessor: vi.fn(async () => setProcessorDeferred.promise),
    };
    const { vm } = makeRoomAndVm(mic);

    const teardown = startLotusDenoise(vm);
    expect(instances).toHaveLength(1);

    teardown();

    expect(instances[0].destroy).toHaveBeenCalledTimes(1);

    // Resolving afterwards must not throw/reject unhandled.
    setProcessorDeferred.resolve();
    await Promise.resolve();
  });

  test("reports the active model to the host once attached (#8)", async () => {
    let attached: unknown;
    const mic = {
      getProcessor: (): unknown => attached,
      setProcessor: vi.fn(async (p: unknown) => {
        await Promise.resolve();
        attached = p;
      }),
    };
    const { vm } = makeRoomAndVm(mic);
    startLotusDenoise(vm);
    await flush();
    expect(hostSend).toHaveBeenCalledWith(LotusWidgetActions.DenoiseState, {
      active: true,
      model: "rnnoise",
      mode: "ml",
    });
  });

  test("retries once with rnnoise when the requested model fails to init (#8)", async () => {
    window.location.hash =
      "#/room?lotusDenoiseSource=1&lotusModel=deepfilternet";
    let attached: unknown;
    const mic = {
      getProcessor: (): unknown => attached,
      setProcessor: vi.fn(async (p: { config: { model: string } }) => {
        await Promise.resolve();
        if (p.config.model === "deepfilternet")
          throw new Error("dfn assets missing");
        attached = p;
      }),
    };
    const { vm } = makeRoomAndVm(mic);
    startLotusDenoise(vm);
    await flush();

    expect(instances.map((i) => i.config.model)).toEqual([
      "deepfilternet",
      "rnnoise",
    ]);
    expect(mic.setProcessor).toHaveBeenCalledTimes(2);
    expect(hostSend).toHaveBeenCalledTimes(1);
    expect(hostSend).toHaveBeenCalledWith(LotusWidgetActions.DenoiseState, {
      active: true,
      model: "rnnoise",
      mode: "ml",
    });
  });

  test("notifies the host with the error when the rnnoise fallback also fails (#8)", async () => {
    window.location.hash = "#/room?lotusDenoiseSource=1&lotusModel=dtln";
    const mic = {
      getProcessor: (): unknown => undefined,
      setProcessor: vi.fn(async () => {
        await Promise.resolve();
        throw new Error("no worklet for you");
      }),
    };
    const { vm } = makeRoomAndVm(mic);
    startLotusDenoise(vm);
    await flush();

    expect(instances.map((i) => i.config.model)).toEqual(["dtln", "rnnoise"]);
    expect(hostSend).toHaveBeenCalledTimes(1);
    expect(hostSend).toHaveBeenCalledWith(LotusWidgetActions.DenoiseState, {
      active: false,
      error: "no worklet for you",
      model: "rnnoise",
      mode: "ml",
    });
  });

  test("mirrors mic TrackMuted/TrackUnmuted onto the processor (#9)", async () => {
    let attached: unknown;
    const mic = {
      isMuted: true,
      getProcessor: (): unknown => attached,
      setProcessor: vi.fn(async (p: unknown) => {
        await Promise.resolve();
        attached = p;
      }),
    };
    const { vm, fire } = makeRoomAndVm(mic);
    startLotusDenoise(vm);
    // Seeded from the mic's current state before setProcessor().
    expect(instances[0].setMicMuted).toHaveBeenLastCalledWith(true);
    await flush();

    const micPub = { source: Track.Source.Microphone };
    const camPub = { source: Track.Source.Camera };
    fire(ParticipantEvent.TrackUnmuted, micPub);
    expect(instances[0].setMicMuted).toHaveBeenLastCalledWith(false);
    fire(ParticipantEvent.TrackMuted, micPub);
    expect(instances[0].setMicMuted).toHaveBeenLastCalledWith(true);

    // Camera mute must not touch the audio context.
    instances[0].setMicMuted.mockClear();
    fire(ParticipantEvent.TrackMuted, camPub);
    fire(ParticipantEvent.TrackUnmuted, camPub);
    expect(instances[0].setMicMuted).not.toHaveBeenCalled();
  });
});

/** A mic whose processor and capture constraints behave like LiveKit's. */
function liveMic(): {
  getProcessor: () => { config: { model: string } } | undefined;
  setProcessor: ReturnType<typeof vi.fn>;
  stopProcessor: ReturnType<typeof vi.fn>;
  restartTrack: ReturnType<typeof vi.fn>;
  constraints: MediaTrackConstraints;
  mediaStreamTrack: { getSettings: () => MediaTrackSettings };
  isMuted: boolean;
} {
  let attached: { config: { model: string } } | undefined;
  return {
    isMuted: false,
    getProcessor: () => attached,
    setProcessor: vi.fn(async (p: { config: { model: string } }) => {
      await Promise.resolve();
      attached = p;
    }),
    stopProcessor: vi.fn(async () => {
      await Promise.resolve();
      attached = undefined;
    }),
    restartTrack: vi.fn().mockResolvedValue(undefined),
    constraints: { echoCancellation: true, autoGainControl: true },
    mediaStreamTrack: { getSettings: () => ({ deviceId: "mic-1" }) },
  };
}

describe("io.lotus.set_denoise (switch mid-call)", () => {
  test("swaps the ML model on the live mic without re-capturing it", async () => {
    window.location.hash =
      "#/room?lotusDenoiseSource=1&lotusModel=dtln&noiseSuppression=false";
    const mic = liveMic();
    const { vm } = makeRoomAndVm(mic);
    startLotusDenoise(vm);
    await flush();
    expect(mic.getProcessor()?.config.model).toBe("dtln");

    sendAction(LotusWidgetActions.SetDenoise, {
      mode: "ml",
      model: "deepfilternet",
    });
    await flush();
    await flush();

    expect(mic.stopProcessor).toHaveBeenCalledTimes(1);
    expect(mic.getProcessor()?.config.model).toBe("deepfilternet");
    expect(mic.restartTrack).not.toHaveBeenCalled();
    expect(hostSend).toHaveBeenLastCalledWith(LotusWidgetActions.DenoiseState, {
      active: true,
      model: "deepfilternet",
      mode: "ml",
    });
  });

  test("ML → browser-native drops the model and re-captures with browser NS on", async () => {
    window.location.hash =
      "#/room?lotusDenoiseSource=1&lotusModel=rnnoise&noiseSuppression=false";
    const mic = liveMic();
    const { vm } = makeRoomAndVm(mic);
    startLotusDenoise(vm);
    await flush();

    sendAction(LotusWidgetActions.SetDenoise, { mode: "browser" });
    await flush();
    await flush();

    expect(mic.stopProcessor).toHaveBeenCalledTimes(1);
    expect(mic.getProcessor()).toBeUndefined();
    expect(mic.restartTrack).toHaveBeenCalledWith({
      deviceId: "mic-1",
      channelCount: 1,
      echoCancellation: true,
      autoGainControl: true,
      noiseSuppression: true,
      voiceIsolation: true,
    });
    expect(hostSend).toHaveBeenLastCalledWith(LotusWidgetActions.DenoiseState, {
      active: false,
      mode: "browser",
    });
  });

  test("browser-native → ML works even when the call joined without ML", async () => {
    window.location.hash = "#/room?noiseSuppression=true";
    const mic = liveMic();
    const { vm } = makeRoomAndVm(mic);
    startLotusDenoise(vm);
    await flush();
    expect(mic.setProcessor).not.toHaveBeenCalled();

    sendAction(LotusWidgetActions.SetDenoise, { mode: "ml", model: "speex" });
    await flush();
    await flush();

    expect(mic.restartTrack).toHaveBeenCalledWith(
      expect.objectContaining({ noiseSuppression: false }),
    );
    expect(mic.getProcessor()?.config.model).toBe("speex");
  });

  test("ignores a payload without a valid mode", async () => {
    window.location.hash = "#/room?lotusDenoiseSource=1";
    const mic = liveMic();
    const { vm } = makeRoomAndVm(mic);
    startLotusDenoise(vm);
    await flush();
    sendAction(LotusWidgetActions.SetDenoise, { mode: "loud" });
    await flush();
    expect(mic.stopProcessor).not.toHaveBeenCalled();
    expect(mic.restartTrack).not.toHaveBeenCalled();
  });
});

describe("applySetDenoise", () => {
  const base: DenoiseTarget = {
    mode: "ml",
    nativeNS: false,
    config: {
      model: "rnnoise",
      assetBase: "https://x/denoise/",
      gate: false,
      gateThreshold: -45,
      floor: 0.15,
    },
  };

  test("keeps fields the payload doesn't name or gets wrong", () => {
    expect(
      applySetDenoise(base, { mode: "ml", model: "nope", gate: "yes" }),
    ).toEqual(base);
  });

  test("moves only the named fields", () => {
    const next = applySetDenoise(base, {
      mode: "ml",
      model: "dtln",
      gateThreshold: -50,
    });
    expect(next?.config).toEqual({
      ...base.config,
      model: "dtln",
      gateThreshold: -50,
    });
  });

  test("browser NS is on for browser mode and for ML with nativeNS", () => {
    expect(browserNsFor({ ...base, mode: "browser" })).toBe(true);
    expect(browserNsFor({ ...base, mode: "off" })).toBe(false);
    expect(browserNsFor(base)).toBe(false);
    expect(browserNsFor({ ...base, nativeNS: true })).toBe(true);
  });
});