rtx-csm: Phase 6c.2 — Rust Unmute MVP, voice conversation round-trip

Full-stack pure-Rust voice conversation server + CLI client:

  Client -> Server  binary frames: 16-bit LE PCM @ 24 kHz mono
  Client -> Server  text "EOT": signal end-of-turn
  Server: STT (Kyutai 1B en/fr) -> transcript
          LLM (OpenAI-compatible OR mock echo) -> token stream
          Converse: sentence buffer -> CSM TTS -> 16-bit LE PCM
  Server -> Client  text {"event":"transcript","text":"..."}
  Server -> Client  binary frames: assistant audio
  Server -> Client  text {"event":"done","assistant":"..."}

Per-connection chat history; multiple turns supported per socket.
--mock-llm mode for testing without API keys (echoes user transcript).

examples/converse_server.rs: axum WebSocket server.
examples/converse_client.rs: CLI; streams WAV in as user turn, saves
  response audio out.

Verified end-to-end on Metal:
  Input: 10.43s LibriSpeech FLAC ("He hoped there would be stew...")
  STT transcript: matched (full sentence captured by 23/25 words)
  Mock LLM: "I heard you say: <transcript>."
  CSM TTS: response audio streamed back via WebSocket
  Round-trip wall-clock: 6.86s (TTFA on first audio chunk: 6.86s; the
  pipeline is sequential per turn — Phase 6c.3 would pipeline LLM
  tokens with TTS to get TTFA much lower).

This is the Rust Unmute MVP: PCM in, voice out, no Python in the
runtime path. Strategic Phase 6 deliverable.

Phase 6 status:
  6a STT: working
  6b LLM client: working
  6c.1 text->LLM->TTS: working
  6c.2 WebSocket duplex MVP: working (this commit)
  6c.3 streaming pipeline + auto-EOT + barge-in: deferred
  6d productionization: deferred

Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>

crates/models/rtx-csm/Cargo.toml

@@ -83,8 +83,10 @@ clap = { version = "4.5", features = ["derive"] }
 tempfile = "3.0"
 approx = "0.5"
 tracing-subscriber = "0.3"
-# For the TTS HTTP server example.
+# For the TTS HTTP server + converse_server WebSocket examples.
-axum = { version = "0.7", features = ["multipart"] }
+axum = { version = "0.7", features = ["multipart", "ws"] }
+# WebSocket client for examples/converse_client.
+tokio-tungstenite = { version = "0.24", default-features = false, features = ["connect", "rustls-tls-webpki-roots"] }
 # tokio with extra features (signal handler) needed by tts_server.
 tokio = { version = "1", features = ["macros", "rt-multi-thread", "signal", "sync"] }
 tower = "0.5"
@@ -199,3 +201,11 @@ path = "examples/llm_chat.rs"
 [[example]]
 name = "converse"
 path = "examples/converse.rs"
+
+[[example]]
+name = "converse_server"
+path = "examples/converse_server.rs"
+
+[[example]]
+name = "converse_client"
+path = "examples/converse_client.rs"

crates/models/rtx-csm/examples/converse_client.rs

