au_whisper_web/src/lib/live.ts

/**
 * Live microphone transcription — the segment loop (Decisions.md D-016).
 *
 * Whisper is not a streaming model, so "live" is approximated by splitting
 * the mic signal into speech segments with a simple RMS energy gate and
 * re-transcribing the current segment on a rolling basis (tentative text)
 * until it commits, at which point it gets one final pass and freezes.
 *
 * Everything here is pure and free of browser globals so it runs under
 * `bun test`: audio arrives via `push()` (mic capture lives in lib/mic.ts)
 * and inference is an injected async function (the UI wraps the worker).
 */

import type { TranscriptChunk } from "../worker/protocol";
import { TARGET_SAMPLE_RATE } from "./audio";

/** Root-mean-square level of a sample block. Pure. */
export function rms(samples: Float32Array): number {
  if (samples.length === 0) return 0;
  let sum = 0;
  for (let i = 0; i < samples.length; i++) {
    sum += samples[i]! * samples[i]!;
  }
  return Math.sqrt(sum / samples.length);
}

function concat(parts: readonly Float32Array[], length: number): Float32Array {
  const out = new Float32Array(length);
  let offset = 0;
  for (const part of parts) {
    out.set(part, offset);
    offset += part.length;
  }
  return out;
}

// ---------------------------------------------------------------------------
// Segmenter: RMS energy gate → speech segments
// ---------------------------------------------------------------------------

export interface LiveSegmenterOptions {
  /** Sample rate of the incoming (and outgoing) PCM. Default 16 kHz. */
  sampleRate?: number;
  /** Analysis frame length in seconds. Default 0.032 (~32 ms). */
  frameSeconds?: number;
  /** A frame at/above this RMS counts as speech. Default 0.015. */
  speechRms?: number;
  /** Trailing silence that commits the current segment. Default 0.75 s. */
  commitSilenceSeconds?: number;
  /**
   * Hard cap per segment — must stay under Whisper's 30 s window so a live
   * segment is always a single-chunk transcription. Default 25 s.
   */
  maxSegmentSeconds?: number;
  /** Audio kept from just before speech onset. Default 0.3 s. */
  prerollSeconds?: number;
  /**
   * Segments with less accumulated speech than this are dropped instead of
   * committed — Whisper hallucinates on (near-)silence. Default 0.25 s.
   */
  minSpeechSeconds?: number;
}

export type CommitReason = "silence" | "max-length" | "flush";

/** One committed speech segment, ready for a final transcription pass. */
export interface CommittedSegment {
  audio: Float32Array;
  /** Absolute offset of `audio[0]` in samples since the session started. */
  startSample: number;
  reason: CommitReason;
}

interface Frame {
  samples: Float32Array;
  /** Absolute offset of this frame's first sample. */
  start: number;
}

/**
 * Splits a continuous PCM stream into speech segments. Push mic audio in any
 * block sizes; committed segments come back from `push()`/`flush()`.
 */
export class LiveSegmenter {
  private readonly sampleRate: number;
  private readonly frameSamples: number;
  private readonly speechRms: number;
  private readonly commitSilenceFrames: number;
  private readonly maxSegmentSamples: number;
  private readonly maxPrerollFrames: number;
  private readonly minSpeechSamples: number;

  /** Leftover samples smaller than one analysis frame. */
  private carry = new Float32Array(0);
  /** Absolute offset of the next unprocessed frame's first sample. */
  private absSample = 0;
  /** Recent silence kept while idle, prepended at speech onset. */
  private preroll: Frame[] = [];
  private segment: Frame[] | null = null;
  private segmentStart = 0;
  private segmentSamples = 0;
  private speechSamples = 0;
  private silenceFrames = 0;
  private nextSegmentId = 1;
  private segmentId: number | null = null;

