Blame for server/pcm-segmenter.ts

/**
 * pcm-segmenter.ts — Continuous PCM16 utterance segmenter for streaming STT.
 *
 * Turns an endless stream of raw PCM16LE mono audio into discrete utterance
 * segments suitable for a request/response STT API (Whisper), without losing
 * speech at segment boundaries:
 *
 *   - PRE-ROLL: audio captured just BEFORE speech onset is prepended to the
 *     utterance so soft word onsets are not clipped.
 *   - HANGOVER: an utterance only ends after a sustained silence, so trailing
 *     syllables and short pauses inside a sentence are kept.
 *   - ADAPTIVE VAD: the speech threshold tracks the noise floor, so it works
 *     for both close-mic speech and external audio (TV, videos, news).
 *   - FORCED CUTS: continuous speech longer than `maxUtteranceMs` is split at
 *     the QUIETEST recent window (least likely to be mid-word) instead of an
 *     arbitrary time boundary.
 *
 * The class is pure (no I/O): callers feed PCM buffers in and receive
 * segment buffers via the `onSegment` callback.
 */

// ─── WAV encoding ────────────────────────────────────────────────────────────

/** Wrap raw PCM16LE mono samples in a 44-byte WAV header. */
export function pcm16ToWav(pcm: Buffer, sampleRate: number): Buffer {
  const header = Buffer.alloc(44);
  const byteRate = sampleRate * 2; // mono, 16-bit
  header.write("RIFF", 0);
  header.writeUInt32LE(36 + pcm.length, 4);
  header.write("WAVE", 8);
  header.write("fmt ", 12);
  header.writeUInt32LE(16, 16); // fmt chunk size
  header.writeUInt16LE(1, 20); // PCM format
  header.writeUInt16LE(1, 22); // mono
  header.writeUInt32LE(sampleRate, 24);
  header.writeUInt32LE(byteRate, 28);
  header.writeUInt16LE(2, 32); // block align
  header.writeUInt16LE(16, 34); // bits per sample
  header.write("data", 36);
  header.writeUInt32LE(pcm.length, 40);
  return Buffer.concat([header, pcm]);
}

// ─── Segmenter ───────────────────────────────────────────────────────────────

export interface PcmSegmenterOptions {
  /** Sample rate of the incoming PCM16LE mono stream. Default: 16000 */
  sampleRate?: number;
  /** Audio kept from before speech onset. Default: 320ms */
  preRollMs?: number;
  /** Sustained silence that ends an utterance. Default: 900ms */
  hangoverMs?: number;
  /** Utterances with less speech than this are discarded as noise. Default: 180ms */
  minSpeechMs?: number;
  /** Continuous speech is force-cut after this long. Default: 7000ms */
  maxUtteranceMs?: number;
  /** Absolute minimum level (dBFS) treated as speech. Default: -48 */
  speechFloorDb?: number;
  /** Speech must exceed the noise floor by this margin (dB). Default: 9 */
  noiseMarginDb?: number;
}

export interface SegmentInfo {
  /** True when the segment was cut mid-speech (maxUtteranceMs reached). */
  forced: boolean;
  /** Audio duration of the emitted segment in milliseconds. */
  durationMs: number;
  /** Milliseconds of the segment classified as speech. */
  speechMs: number;
}

interface Window {
  buf: Buffer;
  db: number;
  speech: boolean;
}

/** VAD analysis window; 20ms is fine-grained enough for word boundaries. */
const WINDOW_MS = 20;
/** Trailing silence kept on an emitted segment (helps STT decode the tail). */
const TAIL_SILENCE_KEEP_MS = 240;
/** Forced cuts search for the quietest window within this recent span. */
const FORCED_CUT_SEARCH_MS = 2000;
/** Forced cuts never split inside the newest audio (still mid-word). */
const FORCED_CUT_GUARD_MS = 200;

export class PcmSegmenter {
  private readonly sampleRate: number;
  private readonly preRollMs: number;
  private readonly hangoverMs: number;
  private readonly minSpeechMs: number;
  private readonly maxUtteranceMs: number;
  private readonly speechFloorDb: number;
  private readonly noiseMarginDb: number;
  private readonly windowBytes: number;

