tts-bridge.js

// tts-bridge.js — observes the <llm-chat> shadow DOM, stream-converts the
// current model message to speech, and feeds audio chunks to an audio worklet
// as they're generated.
//
// Pipeline:
//   MediaPipe LlmService.generateResponse() callback fires per partial result
//     └─ llmService.history$ (RxJS BehaviorSubject) emits with cumulative text
//         └─ HistorySubscriber diffs against last-seen, accumulates new chars
//             └─ on sentence boundary (>= MIN_FIRST_CHARS / MIN_CHUNK_CHARS)
//                 └─ enqueue text chunk -> TTS worker (queue serializes them)
//
//   TTS worker (24kHz PCM chunks)
//     └─ tts.onChunk(Float32Array) per generation step
//         └─ batch up to BATCH_MS, then RESAMPLE in JS to AudioContext.sampleRate
//             └─ AudioWorkletNode.port.postMessage({type:'chunk', samples})
//                 └─ worklet plays gapless in real time (FIFO, no fade)
//
// Cross-message behavior:
//   - HistorySubscriber detects "new generation" via key change
//     (history.length / model-message-count / message-index). When user sends
//     a new prompt, the new model message index changes and we reset state:
//       1. Flush any leftover text buffer from the previous generation.
//       2. Clear the bridge text queue.
//       3. Send 'clear' to the worklet (silences queued OLD-gen audio).
//       4. Set `discarding = true` so in-flight TTS chunks for the OLD
//          generation are dropped until its onDone fires.
//   - This guarantees the user hears ONLY the current message's audio.
//
// Why resample in JS (not the worklet):
//   Pocket TTS outputs 24kHz Float32 mono PCM. AudioContext runs at the
//   device's native rate (typically 48kHz — Chrome silently ignores any
//   sampleRate option). Doing the resample in JS once per batch keeps the
//   worklet simple (just plays pre-resampled audio at native rate) and
//   avoids per-sample interpolation overhead in the audio thread.

import { TtsClient } from './tts-client.js?v=6';

const VOICE = 'azelma';
const TEMPERATURE = 0.5;
const SOURCE_RATE = 24000;          // Pocket TTS output rate
const TTS_MARKER_ATTR = 'data-tts-bridge-state';