  constructor(options: LiveSegmenterOptions = {}) {
    this.sampleRate = options.sampleRate ?? TARGET_SAMPLE_RATE;
    const frameSeconds = options.frameSeconds ?? 0.032;
    this.frameSamples = Math.max(1, Math.round(frameSeconds * this.sampleRate));
    this.speechRms = options.speechRms ?? 0.015;
    this.commitSilenceFrames = Math.max(
      1,
      Math.ceil((options.commitSilenceSeconds ?? 0.75) / frameSeconds),
    );
    this.maxSegmentSamples = Math.round(
      (options.maxSegmentSeconds ?? 25) * this.sampleRate,
    );
    this.maxPrerollFrames = Math.ceil((options.prerollSeconds ?? 0.3) / frameSeconds);
    this.minSpeechSamples = Math.round(
      (options.minSpeechSeconds ?? 0.25) * this.sampleRate,
    );
  }

  /** Identity of the in-progress segment; null while idle (no speech yet). */
  get currentSegmentId(): number | null {
    return this.segmentId;
  }

  /** Samples accumulated in the in-progress segment (0 while idle). */
  get currentSegmentSamples(): number {
    return this.segment ? this.segmentSamples : 0;
  }

  /** Copy of the in-progress segment's audio, for a tentative pass. */
  currentSegment(): { audio: Float32Array; startSample: number } | null {
    if (!this.segment) return null;
    return {
      audio: concat(
        this.segment.map((frame) => frame.samples),
        this.segmentSamples,
      ),
      startSample: this.segmentStart,
    };
  }

  /** Feed mic samples; returns segments that committed during this block. */
  push(samples: Float32Array): CommittedSegment[] {
    const committed: CommittedSegment[] = [];
    let buffer: Float32Array;
    if (this.carry.length > 0) {
      buffer = concat([this.carry, samples], this.carry.length + samples.length);
    } else {
      buffer = samples;
    }

    let offset = 0;
    while (buffer.length - offset >= this.frameSamples) {
      // Copy the frame — the caller may reuse/transfer its buffer.
      const frame: Frame = {
        samples: buffer.slice(offset, offset + this.frameSamples),
        start: this.absSample,
      };
      offset += this.frameSamples;
      this.absSample += this.frameSamples;
      const segment = this.processFrame(frame, rms(frame.samples) >= this.speechRms);
      if (segment) committed.push(segment);
    }
    this.carry = buffer.slice(offset);
    return committed;
  }

  /** Commit whatever is in progress (session stop). */
  flush(): CommittedSegment | null {
    if (this.segment && this.carry.length > 0) {
      // Sub-frame tail — too short to classify, but part of the speech.
      this.segment.push({ samples: this.carry, start: this.absSample });
      this.segmentSamples += this.carry.length;
      this.absSample += this.carry.length;
      this.carry = new Float32Array(0);
    }
    return this.commit("flush");
  }

  private processFrame(frame: Frame, speech: boolean): CommittedSegment | null {
    if (!this.segment) {
      if (!speech) {
        this.preroll.push(frame);
        if (this.preroll.length > this.maxPrerollFrames) this.preroll.shift();
        return null;
      }
      this.segment = [...this.preroll, frame];
      this.segmentStart = this.segment[0]!.start;
      this.segmentSamples = this.segment.length * this.frameSamples;
      this.speechSamples = this.frameSamples;
      this.silenceFrames = 0;
      this.preroll = [];
      this.segmentId = this.nextSegmentId++;
      return null;
    }

    this.segment.push(frame);
    this.segmentSamples += this.frameSamples;
    if (speech) {
      this.speechSamples += this.frameSamples;
      this.silenceFrames = 0;
    } else {
      this.silenceFrames += 1;
      if (this.silenceFrames >= this.commitSilenceFrames) return this.commit("silence");
    }
    if (this.segmentSamples >= this.maxSegmentSamples) return this.commit("max-length");
    return null;
  }

  private commit(reason: CommitReason): CommittedSegment | null {
    const segment = this.segment;
    if (!segment) return null;
    const result: CommittedSegment | null =
      this.speechSamples >= this.minSpeechSamples
        ? {
            audio: concat(
              segment.map((frame) => frame.samples),
              this.segmentSamples,
            ),
            startSample: this.segmentStart,
            reason,
          }
        : null; // near-silence: drop, never transcribe (D-016)
    this.segment = null;
    this.segmentSamples = 0;
    this.speechSamples = 0;
    this.silenceFrames = 0;
    this.segmentId = null;
    return result;
  }
}