@@ -0,0 +1,138 @@
+//! WebSocket client for `converse_server`. Streams a WAV file as the
+//! user turn, prints the transcript, saves the assistant response audio.
+//!
+//! Usage:
+//! ```bash
+//! # in another terminal: examples/converse_server with model loaded
+//! cargo run -p rtx-csm --release --example converse_client -- \
+//!   --url ws://127.0.0.1:18090/v1/converse \
+//!   --in /tmp/asr_test.flac \
+//!   --out /tmp/converse_response.wav
+//! ```
+
+use anyhow::{Context, Result};
+use clap::Parser;
+use futures_util::{SinkExt, StreamExt};
+use rtx_csm::audio_io;
+use std::path::PathBuf;
+use tokio_tungstenite::tungstenite::protocol::Message;
+
+const PCM_RATE: u32 = 24_000;
+const FRAME_SAMPLES: usize = 4800; // 200 ms per frame
+
+#[derive(Debug, Parser)]
+struct Cli {
+    #[arg(long, default_value = "ws://127.0.0.1:18090/v1/converse")]
+    url: String,
+    /// WAV/FLAC user-turn audio (any rate; resampled to 24 kHz mono).
+    #[arg(long = "in")]
+    input: PathBuf,
+    /// Output WAV (assistant response).
+    #[arg(long)]
+    out: PathBuf,
+}
+
+#[tokio::main]
+async fn main() -> Result<()> {
+    tracing_subscriber::fmt().init();
+    let cli = Cli::parse();
+
+    let samples = audio_io::load_mono_at_rate(&cli.input, PCM_RATE)
+        .context("load user audio")?;
+    println!(
+        "loaded {}: {} samples @ {} Hz ({:.2}s)",
+        cli.input.display(),
+        samples.len(),
+        PCM_RATE,
+        samples.len() as f32 / PCM_RATE as f32
+    );
+
+    let (mut ws, _resp) = tokio_tungstenite::connect_async(&cli.url)
+        .await
+        .with_context(|| format!("connect {}", cli.url))?;
+    println!("connected to {}", cli.url);
+
+    // Stream PCM in 200 ms frames. Server can begin transcribing as soon
+    // as it has audio.
+    for chunk in samples.chunks(FRAME_SAMPLES) {
+        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());
+        }
+        ws.send(Message::Binary(buf.into())).await?;
+    }
+    ws.send(Message::Text("EOT".into())).await?;
+    println!("sent {} samples + EOT", samples.len());
+
+    // Read response until "done" event.
+    let mut response_pcm: Vec<f32> = Vec::new();
+    let mut transcript = String::new();
+    let mut assistant = String::new();
+    let t = std::time::Instant::now();
+    let mut first_audio_ms: Option<u128> = None;
+
+    while let Some(msg) = ws.next().await {
+        let msg = msg?;
+        match msg {
+            Message::Binary(bytes) => {
+                if first_audio_ms.is_none() {
+                    first_audio_ms = Some(t.elapsed().as_millis());
+                }
+                for c in bytes.chunks_exact(2) {
+                    let s = i16::from_le_bytes([c[0], c[1]]);
+                    response_pcm.push(s as f32 / i16::MAX as f32);
+                }
+            }
+            Message::Text(t) => {
+                let parsed: serde_json::Value =
+                    serde_json::from_str(&t).context("parse server event")?;
+                let event = parsed.get("event").and_then(|v| v.as_str()).unwrap_or("");
+                match event {
+                    "transcript" => {
+                        transcript = parsed
+                            .get("text")
+                            .and_then(|v| v.as_str())
+                            .unwrap_or("")
+                            .to_string();
+                        println!("transcript: {transcript:?}");
+                    }
+                    "done" => {
+                        assistant = parsed
+                            .get("assistant")
+                            .and_then(|v| v.as_str())
+                            .unwrap_or("")
+                            .to_string();
+                        println!("assistant: {assistant:?}");
+                        break;
+                    }
+                    "error" => {
+                        let msg = parsed
+                            .get("msg")
+                            .and_then(|v| v.as_str())
+                            .unwrap_or("(unknown)");
+                        anyhow::bail!("server error: {msg}");
+                    }
+                    other => println!("(server event {other}: {t})"),
+                }
+            }
+            Message::Close(_) => break,
+            _ => {}
+        }
+    }
+
+    audio_io::write_wav_24k_mono(&cli.out, &response_pcm)?;
+    println!(
+        "wrote {} samples = {:.2}s @ {} Hz to {}",
+        response_pcm.len(),
+        response_pcm.len() as f32 / PCM_RATE as f32,
+        PCM_RATE,
+        cli.out.display()
+    );
+    println!(
+        "ttfa={:?}ms, total wall = {:.2}s",
+        first_audio_ms,
+        t.elapsed().as_secs_f32()
+    );
+    Ok(())
+}

crates/models/rtx-csm/examples/converse_server.rs

