vocero-client/s2s-frontend/src/audio/engine.ts

import type { NoiseGate } from "../lib/contracts";
import playbackUrl from "./audio-playback.worklet.ts?worker&url";
import { decodePcm } from "./codec";
import captureUrl from "./mic-capture.worklet.ts?worker&url";

const BANDS = [2, 5, 9, 16, 28, 52];
export const MIC_CONSTRAINTS: MediaStreamConstraints = {
  audio: { echoCancellation: true, noiseSuppression: true, autoGainControl: true },
  video: false,
};

/** Construct inside the click handler, before awaiting microphone permission. */
export function createAudioContext() {
  const context = new AudioContext({ latencyHint: "interactive" });
  void context.resume().catch(() => undefined);
  return context;
}

export class AudioEngine {
  private readonly context: AudioContext;
  private stream: MediaStream | null = null;
  private capture: AudioWorkletNode | null = null;
  private playback: AudioWorkletNode | null = null;
  private nodes: AudioNode[] = [];
  private frame = 0;
  private closed = false;
  private gate: NoiseGate;
  private muted = false;

  constructor(context: AudioContext, gate: NoiseGate) {
    this.context = context;
    this.gate = gate;
  }

  async start(onChunk: (chunk: ArrayBuffer) => void, onLevel: (rms: number) => void, speaking: () => boolean) {
    const stream = await navigator.mediaDevices.getUserMedia(MIC_CONSTRAINTS);
    if (this.closed) {
      stream.getTracks().forEach((track) => {
        track.stop();
      });
      return;
    }
    this.stream = stream;
    await Promise.all([
      this.context.audioWorklet.addModule(captureUrl),
      this.context.audioWorklet.addModule(playbackUrl),
    ]);
    if (this.closed) {
      return;
    }
    const capture = new AudioWorkletNode(this.context, "mic-capture", { numberOfInputs: 1, numberOfOutputs: 0 });
    const playback = new AudioWorkletNode(this.context, "audio-playback", {
      numberOfInputs: 0,
      numberOfOutputs: 1,
      outputChannelCount: [1],
    });
    // The checked-in speech pipeline emits native 16 kHz PCM, despite older 24 kHz comments.
    playback.port.postMessage({ kind: "config", inputRate: 16000 });
    this.capture = capture;
    this.playback = playback;
    this.setGate(this.gate);
    this.setMuted(this.muted);
    capture.port.onmessage = (event: MessageEvent<ArrayBuffer | { kind: "level"; rms: number }>) => {
      if (this.closed) {
        return;
      }
      if (event.data instanceof ArrayBuffer) {
        onChunk(event.data);
      } else {
        document.documentElement.style.setProperty("--audio-level", String(Math.min(1, event.data.rms * 5)));
        onLevel(event.data.rms);
      }
    };
    const input = this.context.createMediaStreamSource(stream);
    const micAnalyser = this.context.createAnalyser();
    const outAnalyser = this.context.createAnalyser();
    micAnalyser.fftSize = outAnalyser.fftSize = 256;
    micAnalyser.smoothingTimeConstant = 0;
    outAnalyser.smoothingTimeConstant = 0.3;
    input.connect(capture);
    input.connect(micAnalyser);
    playback.connect(outAnalyser);
    outAnalyser.connect(this.context.destination);
    this.nodes = [input, capture, playback, micAnalyser, outAnalyser];
    const buffer = new Uint8Array(micAnalyser.frequencyBinCount);
    const levels = new Float32Array(5);
    const tick = () => {
      const source = speaking() ? outAnalyser : micAnalyser;
      source.getByteFrequencyData(buffer);
      for (let band = 0; band < 5; band++) {
        const lo = BANDS[band] ?? 0;
        const hi = BANDS[band + 1] ?? 52;
        let sum = 0;
        for (let index = lo; index < hi; index++) {
          sum += buffer[index] ?? 0;
        }
        const target = sum / ((hi - lo) * 255);
        const previous = levels[band] ?? 0;
        levels[band] = previous + (target - previous) * (target > previous ? 0.6 : 0.18);
        document.documentElement.style.setProperty(`--bar${band}`, String(levels[band]));
      }
      outAnalyser.getByteFrequencyData(buffer);
      document.documentElement.style.setProperty("--ai-audio-level", String(Math.max(...buffer.subarray(0, 52)) / 255));
      this.frame = requestAnimationFrame(tick);
    };
    this.frame = requestAnimationFrame(tick);
  }
  play(base64: string) {
    if (!base64 || this.closed) {
      return;
    }
    const samples = decodePcm(base64);
    this.playback?.port.postMessage({ kind: "audio", samples }, [samples.buffer]);
  }
  clear() {
    this.playback?.port.postMessage({ kind: "clear" });
  }
  setGate(gate: NoiseGate) {
    this.gate = gate;
    this.capture?.port.postMessage({ kind: "gate", ...gate });
  }
  setMuted(muted: boolean) {
    this.muted = muted;
    this.stream?.getAudioTracks().forEach((track) => {
      track.enabled = !muted;
    });
    this.capture?.port.postMessage({ kind: "enable", value: !muted });
  }
  async close() {
    this.closed = true;
    cancelAnimationFrame(this.frame);
    this.stream?.getTracks().forEach((track) => {
      track.stop();
    });
    this.capture?.port.close();
    this.playback?.port.close();
    this.nodes.forEach((node) => {
      node.disconnect();
    });
    if (this.context.state !== "closed") {
      await this.context.close();
    }
    for (let index = 0; index < 5; index++) {
      document.documentElement.style.removeProperty(`--bar${index}`);
    }
    document.documentElement.style.removeProperty("--ai-audio-level");
    document.documentElement.style.removeProperty("--audio-level");
  }
}