// ---------------------------------------------------------------------------
// Transcriber: schedules inference passes over the segment stream
// ---------------------------------------------------------------------------

export interface LiveTranscribeResult {
  text: string;
  /** Timestamps relative to the transcribed audio's start, Whisper-style. */
  chunks: TranscriptChunk[];
}

/**
 * One inference pass. Takes ownership of `audio` (it may be transferred to
 * the worker). The controller never runs two passes concurrently.
 */
export type LiveTranscribeFn = (audio: Float32Array) => Promise<LiveTranscribeResult>;

export interface LiveUpdate {
  /** Frozen (or provisionally frozen) text of committed segments, joined. */
  committedText: string;
  /** Rolling text of the in-progress segment; replaced every pass. */
  tentativeText: string;
  /** Finalized chunks with absolute session timestamps. */
  chunks: TranscriptChunk[];
  /** Committed segments still waiting for their final pass. */
  pendingFinalPasses: number;
  /** Last inference error, if any (the session keeps going). */
  error: string | null;
}

export interface LiveTranscriberOptions {
  transcribe: LiveTranscribeFn;
  onUpdate: (update: LiveUpdate) => void;
  segmenter?: LiveSegmenterOptions;
  /** Minimum ms between tentative passes. Default 1500. */
  tentativeIntervalMs?: number;
  /** Minimum in-progress segment length before tentative passes. Default 1 s. */
  minTentativeSeconds?: number;
  /** Injectable clock (tests). Default Date.now. */
  now?: () => number;
}

interface SegmentState {
  startSample: number;
  /** Handed to the final pass, then released. */
  audio: Float32Array | null;
  /** Tentative text until the final pass overwrites it. */
  text: string;
  chunks: TranscriptChunk[];
  final: boolean;
}

/**
 * Drives the live session: feeds mic audio to the segmenter, keeps exactly
 * one inference pass in flight (final passes for committed segments first,
 * then throttled tentative refreshes of the current segment), and reports
 * every text change through `onUpdate`. Scheduling is driven entirely by
 * `push()` calls and pass completions — no timers, so it is fully testable.
 */
export class LiveTranscriber {
  private readonly segmenter: LiveSegmenter;
  private readonly transcribe: LiveTranscribeFn;
  private readonly onUpdate: (update: LiveUpdate) => void;
  private readonly tentativeIntervalMs: number;
  private readonly minTentativeSamples: number;
  private readonly sampleRate: number;
  private readonly now: () => number;

  private segments: SegmentState[] = [];
  private tentativeText = "";
  private lastTentative: { segmentId: number; samples: number; at: number } | null =
    null;
  private inFlight = false;
  private stopped = false;
  private error: string | null = null;
  private drainResolvers: (() => void)[] = [];

  constructor(options: LiveTranscriberOptions) {
    this.segmenter = new LiveSegmenter(options.segmenter);
    this.transcribe = options.transcribe;
    this.onUpdate = options.onUpdate;
    this.tentativeIntervalMs = options.tentativeIntervalMs ?? 1500;
    this.sampleRate = options.segmenter?.sampleRate ?? TARGET_SAMPLE_RATE;
    this.minTentativeSamples = Math.round(
      (options.minTentativeSeconds ?? 1) * this.sampleRate,
    );
    this.now = options.now ?? Date.now;
  }

  /** Feed a block of 16 kHz mono mic samples. Ignored after stop(). */
  push(samples: Float32Array): void {
    if (this.stopped) return;
    for (const committed of this.segmenter.push(samples)) {
      this.acceptSegment(committed);
    }
    this.pump();
  }

