rtx-csm: Phase 6c.3a + 6d.{auth,shutdown,metrics} — productionize converse

Production hardening pass on the Rust Unmute MVP:

6c.3a: streaming audio output. Replaced the "collect-then-send" loop
  with a tokio mpsc channel + tokio::join! between conv.run and the
  WS sender. Audio chunks are forwarded to the client AS each sentence
  completes TTS, instead of waiting for the full assistant response.
  Drop the channel sender at end of conv.run to signal the pump exit.

6d.auth: Bearer token auth. --auth-token flag (or RTX_AUTH_TOKEN env)
  on the server requires an Authorization: Bearer <token> header on the
  WebSocket upgrade. Rejected upgrades return 401. Server logs a warn
  if no token is configured (open dev mode). converse_client gains a
  matching --auth-token flag.

6d.shutdown: Graceful SIGINT/SIGTERM. tokio::signal handlers wired
  into axum::serve.with_graceful_shutdown(). Verified: SIGINT log line
  "received SIGINT, shutting down gracefully" + clean exit 0.

6d.metrics: /metrics Prometheus-style endpoint. Counters
  (turns_total, errors_total, connections_total) + gauges
  (connections_active, stt/tts/e2e_first_audio latency averages).
  Verified end-to-end: rtx_csm_turns_total 1 / errors_total 3 (from
  earlier 401 attempts) / e2e_first_audio_ms_avg 21295 / etc.

Verified all four together: 401 on bad/missing auth, 200 + WS upgrade
on correct auth, full round-trip metrics, clean SIGINT exit.

Phase 6 status:
  6a STT: working
  6b LLM client: working
  6c.1 text->LLM->TTS: working
  6c.2 WebSocket duplex MVP: working
  6c.3a streaming TTS chunks: working (this commit)
  6c.3b semantic VAD / barge-in: deferred
  6d.{auth,shutdown,metrics}: shipped (this commit)
  6d.rate-limit: deferred

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

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

@@ -15,6 +15,7 @@ use clap::Parser;
 use futures_util::{SinkExt, StreamExt};
 use rtx_csm::audio_io;
 use std::path::PathBuf;
+use tokio_tungstenite::tungstenite::client::IntoClientRequest;
 use tokio_tungstenite::tungstenite::protocol::Message;
 
 const PCM_RATE: u32 = 24_000;
@@ -30,6 +31,10 @@ struct Cli {
     /// Output WAV (assistant response).
     #[arg(long)]
     out: PathBuf,
+    /// Bearer token sent in the Authorization header. Reads
+    /// RTX_AUTH_TOKEN env var if not provided.
+    #[arg(long)]
+    auth_token: Option<String>,
 }
 
 #[tokio::main]
@@ -47,7 +52,19 @@ async fn main() -> Result<()> {
         samples.len() as f32 / PCM_RATE as f32
     );
 
-    let (mut ws, _resp) = tokio_tungstenite::connect_async(&cli.url)
+    let mut request = (&cli.url)
+        .into_client_request()
+        .with_context(|| format!("parse url {}", cli.url))?;
+    if let Some(token) = cli
+        .auth_token
+        .or_else(|| std::env::var("RTX_AUTH_TOKEN").ok())
+    {
+        request.headers_mut().insert(
+            "Authorization",
+            format!("Bearer {token}").parse().context("auth header")?,
+        );
+    }
+    let (mut ws, _resp) = tokio_tungstenite::connect_async(request)
         .await
         .with_context(|| format!("connect {}", cli.url))?;
     println!("connected to {}", cli.url);