  /** Carry-over for incoming buffers that don't align to whole windows. */
  private remainder: Buffer = Buffer.alloc(0);
  /** Rolling buffer of recent windows while idle (pre-roll source). */
  private preRoll: Window[] = [];
  /** Windows of the in-progress utterance (empty while idle). */
  private active: Window[] = [];
  private activeSpeechMs = 0;
  private silenceRunMs = 0;
  /** EMA noise floor in dBFS; rises slowly, falls quickly. */
  private noiseFloorDb = -65;

  constructor(
    options: PcmSegmenterOptions,
    private readonly onSegment: (pcm: Buffer, info: SegmentInfo) => void,
  ) {
    this.sampleRate = options.sampleRate ?? 16000;
    this.preRollMs = options.preRollMs ?? 320;
    this.hangoverMs = options.hangoverMs ?? 900;
    this.minSpeechMs = options.minSpeechMs ?? 180;
    this.maxUtteranceMs = options.maxUtteranceMs ?? 7000;
    this.speechFloorDb = options.speechFloorDb ?? -48;
    this.noiseMarginDb = options.noiseMarginDb ?? 9;
    this.windowBytes = Math.round((this.sampleRate * WINDOW_MS) / 1000) * 2;
  }

  /** Feed the next PCM16LE buffer from the stream. */
  feed(pcm: Buffer): void {
    const data = this.remainder.length
      ? Buffer.concat([this.remainder, pcm])
      : pcm;
    const wholeWindows = Math.floor(data.length / this.windowBytes);
    const usedBytes = wholeWindows * this.windowBytes;
    this.remainder = data.subarray(usedBytes);

    for (let i = 0; i < wholeWindows; i++) {
      const buf = data.subarray(i * this.windowBytes, (i + 1) * this.windowBytes);
      this.processWindow(buf);
    }
  }

  /**
   * Finalize the in-progress utterance immediately (stream stopped).
   * Emits the tail segment if it contains enough speech.
   */
  flush(): void {
    if (this.remainder.length >= 4) {
      // Analyse the sub-window remainder too — a stream can stop mid-window.
      this.processWindow(this.remainder);
    }
    this.remainder = Buffer.alloc(0);
    this.endUtterance(false);
  }

  /** Drop all buffered state (new session). Keeps the learned noise floor. */
  reset(): void {
    this.remainder = Buffer.alloc(0);
    this.preRoll = [];
    this.active = [];
    this.activeSpeechMs = 0;
    this.silenceRunMs = 0;
  }

  /**
   * Snapshot of the in-progress utterance, for rolling partial transcription.
   * Returns null while idle or when the utterance is still too short.
   */
  getActiveSnapshot(): { pcm: Buffer; durationMs: number; speechMs: number } | null {
    if (!this.active.length || this.activeSpeechMs < this.minSpeechMs) return null;
    const pcm = Buffer.concat(this.active.map((w) => w.buf));
    return {
      pcm,
      durationMs: this.msOf(pcm),
      speechMs: this.activeSpeechMs,
    };
  }

  // ── Internals ────────────────────────────────────────────────────────────

  private msOf(buf: Buffer): number {
    return (buf.length / 2 / this.sampleRate) * 1000;
  }

  private speechThresholdDb(): number {
    return Math.max(this.speechFloorDb, this.noiseFloorDb + this.noiseMarginDb);
  }

  private processWindow(buf: Buffer): void {
    const db = rmsDb(buf);
    const speech = db >= this.speechThresholdDb();
    const win: Window = { buf, db, speech };
    const winMs = this.msOf(buf);

    // Track the noise floor from non-speech windows only, capped so loud
    // continuous audio can never drag the floor up into speech territory.
    if (!speech) {
      const alpha = db < this.noiseFloorDb ? 0.2 : 0.02;
      this.noiseFloorDb = Math.min(
        -35,
        this.noiseFloorDb + alpha * (db - this.noiseFloorDb),
      );
    }

    if (!this.active.length) {
      if (speech) {
        // Speech onset: start the utterance with the pre-roll prepended.
        this.active = [...this.preRoll, win];
        this.preRoll = [];
        this.activeSpeechMs = winMs;
        this.silenceRunMs = 0;
      } else {
        this.preRoll.push(win);
        let total = this.preRoll.reduce((s, w) => s + this.msOf(w.buf), 0);
        while (this.preRoll.length > 1 && total > this.preRollMs) {
          total -= this.msOf(this.preRoll[0].buf);
          this.preRoll.shift();
        }
      }
      return;
    }

    // Utterance in progress: keep everything (speech AND intra-speech silence).
    this.active.push(win);
    if (speech) {
      this.activeSpeechMs += winMs;
      this.silenceRunMs = 0;
    } else {
      this.silenceRunMs += winMs;
    }

    if (this.silenceRunMs >= this.hangoverMs) {
      this.endUtterance(false);
      return;
    }

    const activeMs = this.active.reduce((s, w) => s + this.msOf(w.buf), 0);
    if (activeMs >= this.maxUtteranceMs) {
      this.forceCut();
    }
  }