// Streaming chunk thresholds. The LLM emits text incrementally via the
// LlmService history stream. We accumulate until we have either a complete
// sentence of >= MIN_FIRST_CHARS (first chunk) / MIN_CHUNK_CHARS (subsequent),
// or we hit MAX_CHUNK_CHARS as a safety cap.
const MIN_FIRST_CHARS = 25;
const MIN_CHUNK_CHARS = 15;
const MAX_CHUNK_CHARS = 220;
const SENTENCE_END = /([.!?]+['")\]]*(?:\s+|$))/g;

// Audio batch window: collect diffusion chunks for this many ms, then
// post one combined Float32 to the worklet. Matches tts-site's
// flushAudioBatch pattern. Without batching, the worklet receives one
// tiny chunk per diffusion step, which can produce audible artifacts at
// the worklet's per-chunk processing edges and adds main-thread churn.
const BATCH_MS = 30;

// ---- Status banner ----
function makeBanner() {
    const el = document.createElement('div');
    el.id = 'tts-bridge-banner';
    el.style.cssText = [
        'position:fixed', 'bottom:12px', 'right:12px',
        'z-index:99999', 'padding:8px 14px',
        'font:12px/1.4 -apple-system,system-ui,sans-serif',
        'background:#1976d2', 'color:#fff',
        'border-radius:6px', 'box-shadow:0 2px 8px rgba(0,0,0,0.2)',
        'transition:opacity 0.3s', 'max-width:360px', 'pointer-events:none'
    ].join(';');
    document.body.appendChild(el);
    let hideTimer = null;
    return {
        set(text) {
            el.textContent = text;
            el.style.opacity = '1';
            if (hideTimer) clearTimeout(hideTimer);
            hideTimer = setTimeout(() => { el.style.opacity = '0.45'; }, 3000);
        }
    };
}

// ---- Linear-interpolation resampler: SOURCE_RATE -> targetRate ----
function makeResampler(targetRate) {
    const ratio = SOURCE_RATE / targetRate; // source samples per output sample
    return function resample(input) {
        if (input.length < 2) return input;
        // output length: how many output samples fit in (input.length - 1)
        // source samples at ratio source-per-output?
        const outLen = Math.max(1, Math.floor((input.length - 1) / ratio));
        const out = new Float32Array(outLen);
        for (let i = 0; i < outLen; i++) {
            const srcPos = i * ratio;
            const i0 = Math.floor(srcPos);
            const i1 = Math.min(i0 + 1, input.length - 1);
            const frac = srcPos - i0;
            out[i] = input[i0] * (1 - frac) + input[i1] * frac;
        }
        return out;
    };
}

// =========================================================================
// HistorySubscriber — subscribes to llmChat.llmService.history (an RxJS
// BehaviorSubject of the full chat history). Each emission carries the
// latest snapshot, including the model message whose `text` field grows
// incrementally as MediaPipe streams tokens.
//
// This replaces the old MutationObserver-on-DOM approach which was defeated
// by LitElement's `repeat()` directive recreating `.message-wrapper`
// elements on every text update.
// =========================================================================

function makeHistorySubscriber(getHistory$, onChunk, onNewGen, onGenDone) {
    // stats lives inside main(); we access it via window.ttsBridgeStats which
    // main() publishes before calling us. Without indirection through window,
    // this callback throws "stats is not defined" (stats is not in our closure).
    const stats = window.ttsBridgeStats;
    if (!stats) {
        console.error('[tts-bridge] makeHistorySubscriber called before window.ttsBridgeStats was published');
        return { subscription: null, get genKey() { return null; }, get isBusy() { return false; } };
    }

    let genKey = null;             // identifies which model message we're tracking
    let lastSeenText = '';
    let buffer = '';
    let sentFirst = false;
    let prevDoneGenerating = false;

    function tryExtractChunks() {
        while (true) {
            const minChars = sentFirst ? MIN_CHUNK_CHARS : MIN_FIRST_CHARS;
            let cut = -1;
            SENTENCE_END.lastIndex = 0;
            let m;
            while ((m = SENTENCE_END.exec(buffer)) !== null) {
                const endIdx = m.index + m[0].length;
                if (endIdx >= minChars) {
                    cut = endIdx;
                    break;
                }
            }
            if (cut < 0 && buffer.length >= MAX_CHUNK_CHARS) {
                cut = MAX_CHUNK_CHARS;
            }
            if (cut < 0) return;

            const text = buffer.slice(0, cut).trim();
            buffer = buffer.slice(cut);
            if (!text) continue;

            sentFirst = true;
            console.log(`[tts-bridge] extracted chunk (${text.length} chars): "${text.slice(0, 60)}${text.length > 60 ? '…' : ''}"`);
            onChunk(text);
        }
    }

    const subscription = getHistory$().subscribe((history) => {
        try {
            const modelMsgs = history.filter(m => m.role === 'model');
            if (modelMsgs.length === 0) return;
            const last = modelMsgs[modelMsgs.length - 1];

            // Key = index of the latest model message. When this changes, the
            // user sent a new prompt and a new model message was pushed.
            const newKey = `${history.length}|${modelMsgs.length}|${history.indexOf(last)}`;
            if (newKey !== genKey) {
                // New generation. Flush any leftover buffer from the previous one
                // before resetting state.
                if (genKey !== null && buffer.trim()) {
                    onChunk(buffer.trim());
                }
                genKey = newKey;
                lastSeenText = '';
                buffer = '';
                sentFirst = false;
                prevDoneGenerating = false;
                stats.historyGenStarts++;
                console.log(`[tts-bridge] new generation (key=${genKey})`);
                onNewGen();
            }

            const text = last.text || '';
            const newChars = text.slice(lastSeenText.length);
            lastSeenText = text;
            stats.historyEmits++;
            if (newChars) {
                stats.historyNewChars += newChars.length;
                buffer += newChars;
                tryExtractChunks();
            }

            // Detect generation completion: doneGenerating transitioned false→true.
            if (last.doneGenerating && !prevDoneGenerating) {
                prevDoneGenerating = true;
                if (buffer.trim()) {
                    console.log(`[tts-bridge] doneGenerating; flushing ${buffer.length} char buffer`);
                    onChunk(buffer.trim());
                    buffer = '';
                }
                stats.historyGenCompletes++;
                onGenDone();
            }
        } catch (err) {
            console.error('[tts-bridge] subscribe callback threw:', err);
        }
    });
    console.log(`[tts-bridge] subscribe returned: ${typeof subscription}, historyEmits=${stats.historyEmits}`);

    return {
        subscription,
        get genKey() { return genKey; },
        get isBusy() { return buffer.length > 0 || sentFirst; },
    };
}

// =========================================================================
// Main
// =========================================================================

async function main() {
    await customElements.whenDefined('llm-chat');
    const llmChat = document.querySelector('llm-chat');
    if (!llmChat) {
        console.error('[tts-bridge] no <llm-chat> element found');
        return;
    }

    const banner = makeBanner();
    banner.set('TTS: initializing...');

    // ---- Audio worklet (just plays chunks at native rate) ----
    let audioCtx = null;
    let workletNode = null;
    try {
        audioCtx = new AudioContext();
        await audioCtx.audioWorklet.addModule('./audio-worklet.js?v=fix2');
        workletNode = new AudioWorkletNode(audioCtx, 'streaming-audio');
        workletNode.connect(audioCtx.destination);
        console.log(`[tts-bridge] audio context native rate: ${audioCtx.sampleRate}Hz`);
        banner.set(`TTS: audio ready (${audioCtx.sampleRate}Hz)`);
    } catch (err) {
        console.error('[tts-bridge] AudioContext/worklet init failed:', err);
        banner.set(`TTS audio init failed: ${err.message}`);
        return;
    }

    // Build a resampler to the AudioContext's native rate
    const resample = makeResampler(audioCtx.sampleRate);

    // Resume on first user gesture (autoplay policy)
    const resumeOnGesture = () => {
        if (audioCtx.state === 'suspended') {
            audioCtx.resume().then(() => console.log('[tts-bridge] audio resumed'));
        }
        document.removeEventListener('click', resumeOnGesture);
        document.removeEventListener('keydown', resumeOnGesture);
    };
    document.addEventListener('click', resumeOnGesture);
    document.addEventListener('keydown', resumeOnGesture);

    // ---- TTS client ----
    let voiceReady = false;
    const tts = new TtsClient({
        modelType: 'pocket-tts',
        baseUrl: '',
        workerUrl: './worker.js?v=3',
        onStatus: (text, ready, progress) => {
            if (progress && progress.total) {
                const pct = Math.round((progress.loaded / progress.total) * 100);
                banner.set(`TTS: ${text} ${pct}%`);
            } else {
                banner.set(`TTS: ${text}`);
            }
        },
        onError: (err) => {
            console.error('[tts-bridge]', err);
            banner.set(`TTS error: ${err.message}`);
        },
        onVoiceLoaded: () => {
            voiceReady = true;
            banner.set('TTS: ready');
            pump();
        },
    });

    tts.init()
        .then(() => tts.loadVoice(VOICE))
        .catch(err => {
            console.error('[tts-bridge] init failed:', err);
            banner.set(`TTS init failed: ${err.message}`);
        });

    // ---- Queue + state ----
    const queue = [];
    let busy = false;
    let discarding = false;          // true between active-switch and old-gen-done

    // ---- Audio batch state ----
    // Collect diffusion chunks for BATCH_MS, then post one combined Float32
    // to the worklet. Smaller surface area = fewer per-chunk discontinuities
    // and less main-thread churn.
    let pendingResampled = [];
    let pendingSampleCount = 0;
    let batchTimer = null;
    let lastChunkTime = 0;           // for instrumentation
    let lastChunkGapMs = 0;          // gap since previous chunk

    // ---- Instrumentation ----
    // Exposed via window.ttsBridgeStats so the daemon's snapshot endpoint
    // can read it without modifying the page. Read by GET /snapshot.
    const stats = {
        diffusionChunks: 0,          // total raw chunks from worker
        workletPosts: 0,             // total batches posted to worklet
        workletSamples: 0,           // total samples posted to worklet (post-resample)
        resampleMs: 0,               // total ms spent resampling (accumulated)
        batchFlushMs: 0,             // total ms spent in flushAudioBatch
        lastChunkGapMs: 0,           // gap between most recent two diffusion chunks
        maxChunkGapMs: 0,            // largest gap seen this generation
        sumChunkGapMs: 0,            // sum of gaps this generation
        currentGenStart: 0,          // perf.now() when current tts.generate() began
        currentGenChunks: 0,         // diffusion chunks in current generation
        currentGenSamples24k: 0,     // source-rate samples in current generation
        lastGenDurationMs: 0,        // wall time of previous generation
        lastGenSamples24k: 0,        // source samples produced by previous generation
        discardDrops: 0,             // chunks dropped because discarding=true
        discardStuckResets: 0,       // times we force-reset because discarding > 45s
        textChunksQueued: 0,         // text chunks pushed to TTS queue total
        textChunksGenerated: 0,      // text chunks fully TTS'd total
        historyEmits: 0,             // emissions from llmService.history$
        historyNewChars: 0,          // total incremental chars observed from history$
        historyGenStarts: 0,         // generations detected
        historyGenCompletes: 0,      // generations marked done (doneGenerating=true)
    };
    window.ttsBridgeStats = stats;

    function flushAudioBatch() {
        if (pendingResampled.length === 0) return;
        const t0 = performance.now();
        const totalSamples = pendingSampleCount;
        if (totalSamples === 0) {
            pendingResampled = [];
            batchTimer = null;
            return;
        }
        const pcm = new Float32Array(totalSamples);
        let off = 0;
        for (const c of pendingResampled) { pcm.set(c, off); off += c.length; }
        try {
            workletNode.port.postMessage(
                { type: 'chunk', samples: pcm },
                [pcm.buffer]
            );
            stats.workletPosts++;
            stats.workletSamples += totalSamples;
        } catch (err) {
            console.error('[tts-bridge] worklet postMessage failed:', err);
        }
        pendingResampled = [];
        pendingSampleCount = 0;
        batchTimer = null;
        stats.batchFlushMs += performance.now() - t0;
    }

    function scheduleFlush() {
        if (batchTimer !== null) return;
        batchTimer = setTimeout(flushAudioBatch, BATCH_MS);
    }

    tts.onChunk = (data) => {
        const now = performance.now();
        if (lastChunkTime > 0) {
            const gap = now - lastChunkTime;
            stats.lastChunkGapMs = gap;
            if (gap > stats.maxChunkGapMs) stats.maxChunkGapMs = gap;
            stats.sumChunkGapMs += gap;
        }
        lastChunkTime = now;
        stats.diffusionChunks++;
        stats.currentGenChunks++;
        if (stats.currentGenStart === 0) stats.currentGenStart = now;

        if (discarding) {
            stats.discardDrops++;
            // Safety: if we've been discarding for over 45s, the OLD worker's
            // onDone never fired (worker crash, etc.). Force a reset so the
            // NEW message can be heard.
            if (Date.now() - discardingSince > 45000) {
                console.warn('[tts-bridge] discarding stuck >45s; forcing reset');
                stats.discardStuckResets++;
                discarding = false;
                flushAudioBatch();      // drop any half-batched OLD-gen audio
                try { workletNode.port.postMessage({ type: 'clear' }); } catch (e) {}
                pump();
            }
            return;
        }
        const samples = data instanceof Float32Array ? data : new Float32Array(data);
        if (samples.length < 2) return;
        stats.currentGenSamples24k += samples.length;
        const rt0 = performance.now();
        const resampled = resample(samples);
        stats.resampleMs += performance.now() - rt0;
        if (resampled.length === 0) return;
        pendingResampled.push(resampled);
        pendingSampleCount += resampled.length;
        scheduleFlush();
    };

    tts.onDone = () => {
        // Flush any in-flight batch immediately so we don't lose the tail.
        flushAudioBatch();

        // Record generation metrics
        if (stats.currentGenStart > 0) {
            stats.lastGenDurationMs = performance.now() - stats.currentGenStart;
            stats.lastGenSamples24k = stats.currentGenSamples24k;
        }
        stats.textChunksGenerated++;
        stats.currentGenStart = 0;
        stats.currentGenChunks = 0;
        stats.currentGenSamples24k = 0;
        stats.maxChunkGapMs = 0;
        stats.sumChunkGapMs = 0;
        lastChunkTime = 0;

        busy = false;

        if (discarding) {
            // OLD generation finished. Clear any OLD audio that snuck in,
            // reset state, and start pumping new chunks.
            discarding = false;
            console.log('[tts-bridge] discarded old generation done; ready for new');
            try { workletNode.port.postMessage({ type: 'clear' }); } catch (e) {}
            pump();
            return;
        }

        try { workletNode.port.postMessage({ type: 'finish' }); } catch (e) {}
        banner.set(`TTS: chunk done (queue: ${queue.length})`);
        pump();
    };

    function pump() {
        if (busy || !tts.isReady() || !voiceReady) return;
        if (discarding) return;
        const next = queue.shift();
        if (!next) return;
        busy = true;
        banner.set(`TTS: gen chunk (${next.text.length} chars)`);
        try {
            tts.generate(next.text, TEMPERATURE);
        } catch (err) {
            console.error('[tts-bridge] generate() threw:', err);
            busy = false;
        }
    }

    // ---- Stream management ----
    // Subscribe to llmChat.llmService.history (RxJS BehaviorSubject).
    // Each emission has the latest chat history; the model message's `text`
    // grows incrementally as MediaPipe streams tokens. We diff and extract
    // sentences as soon as they form.
    let discardingSince = 0;
    function onNewGen() {
        // New generation started. If the OLD generation has TTS in flight
        // (busy=true), set discarding so its chunks are dropped, then reset
        // when its onDone fires. If nothing is in flight, just clear state
        // and continue — otherwise discarding would never reset (no onDone
        // would ever fire to clear it).
        const cleared = queue.length;
        queue.length = 0;

        // Drop any in-flight OLD-gen audio batch so it doesn't get posted
        // to the worklet after the new generation starts.
        pendingResampled = [];
        pendingSampleCount = 0;
        if (batchTimer !== null) {
            clearTimeout(batchTimer);
            batchTimer = null;
        }

        try { workletNode.port.postMessage({ type: 'clear' }); } catch (e) {}
        if (busy) {
            discarding = true;
            discardingSince = Date.now();
            console.log(`[tts-bridge] new gen — OLD TTS in flight, discarding=true; cleared ${cleared} pending text chunks`);
        } else {
            discarding = false;
            discardingSince = 0;
            console.log(`[tts-bridge] new gen — no OLD TTS in flight; cleared ${cleared} pending text chunks`);
        }
    }

    function onGenDone() {
        // generation completed; nothing else to do here, pump() handles next item
    }

    // Wait for llmService to be available on the element (set in LlmChat constructor).
    const waitForLlmService = () => {
        if (llmChat.llmService && llmChat.llmService.history) {
            console.log('[tts-bridge] llmService exposed; subscribing to history$');
            makeHistorySubscriber(
                () => llmChat.llmService.history,
                (text) => {
                    stats.textChunksQueued++;
                    queue.push({ text });
                    pump();
                },
                onNewGen,
                onGenDone,
            );
            banner.set('TTS: bridge wired');
        } else {
            requestAnimationFrame(waitForLlmService);
        }
    };
    waitForLlmService();
}

main().catch(err => {
    console.error('[tts-bridge] fatal:', err);
});