//! Rust Unmute MVP — WebSocket conversational server. //! //! ## ⚠ Superseded by zeroclaw-channel-voice //! //! The canonical voice-loop server lives at //! `~/projects/zeroclaw/crates/zeroclaw-channel-voice` (binary //! `voice_server`). It routes the LLM path through the zeroclaw agent //! runtime — multi-turn history, tools, memory, provider routing — //! instead of the OpenAI-compatible direct path used here. Same WS //! wire protocol so `examples/converse_client.rs` drives both. //! //! This binary is kept for two reasons: //! 1. Reproducibility of the perf benches in `docs/perf_history.md` //! Phase 8.10 (Moonshine TTFA -80 %, mock-LLM 3-turn loop, etc.). //! 2. Standalone use when zeroclaw isn't desired (one-shot, no agent). //! //! For new conversational deployments, prefer: //! //! ```bash //! cd ~/projects/zeroclaw //! cargo run -p zeroclaw-channel-voice --release --bin voice_server -- \ //! --bind 127.0.0.1:18090 //! ``` //! //! ## Original behavior //! //! Loads CSM-1B + Kyutai STT + an OpenAI-compatible LLM at startup. //! Each WebSocket connection is a turn-based voice conversation: //! //! ```text //! Client -> Server binary frames: 16-bit LE PCM @ 24 kHz mono //! Client -> Server text "EOT" : signal end-of-turn //! Server -> Client text {"event":"transcript","text":"..."} //! Server -> Client binary frames: 16-bit LE PCM @ 24 kHz mono //! Server -> Client text {"event":"done","assistant":"..."} //! ``` //! //! Conversation history is kept per-connection. Multiple turns supported //! over a single socket; each turn ends when the client sends "EOT". //! //! ## Status //! //! Half-duplex turn-based. Auto end-of-turn (semantic VAD via the //! Kyutai 1B en/fr `extra_heads` outputs) is the obvious 6c.3 follow-up. //! Barge-in (user interrupts assistant) is also deferred. //! //! ## Usage //! //! ```bash //! export OPENAI_API_KEY=... //! cargo run -p rtx-csm --release --features metal --example converse_server -- \ //! --bind 127.0.0.1:18090 \ //! --llm-base "https://api.openai.com/v1" --llm-model gpt-4o-mini \ //! --system "You are a concise voice assistant. Keep answers under 2 sentences." //! ``` //! //! Drive it with `examples/converse_client.rs`. // Without the `asr` feature, AsrEngine has only the Kyutai variant so // `if let AsrEngine::Kyutai(_) = ...` patterns are technically // irrefutable. With `asr`, they're refutable. Silence the warning so the // code compiles cleanly under both cfgs. #![allow(irrefutable_let_patterns)] use anyhow::{Context, Result}; use async_trait::async_trait; use axum::{ Router, extract::{ State, WebSocketUpgrade, ws::{Message, WebSocket}, }, http::{HeaderMap, StatusCode}, response::IntoResponse, routing::get, }; use clap::Parser; use futures_util::{StreamExt, stream}; use rtx_csm::{ GenerateOptions, Generator, converse::{Converse, ConverseOptions, FlushPolicy}, error::Result as CsmResult, llm_client::{ChatMessage, GenConfig, LlmClient, OpenAiCompatibleClient, TokenStream}, stt::{ASR_DELAY_FRAMES, AsrEvent, SAMPLE_RATE as STT_SR, Stt}, }; use std::net::SocketAddr; use std::path::PathBuf; use std::sync::Arc; use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use tokio::sync::Mutex; const PCM_RATE: u32 = 24_000; #[derive(Debug, Parser)] struct Cli { #[arg(long, default_value = "127.0.0.1:18090")] bind: SocketAddr, #[arg(long, default_value = "https://api.openai.com/v1")] llm_base: String, #[arg(long, default_value = "gpt-4o-mini")] llm_model: String, #[arg(long)] llm_api_key: Option, /// Skip the real LLM; canned echo responses for end-to-end testing /// without an API key. #[arg(long)] mock_llm: bool, #[arg( long, default_value = "You are a concise voice assistant. Reply in one or two short sentences." )] system: String, #[arg(long, default_value_t = 0)] speaker: u32, #[arg(long)] cpu: bool, /// Optional emotion control tag prepended to every TTS turn's text, e.g. /// `--emotion-hint "[whisper]"`. Only meaningful when the loaded LoRA /// adapter was fine-tuned with matching tags (see Phase 12.2 + /// `docs/personal_voice_training_guide.md`); on the un-adapted base it /// just adds extra cosmetic prefix tokens. Acts as a fallback when /// `--reactive-emotion` is also set and detection produces no label. #[arg(long)] emotion_hint: Option, /// Detect the user's emotion from each incoming audio buffer and use /// the resulting tag as the per-turn emotion_hint. Default classifier /// is the Phase 13.3 prosody-rule placeholder; pair with /// `--use-emotion2vec` for the real emotion2vec_plus_base classifier /// (Phase 13.8). Static `--emotion-hint` is the fallback when /// detection returns Neutral. #[arg(long, default_value_t = false)] reactive_emotion: bool, /// With `--reactive-emotion`, swap the prosody-rule classifier for /// the emotion2vec_plus_base candle port. Adds a one-time ~1.1 GB /// download on first run and ~150 ms per turn vs <1 ms for prosody /// rules, in exchange for SOTA-comparable accuracy. #[arg(long, default_value_t = false)] use_emotion2vec: bool, /// When `--reactive-emotion` produces a non-Neutral label, ALSO append /// a per-turn signal to the LLM-facing user message so the LLM's /// response text adapts (not just TTS prosody). Per turn only — the /// annotation is NOT pushed to persistent history, so subsequent /// turns aren't biased by stale signals. Has no effect without /// `--reactive-emotion`. #[arg(long, default_value_t = false)] emotion_aware_llm: bool, /// Bearer token required on the WebSocket Authorization header. If /// unset the server is open (suitable for local dev only). Reads /// RTX_AUTH_TOKEN env var if not provided. #[arg(long)] auth_token: Option, /// Per-connection rate limit: max user audio seconds per 60-second /// rolling window. Excess audio frames are rejected with an error /// event and the connection closes. 0 disables. #[arg(long, default_value_t = 600)] rate_audio_secs_per_min: u32, /// Per-connection rate limit: max turns per 60-second rolling window. /// 0 disables. #[arg(long, default_value_t = 60)] rate_turns_per_min: u32, /// Optional path to a quantized Q8/Q4_K_M GGUF for CSM-1B (output of /// `examples/quantize`). When set, loads the quantized TTS model /// instead of FP — typically ~3× faster on Metal at minimal quality /// loss. Mimi codec stays FP regardless. #[arg(long)] quantized_gguf: Option, /// LoRA voice adapter (safetensors, output of `examples/lora_train` or /// `examples/lora_train_emotional`). Loaded into the backbone (FP or /// quantized) before the conversation loop starts so every assistant /// utterance uses the trained voice / prosody. Combines cleanly with /// `--quantized-gguf` (Phase 12.6): the adapter rides on top of the Q8 /// base in F32 — Phase 12.1's `csm_quantized` LoRA hooks make this work /// with the same shared `apply_lora_adapter` helper used for the FP path. #[arg(long)] lora: Option, /// LoRA rank override. When omitted, auto-detected from adapter metadata /// (Phase 12.5); falls back to 8 for older adapters. #[arg(long)] lora_rank: Option, /// LoRA alpha override. Auto-detected from metadata when omitted. #[arg(long)] lora_alpha: Option, /// Force extended LoRA coverage (q+k+v+output_proj + MLP). Auto-detected /// from metadata when omitted; setting it on a classic q+v adapter just /// allocates extra unused B=0 slots. #[arg(long, default_value_t = false)] extended_lora: bool, /// Path to converted AudioSeal generator safetensors (run /// `examples/audioseal_convert` first). When set together with /// `--watermark-detector`, every assistant utterance is watermarked /// for provenance before being sent over the WebSocket. #[arg(long)] watermark_generator: Option, /// Path to converted AudioSeal detector safetensors. Required for the /// watermarker even when only embedding (the detector is part of /// AudioSealWatermarker construction). #[arg(long)] watermark_detector: Option, /// 16-bit watermark message (decimal or 0xHEX). Defaults to 0. /// Used by both AudioSeal (`--watermark-{generator,detector}`) /// and SilentCipher (`--watermark-silentcipher`). #[arg(long, default_value = "0")] watermark_message: String, /// Use Sesame's actual production watermarker — SilentCipher — /// instead of AudioSeal. Pass a directory containing the three /// released checkpoint files (`enc_c.ckpt`, `dec_c.ckpt`, /// `dec_m_0.ckpt`). The 16 kHz model is the one we wire here; /// the 44.1 kHz model is unsupported. /// /// If you don't have a local copy, leave this empty and pass /// the magic value `hf` to download from `sony/silentcipher` /// at boot. /// /// Mutex with `--watermark-generator` / `--watermark-detector` /// (AudioSeal). Pick one watermarker. #[arg(long)] watermark_silentcipher: Option, /// Use the VAD-enabled STT variant (`kyutai/stt-1b-en_fr-candle`). /// Adds 4 extra prediction heads emitting per-frame probabilities; /// /v1/converse uses head-2 probability > `--vad-threshold` for K /// consecutive frames as auto-end-of-turn (no client "EOT" needed). #[arg(long)] vad: bool, /// VAD probability threshold for auto end-of-turn (default 0.5). #[arg(long, default_value_t = 0.5)] vad_threshold: f32, /// Number of consecutive frames above threshold required to fire /// end-of-turn. Higher = less sensitive to noise. #[arg(long, default_value_t = 4)] vad_consecutive: u32, /// Provider-specific JSON merged into every LLM request body. /// Example: `--llm-extra-body '{"thinking":{"type":"disabled"}}'` /// disables Z.AI glm-4.5/4.6/4.7 reasoning_content traversal so /// `content` emits immediately (the cost is one bypassed thinking /// pass; latency drops 30-40s on glm-4.5 voice loops). Validated /// against Z.AI's coding endpoint; see /// `reference_zai_coding_setup.md` for the full provider matrix. #[arg(long)] llm_extra_body: Option, /// Stream TTS chunks to the client as they're generated rather than /// buffering each full sentence. Cuts time-to-first-audio by ~3-5s /// (the first ~320ms chunk leaves the server within ~500-1000ms of /// the LLM's first sentence boundary). Tradeoff: post-process /// (HPF/LUFS) and watermark are SKIPPED — they require /// full-sentence context. Combining `--stream-tts` with /// `--watermark-*` is rejected at boot. #[arg(long)] stream_tts: bool, /// Use Moonshine-tiny (pure candle, English-only) instead of Kyutai /// STT 1B for transcription. ~34× faster than realtime on M-series /// Metal (10 s of audio → ~300 ms transcribe). Batch-only — fires /// the transcript at EOT, not incrementally. Loses semantic VAD. /// /// **Pure-Rust win**: unlike `--whisper`, Moonshine integrates /// cleanly without ggml/protobuf runtime conflicts. No feature /// flag needed (ships in the default build). /// /// English-only — for multilingual deploys, stick with Kyutai. #[arg(long)] moonshine: bool, /// Use Whisper-tiny (whisper-rs / whisper.cpp) instead of Kyutai STT /// 1B for transcription. ~50× faster on Metal in isolation (10 s of /// audio → 0.2 s transcribe) but batch-only — transcript fires at /// EOT, not incrementally. Loses semantic VAD (`--vad`). /// /// **WARNING — perf regression**: linking whisper-rs's C++ runtime /// into the same binary as candle/CSM costs ~2-3× across all CSM /// inference (recv_phase, tts_per_utterance, total_turn) even when /// this flag isn't set. We don't yet understand the conflict /// (BLAS / threadpool / Metal init contention is suspected). For /// production deploys, build WITHOUT `--features asr` and accept /// Kyutai STT's 1× realtime cost; or run Whisper as a sidecar /// process via IPC (not yet implemented). /// /// Requires the crate built with `--features asr-metal` (or `asr`). #[arg(long)] whisper: bool, /// In streaming mode, the number of 80 ms Mimi frames per chunk. /// Default 2 = 160 ms. /// /// Phase 9.2 measurement (Q8 + stream + Moonshine + mock LLM): /// chunk_frames=4 llm_to_first_audio 533 ms e2e 939 ms /// chunk_frames=2 llm_to_first_audio 374 ms e2e 784 ms ← default /// chunk_frames=1 llm_to_first_audio 283 ms e2e 664 ms /// Per-utterance TTS gen is comparable across all three (~5.7-5.9 s /// for the 4-sentence mock LLM), so smaller chunks don't add /// meaningful decode overhead. Drop to 1 for tightest TTFA at the /// cost of slightly more per-chunk send overhead. #[arg(long, default_value_t = 2)] stream_chunk_frames: usize, /// Enable an RMS-amplitude VAD pre-filter on STT work. Each /// incoming 200 ms binary chunk gets its RMS computed; chunks /// below `--vad-gate-threshold` skip the (expensive) Kyutai /// step_pcm call entirely. Pure-Rust energy detector — no external /// runtime, no protobuf conflict (the original plan was Silero V5 /// via `ort`, but `ort` 3.21 protobuf clashes with sentencepiece /// 3.14; energy VAD sidesteps that). /// /// Real-world voice-agent audio is 30-50% silence (pauses, /// breathing, room tone), giving proportional `recv_phase` /// reduction. Quality tradeoff: energy VAD misses quiet speech /// (whispering, distant). Acceptable for mic-distance loops. /// /// Orthogonal to `--vad` (Kyutai's semantic-VAD-for-EOT). #[arg(long)] vad_gate: bool, /// RMS threshold for `--vad-gate` in [0, 1] (post-normalization /// f32 PCM range). 0.01 default; lower = more chunks pass through /// STT (less aggressive). Tune per microphone gain. #[arg(long, default_value_t = 0.01)] vad_gate_threshold: f32, } /// Multi-sentence mock LLM. Reads the most recent user message and /// returns a fixed multi-sentence acknowledgment with a small delay /// between sentence boundaries so the streaming pipeline runs long /// enough to be barge-in-testable. struct MockLlm; #[async_trait] impl LlmClient for MockLlm { async fn generate_stream( &self, messages: Vec, _config: GenConfig, ) -> CsmResult { let user_text = messages .iter() .rev() .find_map(|m| { if matches!(m.role, rtx_csm::llm_client::Role::User) { Some(m.content.clone()) } else { None } }) .unwrap_or_default(); let response = format!( "I heard you. You said: {}. That is interesting. Tell me more about it.", user_text ); // Stream tokens one whitespace-split chunk at a time, with a small // sleep between chunks so the pipeline takes a few seconds end to // end (lets barge-in tests fire mid-stream). let chunks: Vec = response .split_inclusive(' ') .map(|s| s.to_string()) .collect(); let s = stream::iter(chunks).then(|s| async move { tokio::time::sleep(std::time::Duration::from_millis(50)).await; Ok::<_, rtx_csm::CsmError>(s) }); Ok(Box::pin(s)) } } enum AnyLlm { Real(OpenAiCompatibleClient), Mock(MockLlm), } #[async_trait] impl LlmClient for AnyLlm { async fn generate_stream( &self, messages: Vec, config: GenConfig, ) -> CsmResult { match self { AnyLlm::Real(c) => c.generate_stream(messages, config).await, AnyLlm::Mock(c) => c.generate_stream(messages, config).await, } } } #[derive(Default)] struct Metrics { /// Successful turns completed. turns_total: AtomicU64, /// Errors raised during a turn (auth, STT, LLM, TTS). errors_total: AtomicU64, /// Connections opened. connections_total: AtomicU64, /// Connections currently active. connections_active: AtomicU64, /// Sum of STT latencies in ms (for avg = sum / count). stt_latency_ms_sum: AtomicU64, stt_latency_ms_count: AtomicU64, /// Sum of TTS latencies in ms (per-utterance, summed across utterances). tts_latency_ms_sum: AtomicU64, tts_latency_ms_count: AtomicU64, /// End-to-end turn latency (audio-in → first-audio-out): sum + count. e2e_first_ms_sum: AtomicU64, e2e_first_ms_count: AtomicU64, /// Receive phase: turn start → EOT (or VAD auto-EOT) received. recv_phase_ms_sum: AtomicU64, recv_phase_ms_count: AtomicU64, /// LLM-to-first-audio: from conv_fut start (after STT done) to first /// PCM chunk reaching the pump. Bundles LLM stream + first sentence /// TTS gen. llm_to_first_audio_ms_sum: AtomicU64, llm_to_first_audio_ms_count: AtomicU64, /// Full conversation phase: conv_fut start → all sentences TTS'd. conv_total_ms_sum: AtomicU64, conv_total_ms_count: AtomicU64, /// Total turn wall time: turn start → done event sent. total_turn_ms_sum: AtomicU64, total_turn_ms_count: AtomicU64, /// Reactive emotion observability (Phase 13.6). Each turn that ran /// the detector increments `_calls_total`; the per-label counters /// add to one of the five labels. `emotion_aware_llm_applied_total` /// counts turns where the LLM-history augmentation actually fired /// (only when the label is non-Neutral AND --emotion-aware-llm is on). reactive_emotion_calls_total: AtomicU64, reactive_emotion_neutral_total: AtomicU64, reactive_emotion_calm_total: AtomicU64, reactive_emotion_sad_total: AtomicU64, reactive_emotion_angry_total: AtomicU64, reactive_emotion_excited_total: AtomicU64, /// Phase 13.10 raw emotion2vec classes — added when the 9→5 fold /// was removed. The original 5 buckets above are still emitted /// (prosody-rule SER produces those) so these are emotion2vec-only /// in practice. reactive_emotion_disgusted_total: AtomicU64, reactive_emotion_fearful_total: AtomicU64, reactive_emotion_happy_total: AtomicU64, reactive_emotion_surprised_total: AtomicU64, reactive_emotion_unk_total: AtomicU64, emotion_aware_llm_applied_total: AtomicU64, } /// ASR backend wrapped behind a single Mutex so the receive loop can /// branch on its kind. Whisper variant is feature-gated; without `asr` /// at compile time only Kyutai + Moonshine are available. /// /// Moonshine is pure candle (no external runtime), so it ships in the /// default build alongside Kyutai. English-only. enum AsrEngine { Kyutai(Stt), Moonshine(MoonshineAsr), #[cfg(feature = "asr")] Whisper(rtx_csm::asr::WhisperAsr), } /// Bundle of Moonshine encoder + decoder + tokenizer + config so the /// receive-loop's batch transcribe call has everything it needs. struct MoonshineAsr { encoder: rtx_csm::moonshine::Encoder, decoder: rtx_csm::moonshine::Decoder, tokenizer: tokenizers::Tokenizer, cfg: rtx_csm::moonshine::MoonshineConfig, device: candle_core::Device, } impl MoonshineAsr { /// Resample 24 kHz user audio to 16 kHz, encode, greedy decode, /// detokenize. Returns the transcript text. Batch-only (no /// streaming, no per-frame VAD). fn transcribe_24k(&self, samples_24k: &[f32]) -> anyhow::Result { let samples_16k = rtx_csm::audio_io::resample(samples_24k, 24_000, 16_000) .map_err(|e| anyhow::anyhow!("resample 24k -> 16k: {e}"))?; let pcm = candle_core::Tensor::from_vec( samples_16k.clone(), (1, 1, samples_16k.len()), &self.device, )?; let enc = self.encoder.forward(&pcm)?; let token_ids = self.decoder.generate_cached( &enc, &self.cfg, self.cfg.max_position_embeddings.min(180), )?; let text = self .tokenizer .decode(&token_ids, true) .map_err(|e| anyhow::anyhow!("detok: {e}"))?; Ok(text) } } struct Shared { generator: Mutex, stt: Mutex, llm: AnyLlm, system_prompt: String, speaker: u32, auth_token: Option, rate_audio_secs_per_min: u32, rate_turns_per_min: u32, /// VAD: when Some, the WS receive loop monitors per-frame end-of-turn /// probability and auto-fires EOT when prs[2][0] > threshold for /// `consecutive` frames in a row. vad_eot: Option<(f32, u32)>, /// Provider-specific JSON merged into every LLM request body. None /// means the OpenAI-schema fields are sent as-is. llm_extra_body: serde_json::Map, /// When true, conv_fut uses Converse::run_streaming for low-TTFA /// chunked PCM delivery. Skips post-process + watermark. stream_tts: bool, /// Mimi frames per stream chunk (4 = 320ms default). stream_chunk_frames: usize, /// True when the ASR backend is batch-only (Whisper or Moonshine). /// The receive loop skips parallel STT in this case and runs one /// transcribe call at EOT. (Kyutai uses incremental step_pcm.) batch_asr: bool, /// Optional VAD gate. When `Some`, each incoming binary chunk is /// classified speech/silence before being sent to STT; silence /// chunks skip step_pcm. Behind a Mutex because is_speech() takes /// `&mut self` to update internal counters. vad_gate: Option>, /// Optional control-token tag prepended to every TTS turn's text. Only /// useful when the loaded LoRA was fine-tuned with the same tag (Phase 12.2 /// plumbing); on the un-adapted base it is a no-op cosmetic prefix. emotion_hint: Option, /// Pre-built reactive-emotion classifier shared across turns. Boxed /// `dyn` so the same call site handles both ProsodyDetector (cheap, /// zero-init) and Emotion2Vec (loaded once on boot, ~150 ms/turn). /// `Some` ⇔ `--reactive-emotion` was set; `None` disables the path. emotion_detector: Option>, /// When true (and `reactive_emotion` is also true), append a per-turn /// emotion-signal annotation to the LLM-facing user message so the /// response text adapts in addition to the TTS prosody. Per-turn, /// not pushed to persistent history (Phase 13.5). emotion_aware_llm: bool, metrics: Metrics, } /// Sliding-window rate limiter (60-second window). Stores timestamps /// (Instant) of recent events; trims to the window on each `try_consume`. struct RateBucket { /// (timestamp, weight) for each event in the window. events: std::collections::VecDeque<(std::time::Instant, u32)>, weight_in_window: u32, limit: u32, window: std::time::Duration, } impl RateBucket { fn new(limit: u32) -> Self { Self { events: std::collections::VecDeque::new(), weight_in_window: 0, limit, window: std::time::Duration::from_secs(60), } } /// Returns true if `weight` units fit within the budget; records the /// event if so. fn try_consume(&mut self, weight: u32) -> bool { if self.limit == 0 { return true; // disabled } let now = std::time::Instant::now(); while let Some((t, w)) = self.events.front() { if now.duration_since(*t) > self.window { self.weight_in_window = self.weight_in_window.saturating_sub(*w); self.events.pop_front(); } else { break; } } if self.weight_in_window + weight > self.limit { return false; } self.weight_in_window += weight; self.events.push_back((now, weight)); true } } #[tokio::main] async fn main() -> Result<()> { tracing_subscriber::fmt().init(); tracing::warn!( "converse_server is superseded by zeroclaw-channel-voice (see \ ~/projects/zeroclaw/crates/zeroclaw-channel-voice). This binary \ remains for reproducibility of docs/perf_history.md Phase 8.10 \ benches and standalone (no-agent) use." ); let cli = Cli::parse(); let llm = if cli.mock_llm { tracing::info!("LLM: mock (echoes user transcript)"); AnyLlm::Mock(MockLlm) } else { let api_key = cli .llm_api_key .clone() .or_else(|| std::env::var("OPENAI_API_KEY").ok()) .ok_or_else(|| { anyhow::anyhow!("set --llm-api-key or OPENAI_API_KEY (or pass --mock-llm)") })?; AnyLlm::Real(OpenAiCompatibleClient::new( &cli.llm_base, api_key, &cli.llm_model, )) }; let device = if cli.cpu { candle_core::Device::Cpu } else { Generator::default_device()? }; tracing::info!("device: {device:?}"); let mut generator = if let Some(gguf) = cli.quantized_gguf.as_ref() { tracing::info!("loading CSM-1B (quantized, {})...", gguf.display()); Generator::load_csm_1b_quantized(gguf, &device, /* enable_cfg */ false)? } else { tracing::info!("loading CSM-1B (FP)..."); Generator::load_csm_1b(&device)? }; if let Some(lora_path) = cli.lora.as_ref() { let extended_override = if cli.extended_lora { Some(true) } else { None }; rtx_csm::training::apply_lora_adapter( &mut generator, lora_path, cli.lora_rank, cli.lora_alpha, extended_override, &device, )?; } if let (Some(gen_path), Some(det_path)) = ( cli.watermark_generator.as_ref(), cli.watermark_detector.as_ref(), ) { let msg_str = cli.watermark_message.trim(); let message: u16 = if let Some(rest) = msg_str .strip_prefix("0x") .or_else(|| msg_str.strip_prefix("0X")) { u16::from_str_radix(rest, 16)? } else { msg_str.parse::()? }; let gen_vb = unsafe { candle_nn::VarBuilder::from_mmaped_safetensors( &[gen_path], candle_core::DType::F32, &device, ) }?; let det_vb = unsafe { candle_nn::VarBuilder::from_mmaped_safetensors( &[det_path], candle_core::DType::F32, &device, ) }?; let inner = rtx_csm::AudioSealWatermarker::from_var_builders( gen_vb, det_vb, device.clone(), message, )?; // CSM is 24 kHz; AudioSeal is 16 kHz. Wrap so the watermarker // resamples around its forward pass. let wm = rtx_csm::ResampledWatermarker::new(inner, generator.config.sample_rate, 16_000); generator.set_watermarker(Box::new(wm)); tracing::info!( "watermarker installed (message=0x{:04X}, model 16 kHz, output {} Hz)", message, generator.config.sample_rate, ); } else if cli.watermark_generator.is_some() || cli.watermark_detector.is_some() { anyhow::bail!( "--watermark-generator and --watermark-detector must be set together \ (AudioSealWatermarker requires both)." ); } // SilentCipher (Sesame's actual production watermarker). Mutex // with --watermark-generator/-detector (AudioSeal): the Generator // can only carry one watermarker. if cli.watermark_silentcipher.is_some() && (cli.watermark_generator.is_some() || cli.watermark_detector.is_some()) { anyhow::bail!( "--watermark-silentcipher and --watermark-generator/-detector are mutually \ exclusive — pick one watermarker (Generator only carries one)." ); } if let Some(silent_path) = cli.watermark_silentcipher.as_ref() { let msg_str = cli.watermark_message.trim(); let message: u16 = if let Some(rest) = msg_str .strip_prefix("0x") .or_else(|| msg_str.strip_prefix("0X")) { u16::from_str_radix(rest, 16)? } else { msg_str.parse::()? }; // Resolve checkpoint paths: either an `hf` magic value (download // from sony/silentcipher), or a local directory containing the // three .ckpt files. let (enc_c, dec_c, dec_m_0) = if silent_path == "hf" { tracing::info!("downloading SilentCipher 16 kHz from sony/silentcipher..."); let api = hf_hub::api::sync::Api::new()?; let repo = api.model("sony/silentcipher".to_string()); let dir = "16_khz/97561_iteration"; let enc = repo.get(&format!("{dir}/enc_c.ckpt"))?; let dec = repo.get(&format!("{dir}/dec_c.ckpt"))?; let dem = repo.get(&format!("{dir}/dec_m_0.ckpt"))?; (enc, dec, dem) } else { let dir = std::path::PathBuf::from(silent_path); ( dir.join("enc_c.ckpt"), dir.join("dec_c.ckpt"), dir.join("dec_m_0.ckpt"), ) }; let cfg_sc = rtx_csm::silentcipher::SilentCipherConfig::sixteen_khz(); let inner = rtx_csm::silentcipher::SilentCipherWatermarker::from_ckpts( cfg_sc, &enc_c, &dec_c, &dec_m_0, &device, )?; let wm = rtx_csm::silentcipher::SilentCipherWatermark::new(inner, message as u32); // CSM is 24 kHz; SilentCipher 16 kHz. Wrap with the existing // ResampledWatermarker so callers feed 24 kHz directly. let wrapped = rtx_csm::ResampledWatermarker::new(wm, generator.config.sample_rate, 16_000); generator.set_watermarker(Box::new(wrapped)); tracing::info!( "SilentCipher watermarker installed (message=0x{:04X}, model 16 kHz, output {} Hz)", message, generator.config.sample_rate, ); } if cli.stream_tts && generator.watermarker.is_some() { anyhow::bail!( "--stream-tts cannot combine with --watermark-* (watermark needs \ full-sentence context). Pick one: low-latency streaming OR \ provenance-marked audio." ); } if cli.moonshine && cli.whisper { anyhow::bail!("--moonshine and --whisper are mutually exclusive (pick one ASR backend)."); } if cli.moonshine && cli.vad { anyhow::bail!( "--moonshine and --vad cannot combine: Moonshine has no per-frame probability \ output." ); } let stt: AsrEngine = if cli.moonshine { tracing::info!("loading Moonshine-tiny (pure candle, English-only)..."); let api = hf_hub::api::sync::Api::new()?; let repo = api.model("UsefulSensors/moonshine-tiny".to_string()); let weights = repo.get("model.safetensors")?; let tok_path = repo.get("tokenizer.json")?; let cfg = rtx_csm::moonshine::MoonshineConfig::tiny(); let (encoder, decoder) = rtx_csm::moonshine::load_full(&weights, &device, &cfg)?; let tokenizer = rtx_csm::moonshine::load_tokenizer(&tok_path) .map_err(|e| anyhow::anyhow!("moonshine tokenizer: {e}"))?; AsrEngine::Moonshine(MoonshineAsr { encoder, decoder, tokenizer, cfg, device: device.clone(), }) } else if cli.whisper { if cli.vad { anyhow::bail!( "--whisper and --vad cannot combine: Whisper has no per-frame probability \ output. Pick one." ); } #[cfg(feature = "asr")] { tracing::info!("loading Whisper-tiny (whisper-rs)..."); AsrEngine::Whisper(rtx_csm::asr::WhisperAsr::load_default()?) } #[cfg(not(feature = "asr"))] { anyhow::bail!( "--whisper requires the `asr-metal` (or `asr`) feature. Rebuild with \ `cargo build --features metal,asr-metal --example converse_server`." ); } } else if cli.vad { tracing::info!("loading Kyutai STT 1B en/fr-candle with VAD (~3 GB)..."); AsrEngine::Kyutai(Stt::load_default_with_vad(&device)?) } else { tracing::info!("loading Kyutai STT 1B en/fr (~3 GB)..."); AsrEngine::Kyutai(Stt::load_default(&device)?) }; tracing::info!("models loaded"); // Warm up the Generator: a longer throwaway generate() so the full // Metal kernel set (backbone + depth decoder + Mimi decode + every // sampler path) compiles up-front. Empirically eliminates ~4-6 s // from the first-sentence TTS gen latency. Cost is paid once at // boot. Phase 8.2 raised this from 200 ms to 2000 ms (~25 frames) // because at 200 ms only the early code paths fired — variable // first-sentence cost still showed up on first user turn. let warmup_t = std::time::Instant::now(); { let opts = GenerateOptions { max_audio_ms: 2000, seed: 0, ..GenerateOptions::default() }; match generator.generate( "Warming up the speech model with a slightly longer prompt.", cli.speaker, &[], opts, ) { Ok(_) => { tracing::info!( "generator warm-up complete ({}ms)", warmup_t.elapsed().as_millis() ); } Err(e) => { tracing::warn!( "generator warm-up failed (non-fatal, first turn will be slower): {e}" ); } } } let auth_token = cli .auth_token .or_else(|| std::env::var("RTX_AUTH_TOKEN").ok()); if auth_token.is_none() { tracing::warn!( "auth disabled — anyone who can reach this address can use the service. \ Set --auth-token or RTX_AUTH_TOKEN for production." ); } let llm_extra_body = match cli.llm_extra_body.as_deref() { Some(s) => { let parsed: serde_json::Value = serde_json::from_str(s) .with_context(|| format!("--llm-extra-body must be a JSON object: {s}"))?; match parsed { serde_json::Value::Object(m) => m, other => anyhow::bail!("--llm-extra-body must be a JSON object, got {}", other), } } None => serde_json::Map::new(), }; if !llm_extra_body.is_empty() { tracing::info!( "LLM extra_body: {} field(s) — merged into every request", llm_extra_body.len() ); } // Build the reactive-emotion classifier once at boot. Cheap for // ProsodyDetector; pays the ~1.1 GB download + Metal allocation // upfront for Emotion2Vec so per-turn cost is just ~150 ms forward. let emotion_detector: Option> = if cli.reactive_emotion { if cli.use_emotion2vec { let api = hf_hub::api::sync::Api::new()?; let path = api .model("emotion2vec/emotion2vec_plus_base".into()) .get("model.pt") .context("download emotion2vec_plus_base/model.pt")?; let model = rtx_csm::emotion2vec::Emotion2Vec::load_from_pickle(&path, &device)?; tracing::info!("loaded emotion2vec_plus_base from {}", path.display()); Some(Box::new(model)) } else { Some(Box::new(rtx_csm::ser::ProsodyDetector::default())) } } else { None }; let shared = Arc::new(Shared { generator: Mutex::new(generator), stt: Mutex::new(stt), llm, system_prompt: cli.system, speaker: cli.speaker, auth_token, rate_audio_secs_per_min: cli.rate_audio_secs_per_min, rate_turns_per_min: cli.rate_turns_per_min, vad_eot: if cli.vad { Some((cli.vad_threshold, cli.vad_consecutive)) } else { None }, llm_extra_body, stream_tts: cli.stream_tts, stream_chunk_frames: cli.stream_chunk_frames.max(1), batch_asr: cli.whisper || cli.moonshine, vad_gate: if cli.vad_gate { tracing::info!( "energy VAD gate enabled (rms_threshold={:.4})", cli.vad_gate_threshold ); Some(Mutex::new(rtx_csm::stt::VadGate::new( cli.vad_gate_threshold, ))) } else { None }, emotion_hint: cli.emotion_hint.clone(), emotion_detector, emotion_aware_llm: cli.emotion_aware_llm, metrics: Metrics::default(), }); let app = Router::new() .route("/health", get(|| async { "ok" })) .route("/metrics", get(metrics_handler)) .route("/v1/converse", get(ws_handler)) .with_state(shared); let listener = tokio::net::TcpListener::bind(&cli.bind).await?; tracing::info!("listening on http://{}/v1/converse (WebSocket)", cli.bind); let shutdown = shutdown_signal(); axum::serve(listener, app) .with_graceful_shutdown(shutdown) .await?; tracing::info!("server stopped"); Ok(()) } /// SIGINT (Ctrl+C) or SIGTERM stops accepting new connections; in-flight /// turns finish naturally before the server exits. async fn shutdown_signal() { let ctrl_c = async { let _ = tokio::signal::ctrl_c().await; }; #[cfg(unix)] let term = async { let mut sig = tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()) .expect("install SIGTERM handler"); sig.recv().await; }; #[cfg(not(unix))] let term = std::future::pending::<()>(); tokio::select! { _ = ctrl_c => tracing::info!("received SIGINT, shutting down gracefully"), _ = term => tracing::info!("received SIGTERM, shutting down gracefully"), } } /// Prometheus-style /metrics endpoint. Counter + summary lines. async fn metrics_handler(State(shared): State>) -> impl IntoResponse { let m = &shared.metrics; let stt_count = m.stt_latency_ms_count.load(Ordering::Relaxed).max(1); let tts_count = m.tts_latency_ms_count.load(Ordering::Relaxed).max(1); let e2e_count = m.e2e_first_ms_count.load(Ordering::Relaxed).max(1); let recv_count = m.recv_phase_ms_count.load(Ordering::Relaxed).max(1); let llm_to_audio_count = m.llm_to_first_audio_ms_count.load(Ordering::Relaxed).max(1); let conv_count = m.conv_total_ms_count.load(Ordering::Relaxed).max(1); let total_count = m.total_turn_ms_count.load(Ordering::Relaxed).max(1); let body = format!( "# TYPE rtx_csm_turns_total counter\n\ rtx_csm_turns_total {}\n\ # TYPE rtx_csm_errors_total counter\n\ rtx_csm_errors_total {}\n\ # TYPE rtx_csm_connections_total counter\n\ rtx_csm_connections_total {}\n\ # TYPE rtx_csm_connections_active gauge\n\ rtx_csm_connections_active {}\n\ # TYPE rtx_csm_stt_latency_ms_avg gauge\n\ rtx_csm_stt_latency_ms_avg {}\n\ # TYPE rtx_csm_tts_latency_ms_avg gauge\n\ rtx_csm_tts_latency_ms_avg {}\n\ # TYPE rtx_csm_e2e_first_audio_ms_avg gauge\n\ rtx_csm_e2e_first_audio_ms_avg {}\n\ # TYPE rtx_csm_recv_phase_ms_avg gauge\n\ rtx_csm_recv_phase_ms_avg {}\n\ # TYPE rtx_csm_llm_to_first_audio_ms_avg gauge\n\ rtx_csm_llm_to_first_audio_ms_avg {}\n\ # TYPE rtx_csm_conv_total_ms_avg gauge\n\ rtx_csm_conv_total_ms_avg {}\n\ # TYPE rtx_csm_total_turn_ms_avg gauge\n\ rtx_csm_total_turn_ms_avg {}\n\ # TYPE rtx_csm_reactive_emotion_calls_total counter\n\ rtx_csm_reactive_emotion_calls_total {}\n\ # TYPE rtx_csm_reactive_emotion_total counter\n\ rtx_csm_reactive_emotion_total{{label=\"neutral\"}} {}\n\ rtx_csm_reactive_emotion_total{{label=\"calm\"}} {}\n\ rtx_csm_reactive_emotion_total{{label=\"sad\"}} {}\n\ rtx_csm_reactive_emotion_total{{label=\"angry\"}} {}\n\ rtx_csm_reactive_emotion_total{{label=\"excited\"}} {}\n\ rtx_csm_reactive_emotion_total{{label=\"disgusted\"}} {}\n\ rtx_csm_reactive_emotion_total{{label=\"fearful\"}} {}\n\ rtx_csm_reactive_emotion_total{{label=\"happy\"}} {}\n\ rtx_csm_reactive_emotion_total{{label=\"surprised\"}} {}\n\ rtx_csm_reactive_emotion_total{{label=\"unk\"}} {}\n\ # TYPE rtx_csm_emotion_aware_llm_applied_total counter\n\ rtx_csm_emotion_aware_llm_applied_total {}\n", m.turns_total.load(Ordering::Relaxed), m.errors_total.load(Ordering::Relaxed), m.connections_total.load(Ordering::Relaxed), m.connections_active.load(Ordering::Relaxed), m.stt_latency_ms_sum.load(Ordering::Relaxed) / stt_count, m.tts_latency_ms_sum.load(Ordering::Relaxed) / tts_count, m.e2e_first_ms_sum.load(Ordering::Relaxed) / e2e_count, m.recv_phase_ms_sum.load(Ordering::Relaxed) / recv_count, m.llm_to_first_audio_ms_sum.load(Ordering::Relaxed) / llm_to_audio_count, m.conv_total_ms_sum.load(Ordering::Relaxed) / conv_count, m.total_turn_ms_sum.load(Ordering::Relaxed) / total_count, m.reactive_emotion_calls_total.load(Ordering::Relaxed), m.reactive_emotion_neutral_total.load(Ordering::Relaxed), m.reactive_emotion_calm_total.load(Ordering::Relaxed), m.reactive_emotion_sad_total.load(Ordering::Relaxed), m.reactive_emotion_angry_total.load(Ordering::Relaxed), m.reactive_emotion_excited_total.load(Ordering::Relaxed), m.reactive_emotion_disgusted_total.load(Ordering::Relaxed), m.reactive_emotion_fearful_total.load(Ordering::Relaxed), m.reactive_emotion_happy_total.load(Ordering::Relaxed), m.reactive_emotion_surprised_total.load(Ordering::Relaxed), m.reactive_emotion_unk_total.load(Ordering::Relaxed), m.emotion_aware_llm_applied_total.load(Ordering::Relaxed), ); ( StatusCode::OK, [( axum::http::header::CONTENT_TYPE, "text/plain; version=0.0.4", )], body, ) } async fn ws_handler( ws: WebSocketUpgrade, State(shared): State>, headers: HeaderMap, ) -> impl IntoResponse { if let Some(expected) = shared.auth_token.as_ref() { let supplied = headers .get("authorization") .and_then(|v| v.to_str().ok()) .and_then(|h| h.strip_prefix("Bearer ")) .unwrap_or(""); if supplied != expected { shared.metrics.errors_total.fetch_add(1, Ordering::Relaxed); return (StatusCode::UNAUTHORIZED, "unauthorized").into_response(); } } ws.on_upgrade(move |socket| handle_connection(socket, shared)) .into_response() } async fn handle_connection(mut socket: WebSocket, shared: Arc) { shared .metrics .connections_total .fetch_add(1, Ordering::Relaxed); shared .metrics .connections_active .fetch_add(1, Ordering::Relaxed); tracing::info!("WS connection opened"); let mut history: Vec = Vec::new(); history.push(ChatMessage::system(&shared.system_prompt)); // Per-connection rate limits. let mut audio_bucket = RateBucket::new(shared.rate_audio_secs_per_min); let mut turns_bucket = RateBucket::new(shared.rate_turns_per_min); // If the previous turn ended in a barge-in, the carried-over user audio // bytes seed the next turn so we don't drop them. let mut carry_over: Option> = None; 'session: loop { let turn_start = std::time::Instant::now(); // Charge the turn bucket up front; reject if over budget. if !turns_bucket.try_consume(1) { shared.metrics.errors_total.fetch_add(1, Ordering::Relaxed); send_text( &mut socket, "{\"event\":\"error\",\"msg\":\"rate limit: turns per minute exceeded\"}", ) .await; break 'session; } // Reset Kyutai STT state for each new turn (Whisper is stateless // so this is a no-op for Whisper). { let mut stt = shared.stt.lock().await; if let AsrEngine::Kyutai(s) = &mut *stt { if let Err(e) = s.reset() { send_text( &mut socket, &format!("{{\"event\":\"error\",\"msg\":\"stt reset: {e}\"}}"), ) .await; break 'session; } } } let mut user_audio_24k: Vec = Vec::new(); // Barge-in carryover from the previous turn. if let Some(co) = carry_over.take() { user_audio_24k.extend(co); } // Same reset (some legacy duplication; harmless and Kyutai-only). { let mut stt = shared.stt.lock().await; if let AsrEngine::Kyutai(s) = &mut *stt { if let Err(e) = s.reset() { send_text( &mut socket, &format!("{{\"event\":\"error\",\"msg\":\"stt reset: {e}\"}}"), ) .await; break 'session; } } } // VAD state (only used when --vad is enabled). let mut consecutive_eot: u32 = 0; // Tracks how much of user_audio_24k has been pushed to STT. let mut stt_streaming_offset: usize = 0; // Accumulate Word/EndWord events as they arrive. let mut pending_word: Option> = None; let mut transcript_words: Vec = Vec::new(); // If we have carry-over audio, feed it first (it's already in // user_audio_24k from the barge-in stash). Kyutai-only — Whisper // skips parallel ingest entirely. STT runs on the blocking pool // so the async worker can poll other tasks during the ~80ms // per-frame work (multi-tenant concurrency win). if !user_audio_24k.is_empty() && !shared.batch_asr { let carry_slice = user_audio_24k.clone(); stt_streaming_offset = carry_slice.len(); let shared_carry = shared.clone(); let events = tokio::task::spawn_blocking(move || { let mut guard = shared_carry.stt.blocking_lock(); match &mut *guard { AsrEngine::Kyutai(s) => match s.step_pcm(&carry_slice) { Ok(es) => es, Err(e) => { tracing::warn!("carry-over step_pcm: {e}"); Vec::new() } }, AsrEngine::Moonshine(_) => Vec::new(), #[cfg(feature = "asr")] AsrEngine::Whisper(_) => Vec::new(), } }) .await .unwrap_or_else(|e| { tracing::warn!("carry-over spawn_blocking: {e}"); Vec::new() }); collect_transcript_events( &events, &shared.stt, &mut pending_word, &mut transcript_words, ) .await; } else if !user_audio_24k.is_empty() { // Whisper mode: just track that carry-over is in the buffer. stt_streaming_offset = user_audio_24k.len(); } // Receive frames until EOT. loop { match socket.recv().await { Some(Ok(Message::Binary(bytes))) => { if bytes.len() % 2 != 0 { continue; } let n = bytes.len() / 2; user_audio_24k.reserve(n); for c in bytes.chunks_exact(2) { let s = i16::from_le_bytes([c[0], c[1]]); user_audio_24k.push(s as f32 / i16::MAX as f32); } // Kyutai path: feed new samples into STT incrementally // so the transcript builds as audio arrives. STT // runs on the blocking pool — frees the async worker // for other connections during the ~80 ms/frame // compute (multi-tenant concurrency win; no impact // on single-connection latency since STT is // hardware-bound). Whisper path skips this — audio // accumulates for the post-EOT batch transcribe. // // VAD gate (when --vad-gate): classify the new slice // before STT; if silence, skip step_pcm entirely. // Saves ~80-300 ms per silent chunk depending on // chunk size. let events = if !shared.batch_asr { let new_slice = user_audio_24k[stt_streaming_offset..].to_vec(); stt_streaming_offset = user_audio_24k.len(); let skip_stt = if let Some(gate) = shared.vad_gate.as_ref() { let mut g = gate.lock().await; !g.is_speech(&new_slice) } else { false }; if skip_stt { Vec::new() } else { let shared_step = shared.clone(); tokio::task::spawn_blocking(move || { let mut guard = shared_step.stt.blocking_lock(); match &mut *guard { AsrEngine::Kyutai(s) => match s.step_pcm(&new_slice) { Ok(es) => es, Err(e) => { tracing::warn!("step_pcm: {e}"); Vec::new() } }, AsrEngine::Moonshine(_) => Vec::new(), #[cfg(feature = "asr")] AsrEngine::Whisper(_) => Vec::new(), } }) .await .unwrap_or_else(|e| { tracing::warn!("step_pcm spawn_blocking: {e}"); Vec::new() }) } // close `if !skip_stt` else-arm } else { stt_streaming_offset = user_audio_24k.len(); Vec::new() }; collect_transcript_events( &events, &shared.stt, &mut pending_word, &mut transcript_words, ) .await; // VAD: watch Step events for end-of-turn probability. if let Some((threshold, k)) = shared.vad_eot { for ev in &events { if let Some(pr) = Stt::end_of_turn_probability(ev) { if pr > threshold { consecutive_eot += 1; if consecutive_eot >= k { tracing::info!( "VAD auto-EOT (pr={pr:.3} > {threshold} for {k} frames)" ); send_text(&mut socket, "{\"event\":\"vad_eot\"}").await; break; } } else { consecutive_eot = 0; } } } if consecutive_eot >= k { break; } } } Some(Ok(Message::Text(t))) if t.trim() == "EOT" => break, Some(Ok(Message::Text(_))) => continue, Some(Ok(Message::Close(_))) | None => break 'session, Some(Ok(_)) => continue, Some(Err(e)) => { tracing::warn!("WS recv: {e}"); break 'session; } } } // Mark the end of the audio-receive phase. Everything after this // is post-EOT processing (STT flush + LLM + TTS). let recv_done_t = std::time::Instant::now(); let recv_phase_ms = turn_start.elapsed().as_millis() as u64; shared .metrics .recv_phase_ms_sum .fetch_add(recv_phase_ms, Ordering::Relaxed); shared .metrics .recv_phase_ms_count .fetch_add(1, Ordering::Relaxed); // Drain the asr_delay buffer so any trailing Kyutai words // flush. Whisper has no asr_delay buffer to drain. let final_events = { let mut stt = shared.stt.lock().await; match &mut *stt { AsrEngine::Kyutai(s) => s.finish().unwrap_or_default(), AsrEngine::Moonshine(_) => Vec::new(), #[cfg(feature = "asr")] AsrEngine::Whisper(_) => Vec::new(), } }; collect_transcript_events( &final_events, &shared.stt, &mut pending_word, &mut transcript_words, ) .await; if user_audio_24k.is_empty() { send_text(&mut socket, "{\"event\":\"error\",\"msg\":\"empty audio\"}").await; continue 'session; } // Charge the audio-seconds bucket for this turn. let audio_secs = (user_audio_24k.len() as f32 / PCM_RATE as f32).ceil() as u32; if !audio_bucket.try_consume(audio_secs) { shared.metrics.errors_total.fetch_add(1, Ordering::Relaxed); send_text( &mut socket, "{\"event\":\"error\",\"msg\":\"rate limit: audio seconds per minute exceeded\"}", ) .await; break 'session; } // -- STT: transcribe the turn ------------------------------------- // STT runs at 24 kHz natively (Mimi sample rate). Pad with 2s // silence suffix so the asr_delay buffer flushes the last words. // (We use the same SAMPLE_RATE constant as the STT module.) debug_assert_eq!(STT_SR, PCM_RATE); user_audio_24k.extend(std::iter::repeat(0.0f32).take((PCM_RATE as f32 * 2.0) as usize)); // Build user_text. Kyutai has been streaming words incrementally // — just join them. Whisper is batch-only, so we fire one // transcribe call here over the full accumulated audio. (With // Whisper-tiny on Metal, 10s of audio takes ~200ms — much // faster than Kyutai's 5s parallel-STT-during-receive.) // Build user_text. Kyutai has been streaming words incrementally // — just join them. Batch backends (Whisper / Moonshine) fire // one transcribe call here over the full accumulated audio. let stt_t = std::time::Instant::now(); let user_text = if shared.batch_asr { let stt = shared.stt.lock().await; match &*stt { AsrEngine::Moonshine(m) => match m.transcribe_24k(&user_audio_24k) { Ok(t) => t.trim().to_string(), Err(e) => { tracing::warn!("moonshine transcribe: {e}"); String::new() } }, #[cfg(feature = "asr")] AsrEngine::Whisper(w) => match w.transcribe_24k(&user_audio_24k) { Ok(t) => t.trim().to_string(), Err(e) => { tracing::warn!("whisper transcribe: {e}"); String::new() } }, AsrEngine::Kyutai(_) => transcript_words.join(" ").trim().to_string(), } } else { transcript_words.join(" ").trim().to_string() }; let stt_ms = stt_t.elapsed().as_millis() as u64; shared .metrics .stt_latency_ms_sum .fetch_add(stt_ms, Ordering::Relaxed); shared .metrics .stt_latency_ms_count .fetch_add(1, Ordering::Relaxed); let _ = ASR_DELAY_FRAMES; let transcript_msg = serde_json::json!({"event":"transcript","text":&user_text}); send_text(&mut socket, &transcript_msg.to_string()).await; if user_text.is_empty() { shared.metrics.errors_total.fetch_add(1, Ordering::Relaxed); send_text( &mut socket, "{\"event\":\"error\",\"msg\":\"empty transcript\"}", ) .await; continue 'session; } // -- Reactive emotion detection (Phase 13.4) ---------------------- // Reuses the audio buffer STT just consumed. ProsodyDetector // operates at 16 kHz so we resample once on the fly; cheap for a // ~10 s buffer and only paid when --reactive-emotion is on. let reactive_tag: Option = if let Some(det) = shared.emotion_detector.as_ref() { // Trim the 2 s silence pad we appended for STT before // classifying — silence drags the voiced ratio down. let trim_to = user_audio_24k .len() .saturating_sub((PCM_RATE as f32 * 2.0) as usize); let speech = &user_audio_24k[..trim_to]; shared .metrics .reactive_emotion_calls_total .fetch_add(1, Ordering::Relaxed); match rtx_csm::audio_io::resample(speech, PCM_RATE, 16_000) { Ok(speech_16k) => { let label = det .classify(&speech_16k) .unwrap_or(rtx_csm::ser::EmotionLabel::Neutral); // Per-label counters. let bucket = match label { rtx_csm::ser::EmotionLabel::Neutral => { &shared.metrics.reactive_emotion_neutral_total } rtx_csm::ser::EmotionLabel::Calm => { &shared.metrics.reactive_emotion_calm_total } rtx_csm::ser::EmotionLabel::Sad => { &shared.metrics.reactive_emotion_sad_total } rtx_csm::ser::EmotionLabel::Angry => { &shared.metrics.reactive_emotion_angry_total } rtx_csm::ser::EmotionLabel::Excited => { &shared.metrics.reactive_emotion_excited_total } rtx_csm::ser::EmotionLabel::Disgusted => { &shared.metrics.reactive_emotion_disgusted_total } rtx_csm::ser::EmotionLabel::Fearful => { &shared.metrics.reactive_emotion_fearful_total } rtx_csm::ser::EmotionLabel::Happy => { &shared.metrics.reactive_emotion_happy_total } rtx_csm::ser::EmotionLabel::Surprised => { &shared.metrics.reactive_emotion_surprised_total } rtx_csm::ser::EmotionLabel::Unk => { &shared.metrics.reactive_emotion_unk_total } }; bucket.fetch_add(1, Ordering::Relaxed); if label == rtx_csm::ser::EmotionLabel::Neutral { None // fall back to static --emotion-hint } else { let tag = label.as_tag().to_string(); tracing::info!("reactive-emotion: detected {tag}"); Some(tag) } } Err(e) => { tracing::warn!("reactive-emotion resample failed: {e}"); None } } } else { None }; let resolved_hint = reactive_tag.clone().or_else(|| shared.emotion_hint.clone()); // -- LLM + TTS: stream response back as audio chunks -------------- history.push(ChatMessage::user(&user_text)); let opts = ConverseOptions { speaker: shared.speaker, flush: FlushPolicy::Punctuation, generate: GenerateOptions { max_audio_ms: 8_000, emotion_hint: resolved_hint, ..GenerateOptions::default() }, ..ConverseOptions::default() }; let gen_cfg = GenConfig { max_tokens: Some(160), temperature: 0.7, extra_body: shared.llm_extra_body.clone(), ..GenConfig::default() }; // Cancel signal: set by the pump when barge-in is detected; // checked inside the TTS callback to abort early. let cancel = Arc::new(AtomicBool::new(false)); let cancel_for_cb = cancel.clone(); // Bridging channel: TTS callback pushes encoded PCM, pump_and_watch // forwards to the WebSocket while concurrently watching for // user-audio frames (barge-in). let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::>(); // Build the LLM-facing history. When --emotion-aware-llm is on AND // we have a non-Neutral reactive tag (Phase 13.4), augment THIS // turn's user message with a tone signal so the LLM adapts its // text. The augmented message is per-turn-only — it never lands // in persistent `history`, so subsequent turns aren't biased by // stale signals. Phase 13.5. let history_clone = if shared.emotion_aware_llm && reactive_tag.is_some() { let tag = reactive_tag.as_deref().unwrap(); // Strip surrounding [] from the tag for the prose annotation. let label = tag.trim_matches(|c| c == '[' || c == ']'); let mut h = history.clone(); if let Some(last) = h.last_mut() { if matches!(last.role, rtx_csm::llm_client::Role::User) { last.content = format!( "{}\n\n[user audio tone: {label} — adjust your response in tone and content to match]", last.content ); shared .metrics .emotion_aware_llm_applied_total .fetch_add(1, Ordering::Relaxed); } } h } else { history.clone() }; let metrics_for_tts = shared.clone(); // Capture the moment conv_fut starts so we can derive // llm_to_first_audio (LLM stream + first-sentence TTS) and // conv_total (whole conversation) latencies. let conv_start_t = std::time::Instant::now(); let mut llm_to_first_audio_ms: Option = None; let stream_tts = shared.stream_tts; let stream_chunk_frames = shared.stream_chunk_frames; let cancel_for_chunk_cb = cancel.clone(); // Move-captured handles for the spawned conv task. Everything // here is Send + 'static (Arc clones, owned vectors, simple // values, channel senders) so we can lift conv_fut out of the // current task. Why spawn: tokio::join! polls cooperatively in // ONE task, so a sync TTS gen blocks the whole runtime worker // and prevents pump_fut from draining chunks. Spawn lets the // runtime reschedule pump_fut onto a different worker while // conv_fut is mid-synthesize. let shared_conv = shared.clone(); let conv_fut = async move { let mut g = shared_conv.generator.lock().await; let mut conv = Converse::new(&shared_conv.llm, &mut *g); let tx_cb = tx.clone(); let r = if stream_tts { // Streaming: PCM chunks delivered as they're produced // (low TTFA). Per-sentence callback only updates metrics // — the audio has already shipped via on_chunk. let tx_chunk = tx.clone(); let metrics_for_sentence = metrics_for_tts.clone(); conv.run_streaming( history_clone, gen_cfg, opts, stream_chunk_frames, move |chunk: &[f32]| { if cancel_for_chunk_cb.load(Ordering::Relaxed) { return Err(rtx_csm::CsmError::Config( "barge-in: TTS cancelled".into(), )); } let mut buf = Vec::with_capacity(chunk.len() * 2); for &s in chunk { let v = (s.clamp(-1.0, 1.0) * i16::MAX as f32) as i16; buf.extend_from_slice(&v.to_le_bytes()); } let _ = tx_chunk.send(buf); Ok(()) }, move |u| { metrics_for_sentence .metrics .tts_latency_ms_sum .fetch_add(u.tts_latency_ms as u64, Ordering::Relaxed); metrics_for_sentence .metrics .tts_latency_ms_count .fetch_add(1, Ordering::Relaxed); Ok(()) }, ) .await } else { conv.run(history_clone, gen_cfg, opts, move |u| { if cancel_for_cb.load(Ordering::Relaxed) { return Err(rtx_csm::CsmError::Config("barge-in: TTS cancelled".into())); } metrics_for_tts .metrics .tts_latency_ms_sum .fetch_add(u.tts_latency_ms as u64, Ordering::Relaxed); metrics_for_tts .metrics .tts_latency_ms_count .fetch_add(1, Ordering::Relaxed); let mut buf = Vec::with_capacity(u.audio.len() * 2); for &s in &u.audio { let v = (s.clamp(-1.0, 1.0) * i16::MAX as f32) as i16; buf.extend_from_slice(&v.to_le_bytes()); } let _ = tx_cb.send(buf); Ok(()) }) .await }; drop(tx); // close so pump exits naturally r }; // pump_and_watch: send TTS chunks to the WS while concurrently // watching for incoming binary frames (barge-in) or "EOT" text // (user wants to stop). Returns the carried-over user audio when // barge-in fires so the next turn can start with it pre-buffered. enum PumpResult { Done, BargeIn(Vec), Disconnected, } let mut first_audio_recorded = false; let pump_fut = async { loop { tokio::select! { biased; // poll audio_rx first so a fast TTS doesn't get starved chunk = rx.recv() => { match chunk { Some(buf) => { if !first_audio_recorded { first_audio_recorded = true; let e2e_ms = turn_start.elapsed().as_millis() as u64; shared .metrics .e2e_first_ms_sum .fetch_add(e2e_ms, Ordering::Relaxed); shared .metrics .e2e_first_ms_count .fetch_add(1, Ordering::Relaxed); // llm_to_first_audio: from conv_fut // start to the first PCM chunk arriving // here. Bundles LLM stream + first // sentence's TTS gen. let lta = conv_start_t.elapsed().as_millis() as u64; llm_to_first_audio_ms = Some(lta); shared .metrics .llm_to_first_audio_ms_sum .fetch_add(lta, Ordering::Relaxed); shared .metrics .llm_to_first_audio_ms_count .fetch_add(1, Ordering::Relaxed); } if socket.send(Message::Binary(buf)).await.is_err() { return PumpResult::Disconnected; } } None => return PumpResult::Done, } } incoming = socket.recv() => { match incoming { Some(Ok(Message::Binary(bytes))) if !bytes.is_empty() => { cancel.store(true, Ordering::Relaxed); // Decode the barge-in PCM so we can carry it // forward as the first audio of the next turn. let mut samples = Vec::with_capacity(bytes.len() / 2); for c in bytes.chunks_exact(2) { let s = i16::from_le_bytes([c[0], c[1]]); samples.push(s as f32 / i16::MAX as f32); } // Drain pending audio chunks so we don't keep // sending the cancelled TTS to the client. while rx.try_recv().is_ok() {} // Send a 50ms exponential-decay fade-out // chunk so the audio doesn't cut to a click. // (The client's last received sample is // arbitrary; we synthesize a smooth ramp // toward silence from low-amplitude noise.) let fade_len = (PCM_RATE as f32 * 0.05) as usize; let mut fade_buf = Vec::with_capacity(fade_len * 2); for i in 0..fade_len { let t = i as f32 / fade_len as f32; let env = (-4.0 * t).exp(); // 1.0 → ~0.018 // Soft pseudo-noise so the fade isn't dead silence. let n = ((i as u32).wrapping_mul(2654435761) as f32 / u32::MAX as f32 - 0.5) * 0.02; let s = n * env; let v = (s.clamp(-1.0, 1.0) * i16::MAX as f32) as i16; fade_buf.extend_from_slice(&v.to_le_bytes()); } let _ = socket .send(Message::Binary(fade_buf.into())) .await; let _ = socket .send(Message::Text("{\"event\":\"barge_in\"}".into())) .await; tracing::info!( "barge-in detected ({} samples carried over, fade-out sent)", samples.len() ); return PumpResult::BargeIn(samples); } Some(Ok(Message::Text(t))) if t.trim() == "EOT" => { cancel.store(true, Ordering::Relaxed); while rx.try_recv().is_ok() {} return PumpResult::Done; } Some(Ok(Message::Close(_))) | None => { return PumpResult::Disconnected; } Some(Ok(_)) => continue, Some(Err(_)) => return PumpResult::Disconnected, } } } } }; // Spawn conv_fut as a separate task so the runtime can schedule // pump_fut on another worker while conv_fut's sync TTS gen // blocks. Both finish before we move on; pump_fut blocks until // the channel closes (which happens when conv_fut drops `tx`). let conv_handle = tokio::spawn(conv_fut); let pump_res = pump_fut.await; let conv_res: rtx_csm::error::Result = match conv_handle.await { Ok(r) => r, Err(join_err) => Err(rtx_csm::CsmError::Config(format!( "conv task join: {join_err}" ))), }; let conv_total_ms = conv_start_t.elapsed().as_millis() as u64; shared .metrics .conv_total_ms_sum .fetch_add(conv_total_ms, Ordering::Relaxed); shared .metrics .conv_total_ms_count .fetch_add(1, Ordering::Relaxed); let mut barge_in_audio: Option> = None; match pump_res { PumpResult::Done => {} PumpResult::BargeIn(samples) => { barge_in_audio = Some(samples); } PumpResult::Disconnected => break 'session, } // If barge-in cancelled TTS the conv result will be an Err. That's // expected — don't surface it as an error to the client. let assistant_text = if cancel.load(Ordering::Relaxed) { // Cancelled — assistant turn is incomplete. Don't store in history; // the user's next turn will replace what they were responding to. String::new() } else { match conv_res { Ok(t) => t, Err(e) => { shared.metrics.errors_total.fetch_add(1, Ordering::Relaxed); send_text( &mut socket, &format!("{{\"event\":\"error\",\"msg\":\"converse: {e}\"}}"), ) .await; continue 'session; } } }; if !assistant_text.is_empty() { history.push(ChatMessage::assistant(&assistant_text)); shared.metrics.turns_total.fetch_add(1, Ordering::Relaxed); let done_msg = serde_json::json!({"event":"done","assistant":&assistant_text}); send_text(&mut socket, &done_msg.to_string()).await; } let total_turn_ms = turn_start.elapsed().as_millis() as u64; shared .metrics .total_turn_ms_sum .fetch_add(total_turn_ms, Ordering::Relaxed); shared .metrics .total_turn_ms_count .fetch_add(1, Ordering::Relaxed); // Per-turn timing summary. recv + stt_post + conv_total ≈ total // (small slack for mutex acquire, channel close, log formatting). // llm_to_first_audio is a sub-phase of conv_total: it's the time // from conv_fut start to the first PCM byte hitting the pump. let stt_post_ms = conv_start_t .saturating_duration_since(recv_done_t) .as_millis() as u64; tracing::info!( "turn timing: recv={}ms stt_post={}ms llm_to_first_audio={}ms \ conv_total={}ms total={}ms", recv_phase_ms, stt_post_ms, llm_to_first_audio_ms.unwrap_or(0), conv_total_ms, total_turn_ms, ); // Stash any barge-in audio so the next turn picks up where the user started. carry_over = barge_in_audio; } shared .metrics .connections_active .fetch_sub(1, Ordering::Relaxed); tracing::info!("WS connection closed"); } async fn send_text(socket: &mut WebSocket, payload: &str) -> bool { socket .send(Message::Text(payload.to_string())) .await .is_ok() } /// Walks Word/EndWord events: pairs them, detokenizes each word via the /// STT's sentencepiece tokenizer, and appends to `transcript_words`. Step /// events (VAD prs) are ignored here — they're handled by the VAD loop. async fn collect_transcript_events( events: &[AsrEvent], stt_lock: &Mutex, pending: &mut Option>, transcript_words: &mut Vec, ) { if events.is_empty() { return; } let stt = stt_lock.lock().await; let kyutai = match &*stt { AsrEngine::Kyutai(s) => s, // Batch backends (Moonshine / Whisper) don't go through this — // events come from Kyutai's incremental Word/EndWord stream only. AsrEngine::Moonshine(_) => return, #[cfg(feature = "asr")] AsrEngine::Whisper(_) => return, }; for ev in events { match ev { AsrEvent::Word { tokens, .. } => *pending = Some(tokens.clone()), AsrEvent::EndWord { .. } => { if let Some(tokens) = pending.take() { if let Some(w) = kyutai.decode_word_text(&tokens) { let w = w.trim().to_string(); if !w.is_empty() { transcript_words.push(w); } } } } AsrEvent::Step { .. } => {} } } }