  private endUtterance(forced: boolean): void {
    if (!this.active.length) return;

    let windows = this.active;
    this.active = [];
    const speechMs = this.activeSpeechMs;
    this.activeSpeechMs = 0;
    this.silenceRunMs = 0;

    // Trim trailing silence beyond a short tail; the trimmed silence seeds
    // the pre-roll so the next onset still has context.
    let tailSilence = 0;
    while (tailSilence < windows.length && !windows[windows.length - 1 - tailSilence].speech) {
      tailSilence++;
    }
    const keepTail = Math.ceil(TAIL_SILENCE_KEEP_MS / WINDOW_MS);
    if (tailSilence > keepTail) {
      const cut = windows.length - (tailSilence - keepTail);
      this.preRoll = windows.slice(cut);
      windows = windows.slice(0, cut);
    }

    if (speechMs < this.minSpeechMs) {
      // Noise blip — don't transcribe it, but keep its tail as pre-roll so a
      // real onset right after still gets its leading audio.
      this.preRoll = [...windows.slice(-Math.ceil(this.preRollMs / WINDOW_MS)), ...this.preRoll];
      return;
    }

    const pcm = Buffer.concat(windows.map((w) => w.buf));
    this.onSegment(pcm, {
      forced,
      durationMs: this.msOf(pcm),
      speechMs,
    });
  }

  /**
   * Continuous speech reached maxUtteranceMs: split at the quietest window in
   * the recent past (least likely to be mid-word). Audio after the cut point
   * becomes the start of the next utterance — nothing is dropped.
   */
  private forceCut(): void {
    const searchWindows = Math.ceil(FORCED_CUT_SEARCH_MS / WINDOW_MS);
    const guardWindows = Math.ceil(FORCED_CUT_GUARD_MS / WINDOW_MS);
    const end = Math.max(1, this.active.length - guardWindows);
    const start = Math.max(1, end - searchWindows);

    let cut = end - 1;
    for (let i = start; i < end; i++) {
      if (this.active[i].db < this.active[cut].db) cut = i;
    }

    const emitted = this.active.slice(0, cut);
    const carried = this.active.slice(cut);
    const emittedSpeechMs = emitted.reduce(
      (s, w) => s + (w.speech ? this.msOf(w.buf) : 0),
      0,
    );

    this.active = carried;
    this.activeSpeechMs = carried.reduce(
      (s, w) => s + (w.speech ? this.msOf(w.buf) : 0),
      0,
    );
    // silenceRunMs continues counting across the cut; recompute from the tail.
    let silence = 0;
    for (let i = carried.length - 1; i >= 0 && !carried[i].speech; i--) {
      silence += this.msOf(carried[i].buf);
    }
    this.silenceRunMs = silence;

    if (emittedSpeechMs < this.minSpeechMs || !emitted.length) return;

    const pcm = Buffer.concat(emitted.map((w) => w.buf));
    this.onSegment(pcm, {
      forced: true,
      durationMs: this.msOf(pcm),
      speechMs: emittedSpeechMs,
    });
  }
}

/** RMS level of a PCM16LE buffer in dBFS. */
function rmsDb(buf: Buffer): number {
  const samples = buf.length >> 1;
  if (!samples) return -Infinity;
  let sumSq = 0;
  for (let i = 0; i < samples; i++) {
    const s = buf.readInt16LE(i * 2) / 32768;
    sumSq += s * s;
  }
  const rms = Math.sqrt(sumSq / samples);
  return 20 * Math.log10(rms + 1e-10);
}