@@ -0,0 +1,377 @@
+//! Rust Unmute MVP — WebSocket conversational server.
+//!
+//! 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`.
+
+use anyhow::Result;
+use async_trait::async_trait;
+use axum::{
+    extract::{
+        ws::{Message, WebSocket},
+        State, WebSocketUpgrade,
+    },
+    response::IntoResponse,
+    routing::get,
+    Router,
+};
+use clap::Parser;
+use futures_util::stream;
+use rtx_csm::{
+    converse::{Converse, ConverseOptions, FlushPolicy},
+    error::Result as CsmResult,
+    llm_client::{ChatMessage, GenConfig, LlmClient, OpenAiCompatibleClient, TokenStream},
+    stt::{AsrEvent, Stt, SAMPLE_RATE as STT_SR, ASR_DELAY_FRAMES},
+    GenerateOptions, Generator,
+};
+use std::net::SocketAddr;
+use std::sync::Arc;
+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<String>,
+    /// 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,
+}
+
+/// Echo-style mock LLM. Reads the most recent user message from the
+/// chat history and returns a fixed acknowledgment + repeats the heard
+/// transcript. Useful for end-to-end testing without an API key.
+struct MockLlm;
+
+#[async_trait]
+impl LlmClient for MockLlm {
+    async fn generate_stream(
+        &self,
+        messages: Vec<ChatMessage>,
+        _config: GenConfig,
+    ) -> CsmResult<TokenStream> {
+        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 say: {}.", user_text);
+        // Emit it as a small number of chunks to exercise streaming.
+        let chunks: Vec<CsmResult<String>> = response
+            .split_inclusive(' ')
+            .map(|s| Ok(s.to_string()))
+            .collect();
+        Ok(Box::pin(stream::iter(chunks)))
+    }
+}
+
+enum AnyLlm {
+    Real(OpenAiCompatibleClient),
+    Mock(MockLlm),
+}
+
+#[async_trait]
+impl LlmClient for AnyLlm {
+    async fn generate_stream(
+        &self,
+        messages: Vec<ChatMessage>,
+        config: GenConfig,
+    ) -> CsmResult<TokenStream> {
+        match self {
+            AnyLlm::Real(c) => c.generate_stream(messages, config).await,
+            AnyLlm::Mock(c) => c.generate_stream(messages, config).await,
+        }
+    }
+}
+
+struct Shared {
+    generator: Mutex<Generator>,
+    stt: Mutex<Stt>,
+    llm: AnyLlm,
+    system_prompt: String,
+    speaker: u32,
+}
+
+#[tokio::main]
+async fn main() -> Result<()> {
+    tracing_subscriber::fmt().init();
+    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:?}");
+
+    tracing::info!("loading CSM-1B...");
+    let generator = Generator::load_csm_1b(&device)?;
+    tracing::info!("loading Kyutai STT 1B en/fr (~3 GB)...");
+    let stt = Stt::load_default(&device)?;
+    tracing::info!("models loaded");
+
+    let shared = Arc::new(Shared {
+        generator: Mutex::new(generator),
+        stt: Mutex::new(stt),
+        llm,
+        system_prompt: cli.system,
+        speaker: cli.speaker,
+    });
+
+    let app = Router::new()
+        .route("/health", get(|| async { "ok" }))
+        .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);
+    axum::serve(listener, app).await?;
+    Ok(())
+}
+
+async fn ws_handler(
+    ws: WebSocketUpgrade,
+    State(shared): State<Arc<Shared>>,
+) -> impl IntoResponse {
+    ws.on_upgrade(move |socket| handle_connection(socket, shared))
+}
+
+async fn handle_connection(mut socket: WebSocket, shared: Arc<Shared>) {
+    tracing::info!("WS connection opened");
+    let mut history: Vec<ChatMessage> = Vec::new();
+    history.push(ChatMessage::system(&shared.system_prompt));
+
+    'session: loop {
+        // Reset STT state for each new turn so silence buffer + delay
+        // counters start fresh.
+        if let Err(e) = shared.stt.lock().await.reset() {
+            send_text(&mut socket, &format!("{{\"event\":\"error\",\"msg\":\"stt reset: {e}\"}}"))
+                .await;
+            break 'session;
+        }
+
+        let mut user_audio_24k: Vec<f32> = Vec::new();
+        // Receive frames until EOT.
+        loop {
+            match socket.recv().await {
+                Some(Ok(Message::Binary(bytes))) => {
+                    // Decode 16-bit LE PCM @ 24 kHz mono.
+                    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);
+                    }
+                }
+                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;
+                }
+            }
+        }
+
+        if user_audio_24k.is_empty() {
+            send_text(&mut socket, "{\"event\":\"error\",\"msg\":\"empty audio\"}").await;
+            continue '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));
+
+        let user_text = {
+            let mut stt = shared.stt.lock().await;
+            // step_pcm in big chunks; STT is sync but small enough that
+            // we can just hold the lock through the full transcription.
+            let mut events = Vec::new();
+            for chunk in user_audio_24k.chunks(PCM_RATE as usize * 2) {
+                match stt.step_pcm(chunk) {
+                    Ok(es) => events.extend(es),
+                    Err(e) => {
+                        let _ = send_text(
+                            &mut socket,
+                            &format!("{{\"event\":\"error\",\"msg\":\"stt: {e}\"}}"),
+                        )
+                        .await;
+                        continue 'session;
+                    }
+                }
+            }
+            if let Ok(es) = stt.finish() {
+                events.extend(es);
+            }
+
+            // Pair Word/EndWord, detok.
+            let mut text = String::new();
+            let mut pending: Option<Vec<u32>> = None;
+            for ev in events {
+                match ev {
+                    AsrEvent::Word { tokens, .. } => pending = Some(tokens),
+                    AsrEvent::EndWord { .. } => {
+                        if let Some(tokens) = pending.take() {
+                            if let Some(w) = stt.decode_word_text(&tokens) {
+                                if !text.is_empty() && !w.is_empty() {
+                                    text.push(' ');
+                                }
+                                text.push_str(&w);
+                            }
+                        }
+                    }
+                    AsrEvent::Step { .. } => {}
+                }
+            }
+            text.trim().to_string()
+        };
+        // Account for the asr_delay shift (just informational).
+        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() {
+            send_text(&mut socket, "{\"event\":\"error\",\"msg\":\"empty transcript\"}").await;
+            continue 'session;
+        }
+
+        // -- 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,
+                ..GenerateOptions::default()
+            },
+            ..ConverseOptions::default()
+        };
+        let gen_cfg = GenConfig {
+            max_tokens: Some(160),
+            temperature: 0.7,
+            ..GenConfig::default()
+        };
+        let assistant_text = {
+            let mut g = shared.generator.lock().await;
+            let mut conv = Converse::new(&shared.llm, &mut *g);
+            // axum's WebSocket isn't Sync across .await on Mutex, so we
+            // pull events into a Vec, then send between awaits.
+            let mut to_send: Vec<Vec<u8>> = Vec::new();
+            let mut text_full = String::new();
+            let result = conv
+                .run(history.clone(), gen_cfg, opts, |u| {
+                    text_full.push_str(&u.text);
+                    text_full.push(' ');
+                    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());
+                    }
+                    to_send.push(buf);
+                    Ok(())
+                })
+                .await;
+            // Release the generator mutex before sending.
+            drop(g);
+            for buf in to_send {
+                if socket.send(Message::Binary(buf)).await.is_err() {
+                    break 'session;
+                }
+            }
+            match result {
+                Ok(t) => t,
+                Err(e) => {
+                    send_text(
+                        &mut socket,
+                        &format!("{{\"event\":\"error\",\"msg\":\"converse: {e}\"}}"),
+                    )
+                    .await;
+                    continue 'session;
+                }
+            }
+        };
+
+        history.push(ChatMessage::assistant(&assistant_text));
+        let done_msg = serde_json::json!({"event":"done","assistant":&assistant_text});
+        send_text(&mut socket, &done_msg.to_string()).await;
+        // Loop back to accept the next turn.
+    }
+
+    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()
+}