@@ -67,8 +84,6 @@ async fn main() -> Result<()> {
 
     // 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;
 
@@ -90,19 +105,17 @@ async fn main() -> Result<()> {
                 let event = parsed.get("event").and_then(|v| v.as_str()).unwrap_or("");
                 match event {
                     "transcript" => {
-                        transcript = parsed
+                        let transcript = parsed
                             .get("text")
                             .and_then(|v| v.as_str())
-                            .unwrap_or("")
+                            .unwrap_or("");
-                            .to_string();
                         println!("transcript: {transcript:?}");
                     }
                     "done" => {
-                        assistant = parsed
+                        let assistant = parsed
                             .get("assistant")
                             .and_then(|v| v.as_str())
-                            .unwrap_or("")
+                            .unwrap_or("");
-                            .to_string();
                         println!("assistant: {assistant:?}");
                         break;
                     }

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

@@ -39,6 +39,7 @@ use axum::{
         ws::{Message, WebSocket},
         State, WebSocketUpgrade,
     },
+    http::{HeaderMap, StatusCode},
     response::IntoResponse,
     routing::get,
     Router,
@@ -53,6 +54,7 @@ use rtx_csm::{
     GenerateOptions, Generator,
 };
 use std::net::SocketAddr;
+use std::sync::atomic::{AtomicU64, Ordering};
 use std::sync::Arc;
 use tokio::sync::Mutex;
 
@@ -78,6 +80,12 @@ struct Cli {
     speaker: u32,
     #[arg(long)]
     cpu: 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<String>,
 }
 
 /// Echo-style mock LLM. Reads the most recent user message from the
@@ -132,12 +140,35 @@ impl LlmClient for AnyLlm {
     }
 }
 
+#[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,
+}
+
 struct Shared {
     generator: Mutex<Generator>,
     stt: Mutex<Stt>,
     llm: AnyLlm,
     system_prompt: String,
     speaker: u32,
+    auth_token: Option<String>,
+    metrics: Metrics,
 }
 
 #[tokio::main]
@@ -174,38 +205,128 @@ async fn main() -> Result<()> {
     let stt = Stt::load_default(&device)?;
     tracing::info!("models loaded");
 
+    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 shared = Arc::new(Shared {
         generator: Mutex::new(generator),
         stt: Mutex::new(stt),
         llm,
         system_prompt: cli.system,
         speaker: cli.speaker,
+        auth_token,
+        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);
-    axum::serve(listener, app).await?;
+    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<Arc<Shared>>) -> 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 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",
+        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,
+    );
+    (
+        StatusCode::OK,
+        [(axum::http::header::CONTENT_TYPE, "text/plain; version=0.0.4")],
+        body,
+    )
+}
+
 async fn ws_handler(
     ws: WebSocketUpgrade,
     State(shared): State<Arc<Shared>>,
+    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>) {
+    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<ChatMessage> = Vec::new();
     history.push(ChatMessage::system(&shared.system_prompt));
 
     'session: loop {
+        let turn_start = std::time::Instant::now();
         // Reset STT state for each new turn so silence buffer + delay
         // counters start fresh.
         if let Err(e) = shared.stt.lock().await.reset() {
@@ -254,6 +375,7 @@ async fn handle_connection(mut socket: WebSocket, shared: Arc<Shared>) {
         user_audio_24k
             .extend(std::iter::repeat(0.0f32).take((PCM_RATE as f32 * 2.0) as usize));
 
+        let stt_t = std::time::Instant::now();
         let user_text = {
             let mut stt = shared.stt.lock().await;
             // step_pcm in big chunks; STT is sync but small enough that
@@ -297,12 +419,22 @@ async fn handle_connection(mut socket: WebSocket, shared: Arc<Shared>) {
             }
             text.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);
         // 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() {
+            shared.metrics.errors_total.fetch_add(1, Ordering::Relaxed);
             send_text(&mut socket, "{\"event\":\"error\",\"msg\":\"empty transcript\"}").await;
             continue 'session;
         }
@@ -323,36 +455,69 @@ async fn handle_connection(mut socket: WebSocket, shared: Arc<Shared>) {
             temperature: 0.7,
             ..GenConfig::default()
         };
-        let assistant_text = {
+        // Bridging channel: the sync TTS callback pushes encoded PCM
+        // chunks; the async pump forwards them to the WebSocket as they
+        // arrive. tokio::join! lets conv.run (LLM streaming + TTS) run
+        // concurrently with the WebSocket send.
+        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::<Vec<u8>>();
+        let history_clone = history.clone();
+        let metrics_for_tts = shared.clone();
+        let conv_fut = async {
             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
+            let tx_cb = tx.clone();
-            // pull events into a Vec, then send between awaits.
+            let r = conv
-            let mut to_send: Vec<Vec<u8>> = Vec::new();
+                .run(history_clone, gen_cfg, opts, move |u| {
-            let mut text_full = String::new();
+                    metrics_for_tts
-            let result = conv
+                        .metrics
-                .run(history.clone(), gen_cfg, opts, |u| {
+                        .tts_latency_ms_sum
-                    text_full.push_str(&u.text);
+                        .fetch_add(u.tts_latency_ms as u64, Ordering::Relaxed);
-                    text_full.push(' ');
+                    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());
                     }
-                    to_send.push(buf);
+                    let _ = tx_cb.send(buf);
                     Ok(())
                 })
                 .await;
-            // Release the generator mutex before sending.
+            drop(tx); // close the channel so the pump loop exits
-            drop(g);
+            r
-            for buf in to_send {
+        };
+        let mut first_audio_recorded = false;
+        let pump_fut = async {
+            while let Some(buf) = rx.recv().await {
+                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);
+                }
                 if socket.send(Message::Binary(buf)).await.is_err() {
+                    return Err(());
+                }
+            }
+            Ok(())
+        };
+        let (conv_res, pump_res) = tokio::join!(conv_fut, pump_fut);
+        if pump_res.is_err() {
+            // Client disconnected mid-stream.
             break 'session;
         }
-            }
+        let assistant_text = match conv_res {
-            match result {
             Ok(t) => t,
             Err(e) => {
+                shared.metrics.errors_total.fetch_add(1, Ordering::Relaxed);
                 send_text(
                     &mut socket,
                     &format!("{{\"event\":\"error\",\"msg\":\"converse: {e}\"}}"),
@@ -360,15 +525,16 @@ async fn handle_connection(mut socket: WebSocket, shared: Arc<Shared>) {
                 .await;
                 continue 'session;
             }
-            }
         };
 
         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;
         // Loop back to accept the next turn.
     }
 
+    shared.metrics.connections_active.fetch_sub(1, Ordering::Relaxed);
     tracing::info!("WS connection closed");
 }