  /**
   * End the session: commit the in-progress segment, wait for the remaining
   * final passes, and return the full transcript.
   */
  async stop(): Promise<LiveTranscribeResult> {
    if (!this.stopped) {
      this.stopped = true;
      const flushed = this.segmenter.flush();
      if (flushed) this.acceptSegment(flushed);
      // Anything tentative either moved into the flushed segment or belonged
      // to audio the segmenter dropped as near-silence.
      this.tentativeText = "";
      this.pump();
    }
    await this.whenDrained();
    this.emit();
    return { text: this.committedText(), chunks: this.allChunks() };
  }

  private acceptSegment(committed: CommittedSegment): void {
    // The committing segment is the one the tentative text described; carry
    // it as the provisional text until the final pass replaces it.
    this.segments.push({
      startSample: committed.startSample,
      audio: committed.audio,
      text: this.tentativeText,
      chunks: [],
      final: false,
    });
    this.tentativeText = "";
    this.lastTentative = null;
    this.emit();
  }

  /** Single-flight scheduler: pick and start at most one inference pass. */
  private pump(): void {
    if (this.inFlight) return;
    const pending = this.segments.find((segment) => !segment.final);
    if (pending) {
      void this.runFinalPass(pending);
      return;
    }
    if (!this.stopped) {
      this.maybeRunTentativePass();
      return;
    }
    // Stopped and drained — release stop() waiters.
    const resolvers = this.drainResolvers;
    this.drainResolvers = [];
    for (const resolve of resolvers) resolve();
  }

  private async runFinalPass(segment: SegmentState): Promise<void> {
    this.inFlight = true;
    const audio = segment.audio!;
    segment.audio = null; // ownership moves to the transcribe fn
    const offsetSeconds = segment.startSample / this.sampleRate;
    try {
      const result = await this.transcribe(audio);
      segment.text = result.text;
      segment.chunks = result.chunks.map(
        (chunk): TranscriptChunk => ({
          text: chunk.text,
          timestamp: [
            chunk.timestamp[0] + offsetSeconds,
            chunk.timestamp[1] === null ? null : chunk.timestamp[1] + offsetSeconds,
          ],
        }),
      );
    } catch (error) {
      // Keep the provisional (tentative) text rather than losing the segment.
      this.error = error instanceof Error ? error.message : String(error);
    }
    segment.final = true;
    this.inFlight = false;
    this.emit();
    this.pump();
  }

  private maybeRunTentativePass(): void {
    const segmentId = this.segmenter.currentSegmentId;
    if (segmentId === null) return;
    const samples = this.segmenter.currentSegmentSamples;
    if (samples < this.minTentativeSamples) return;
    const last = this.lastTentative;
    if (last && this.now() - last.at < this.tentativeIntervalMs) return;
    if (last && last.segmentId === segmentId && samples <= last.samples) return;

    const snapshot = this.segmenter.currentSegment();
    if (!snapshot) return;
    this.lastTentative = { segmentId, samples, at: this.now() };
    this.inFlight = true;
    void this.transcribe(snapshot.audio)
      .then(
        (result) => {
          // Discard stale passes: the segment may have committed (its final
          // pass supersedes this) or been dropped while we were busy.
          if (this.segmenter.currentSegmentId === segmentId) {
            this.tentativeText = result.text;
          }
        },
        (error: unknown) => {
          this.error = error instanceof Error ? error.message : String(error);
        },
      )
      .then(() => {
        this.inFlight = false;
        this.emit();
        this.pump();
      });
  }

  private committedText(): string {
    return this.segments
      .map((segment) => segment.text.trim())
      .filter(Boolean)
      .join(" ");
  }

  private allChunks(): TranscriptChunk[] {
    return this.segments.flatMap((segment) => segment.chunks);
  }

  private emit(): void {
    this.onUpdate({
      committedText: this.committedText(),
      tentativeText: this.tentativeText,
      chunks: this.allChunks(),
      pendingFinalPasses: this.segments.filter((segment) => !segment.final).length,
      error: this.error,
    });
  }

  private whenDrained(): Promise<void> {
    const drained =
      !this.inFlight && this.segments.every((segment) => segment.final);
    if (drained) return Promise.resolve();
    return new Promise((resolve) => {
      this.drainResolvers.push(resolve);
    });
  }
}