diff options
Diffstat (limited to 'makima/src/server')
| -rw-r--r-- | makima/src/server/handlers/mod.rs | 1 | ||||
| -rw-r--r-- | makima/src/server/handlers/speak.rs | 274 | ||||
| -rw-r--r-- | makima/src/server/messages.rs | 161 | ||||
| -rw-r--r-- | makima/src/server/mod.rs | 3 | ||||
| -rw-r--r-- | makima/src/server/state.rs | 22 |
5 files changed, 460 insertions, 1 deletions
diff --git a/makima/src/server/handlers/mod.rs b/makima/src/server/handlers/mod.rs index b496922..8207399 100644 --- a/makima/src/server/handlers/mod.rs +++ b/makima/src/server/handlers/mod.rs @@ -17,6 +17,7 @@ pub mod mesh_red_team; pub mod mesh_supervisor; pub mod mesh_ws; pub mod repository_history; +pub mod speak; pub mod templates; pub mod transcript_analysis; pub mod users; diff --git a/makima/src/server/handlers/speak.rs b/makima/src/server/handlers/speak.rs new file mode 100644 index 0000000..75e7780 --- /dev/null +++ b/makima/src/server/handlers/speak.rs @@ -0,0 +1,274 @@ +//! WebSocket handler for TTS streaming (direct in-process inference). +//! +//! This module implements the `/api/v1/speak` endpoint which performs +//! text-to-speech synthesis directly using the candle-based TTS engine. +//! No external Python service or proxy — the model runs in-process. +//! +//! ## Architecture +//! +//! The speak handler will: +//! 1. Accept a WebSocket connection from the client +//! 2. Lazily load the TTS model (candle) on first request +//! 3. Parse JSON control messages (start, speak, stop, cancel) +//! 4. Run inference directly and stream audio chunks back +//! +//! See `makima/src/tts/` for the TTS engine implementation. +//! See `docs/specs/qwen3-tts-spec.md` for the full protocol specification. + +use axum::{ + extract::{ws::Message, ws::WebSocket, State, WebSocketUpgrade}, + response::Response, +}; +use futures::{SinkExt, StreamExt}; +use serde::Deserialize; +use uuid::Uuid; + +use crate::server::state::SharedState; + +/// Client-to-server control messages. +#[derive(Debug, Deserialize)] +#[serde(tag = "type", rename_all = "snake_case")] +enum ClientMessage { + /// Request speech synthesis for the given text. + Speak { + text: String, + /// Optional voice ID (e.g., "makima"). Not yet used — reserved for future voice selection. + #[serde(default)] + #[allow(dead_code)] + voice: Option<String>, + }, + /// Cancel any in-progress synthesis. + Cancel, + /// Graceful close. + Stop, +} + +/// WebSocket upgrade handler for TTS streaming. +/// +/// This endpoint accepts WebSocket connections for text-to-speech synthesis. +/// The TTS model runs directly in-process using candle — no external service. +#[utoipa::path( + get, + path = "/api/v1/speak", + responses( + (status = 101, description = "WebSocket connection established"), + (status = 503, description = "TTS engine not available"), + ), + tag = "Speak" +)] +pub async fn websocket_handler( + ws: WebSocketUpgrade, + State(state): State<SharedState>, +) -> Response { + ws.on_upgrade(|socket| handle_speak_socket(socket, state)) +} + +/// Handle TTS WebSocket session with direct in-process inference. +/// +/// Protocol: +/// - Client sends JSON `{ "type": "speak", "text": "..." }` messages +/// - Server responds with binary audio chunks (16-bit PCM @ 24kHz) +/// - Server sends JSON `{ "type": "audio_end" }` when synthesis is complete +/// - Server sends JSON `{ "type": "error", ... }` on failures +async fn handle_speak_socket(socket: WebSocket, state: SharedState) { + let session_id = Uuid::new_v4().to_string(); + tracing::info!(session_id = %session_id, "New TTS WebSocket connection"); + + let (mut sender, mut receiver) = socket.split(); + + // Process incoming messages + while let Some(msg) = receiver.next().await { + let msg = match msg { + Ok(m) => m, + Err(e) => { + tracing::warn!(session_id = %session_id, error = %e, "WebSocket receive error"); + break; + } + }; + + match msg { + Message::Text(text) => { + let client_msg: ClientMessage = match serde_json::from_str(&text) { + Ok(m) => m, + Err(e) => { + let _ = send_error( + &mut sender, + "INVALID_MESSAGE", + &format!("Failed to parse message: {e}"), + ) + .await; + continue; + } + }; + + match client_msg { + ClientMessage::Speak { text, .. } => { + tracing::info!( + session_id = %session_id, + text_len = text.len(), + "TTS speak request" + ); + + // Get or lazily load the TTS engine + let engine = match state.get_tts_engine().await { + Ok(e) => e, + Err(e) => { + tracing::error!( + session_id = %session_id, + error = %e, + "Failed to load TTS engine" + ); + let _ = send_error( + &mut sender, + "TTS_LOAD_FAILED", + &format!("Failed to load TTS engine: {e}"), + ) + .await; + continue; + } + }; + + if !engine.is_ready() { + let _ = send_error( + &mut sender, + "TTS_NOT_READY", + "TTS engine is not ready yet", + ) + .await; + continue; + } + + // Run TTS inference (no voice reference for now — uses default) + match engine.generate(&text, None, None).await { + Ok(chunks) => { + for chunk in &chunks { + // Send binary PCM audio data + let pcm_bytes = chunk.to_pcm16_bytes(); + if sender + .send(Message::Binary(pcm_bytes.into())) + .await + .is_err() + { + tracing::warn!( + session_id = %session_id, + "Failed to send audio chunk — client disconnected" + ); + return; + } + } + + // Signal end of audio + let end_msg = serde_json::json!({ + "type": "audio_end", + "sample_rate": engine.sample_rate(), + "format": "pcm_s16le", + "channels": 1, + }); + let _ = sender + .send(Message::Text(end_msg.to_string().into())) + .await; + } + Err(e) => { + tracing::error!( + session_id = %session_id, + error = %e, + "TTS inference failed" + ); + let _ = send_error( + &mut sender, + "TTS_INFERENCE_FAILED", + &format!("TTS inference failed: {e}"), + ) + .await; + } + } + } + ClientMessage::Cancel => { + tracing::info!(session_id = %session_id, "TTS cancel requested"); + // TODO: support cancellation of in-progress inference + } + ClientMessage::Stop => { + tracing::info!(session_id = %session_id, "TTS stop requested, closing"); + break; + } + } + } + Message::Close(_) => { + tracing::info!(session_id = %session_id, "TTS WebSocket closed by client"); + break; + } + _ => { + // Ignore ping/pong/binary from client + } + } + } + + tracing::info!(session_id = %session_id, "TTS WebSocket connection closed"); +} + +/// Send an error message to the client. +async fn send_error<S>(sender: &mut S, code: &str, message: &str) -> Result<(), axum::Error> +where + S: SinkExt<Message> + Unpin, + <S as futures::Sink<Message>>::Error: std::error::Error, +{ + let error_msg = serde_json::json!({ + "type": "error", + "code": code, + "message": message, + "recoverable": false + }); + + sender + .send(Message::Text(error_msg.to_string().into())) + .await + .ok(); + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_error_message_format() { + let error = serde_json::json!({ + "type": "error", + "code": "TEST_ERROR", + "message": "Test message", + "recoverable": false + }); + + assert_eq!(error["type"], "error"); + assert_eq!(error["code"], "TEST_ERROR"); + assert_eq!(error["message"], "Test message"); + assert_eq!(error["recoverable"], false); + } + + #[test] + fn test_client_message_parse_speak() { + let json = r#"{"type": "speak", "text": "Hello world"}"#; + let msg: ClientMessage = serde_json::from_str(json).unwrap(); + match msg { + ClientMessage::Speak { text, voice } => { + assert_eq!(text, "Hello world"); + assert!(voice.is_none()); + } + _ => panic!("Expected Speak message"), + } + } + + #[test] + fn test_client_message_parse_cancel() { + let json = r#"{"type": "cancel"}"#; + let msg: ClientMessage = serde_json::from_str(json).unwrap(); + assert!(matches!(msg, ClientMessage::Cancel)); + } + + #[test] + fn test_client_message_parse_stop() { + let json = r#"{"type": "stop"}"#; + let msg: ClientMessage = serde_json::from_str(json).unwrap(); + assert!(matches!(msg, ClientMessage::Stop)); + } +} diff --git a/makima/src/server/messages.rs b/makima/src/server/messages.rs index 9c50334..cecb622 100644 --- a/makima/src/server/messages.rs +++ b/makima/src/server/messages.rs @@ -103,3 +103,164 @@ impl ApiError { } } } + +// ============================================================================= +// TTS (Text-to-Speech) Message Types +// ============================================================================= + +/// TTS audio encoding format for WebSocket streaming. +#[derive(Debug, Clone, Copy, Deserialize, Serialize, ToSchema, PartialEq, Default)] +#[serde(rename_all = "lowercase")] +pub enum TtsAudioEncoding { + /// 16-bit signed integer PCM samples + #[default] + Pcm16, + /// 32-bit floating point PCM samples + Pcm32f, +} + +/// TTS synthesis priority level. +#[derive(Debug, Clone, Copy, Deserialize, Serialize, ToSchema, PartialEq, Default)] +#[serde(rename_all = "lowercase")] +pub enum TtsPriority { + /// Low priority - may be queued + Low, + /// Normal priority (default) + #[default] + Normal, + /// High priority - processed immediately + High, +} + +/// TTS session start message from client. +#[derive(Debug, Clone, Deserialize, Serialize, ToSchema)] +#[serde(rename_all = "camelCase")] +pub struct TtsStartMessage { + /// Audio sample rate in Hz (default: 24000) + #[serde(default = "default_tts_sample_rate")] + pub sample_rate: u32, + /// Audio encoding format + #[serde(default)] + pub encoding: TtsAudioEncoding, + /// Voice identifier (default: "makima") + #[serde(default = "default_tts_voice")] + pub voice: String, + /// Language for synthesis (default: "English") + #[serde(default = "default_tts_language")] + pub language: String, +} + +fn default_tts_sample_rate() -> u32 { + 24000 +} + +fn default_tts_voice() -> String { + "makima".to_string() +} + +fn default_tts_language() -> String { + "English".to_string() +} + +/// TTS speak request message from client. +#[derive(Debug, Clone, Deserialize, Serialize, ToSchema)] +#[serde(rename_all = "camelCase")] +pub struct TtsSpeakMessage { + /// Text to synthesize (max 1000 characters) + pub text: String, + /// Synthesis priority + #[serde(default)] + pub priority: TtsPriority, +} + +/// TTS stop request message from client. +#[derive(Debug, Clone, Deserialize, Serialize, ToSchema)] +#[serde(rename_all = "camelCase")] +pub struct TtsStopMessage { + /// Optional reason for stopping + pub reason: Option<String>, +} + +/// Wrapper for all TTS WebSocket messages from client to server. +#[derive(Debug, Clone, Deserialize)] +#[serde(tag = "type", rename_all = "camelCase")] +pub enum TtsClientMessage { + /// Start a new TTS session + Start(TtsStartMessage), + /// Request speech synthesis + Speak(TtsSpeakMessage), + /// Stop the current session + Stop(TtsStopMessage), +} + +/// TTS session ready message sent from server to client. +#[derive(Debug, Clone, Serialize, ToSchema)] +#[serde(rename_all = "camelCase")] +pub struct TtsReadyMessage { + /// Unique session identifier + pub session_id: String, + /// Confirmed sample rate + pub sample_rate: u32, + /// Confirmed encoding format + pub encoding: TtsAudioEncoding, + /// Confirmed voice + pub voice: String, +} + +/// TTS audio chunk message sent from server to client. +#[derive(Debug, Clone, Serialize, ToSchema)] +#[serde(rename_all = "camelCase")] +pub struct TtsAudioChunkMessage { + /// Base64-encoded audio data + pub data: String, + /// Whether this is the final chunk + pub is_final: bool, + /// Timestamp in seconds from start of audio + pub timestamp: f64, +} + +/// TTS synthesis complete message sent from server to client. +#[derive(Debug, Clone, Serialize, ToSchema)] +#[serde(rename_all = "camelCase")] +pub struct TtsCompleteMessage { + /// Total synthesis duration in milliseconds + pub duration_ms: u64, + /// Total number of chunks sent + pub total_chunks: u32, + /// Length of input text + pub text_length: u32, +} + +/// TTS error message sent from server to client. +#[derive(Debug, Clone, Serialize, ToSchema)] +#[serde(rename_all = "camelCase")] +pub struct TtsErrorMessage { + /// Error code for programmatic handling + pub code: String, + /// Human-readable error message + pub message: String, +} + +/// TTS session stopped message sent from server to client. +#[derive(Debug, Clone, Serialize, ToSchema)] +#[serde(rename_all = "camelCase")] +pub struct TtsStoppedMessage { + /// Reason for stopping + pub reason: String, +} + +/// Wrapper for all TTS WebSocket messages from server to client. +#[derive(Debug, Clone, Serialize)] +#[serde(tag = "type", rename_all = "camelCase")] +pub enum TtsServerMessage { + /// Session is ready for synthesis requests + Ready(TtsReadyMessage), + /// Audio chunk (streamed during synthesis) + AudioChunk(TtsAudioChunkMessage), + /// Synthesis completed + Complete(TtsCompleteMessage), + /// Error occurred + Error(TtsErrorMessage), + /// Session has been stopped + Stopped(TtsStoppedMessage), +} diff --git a/makima/src/server/mod.rs b/makima/src/server/mod.rs index b969650..7c13f08 100644 --- a/makima/src/server/mod.rs +++ b/makima/src/server/mod.rs @@ -18,7 +18,7 @@ use tower_http::trace::TraceLayer; use utoipa::OpenApi; use utoipa_swagger_ui::SwaggerUi; -use crate::server::handlers::{api_keys, chat, contract_chat, contract_daemon, contracts, file_ws, files, history, listen, mesh, mesh_chat, mesh_daemon, mesh_merge, mesh_red_team, mesh_supervisor, mesh_ws, repository_history, templates, transcript_analysis, users, versions}; +use crate::server::handlers::{api_keys, chat, contract_chat, contract_daemon, contracts, file_ws, files, history, listen, mesh, mesh_chat, mesh_daemon, mesh_merge, mesh_red_team, mesh_supervisor, mesh_ws, repository_history, speak, templates, transcript_analysis, users, versions}; use crate::server::openapi::ApiDoc; use crate::server::state::SharedState; @@ -44,6 +44,7 @@ pub fn make_router(state: SharedState) -> Router { // API v1 routes let api_v1 = Router::new() .route("/listen", get(listen::websocket_handler)) + .route("/speak", get(speak::websocket_handler)) // Listen/transcript analysis endpoints .route("/listen/analyze", post(transcript_analysis::analyze_transcript)) .route("/listen/create-contract", post(transcript_analysis::create_contract_from_analysis)) diff --git a/makima/src/server/state.rs b/makima/src/server/state.rs index 1bc7d7e..bf8f6f2 100644 --- a/makima/src/server/state.rs +++ b/makima/src/server/state.rs @@ -8,6 +8,7 @@ use uuid::Uuid; use crate::listen::{DiarizationConfig, ParakeetEOU, ParakeetTDT, Sortformer}; use crate::server::auth::{AuthConfig, JwtVerifier}; +use crate::tts::TtsEngine; /// Notification payload for file updates (broadcast to WebSocket subscribers). #[derive(Debug, Clone)] @@ -599,6 +600,8 @@ pub struct AppState { pub jwt_verifier: Option<JwtVerifier>, /// Pending worktree info requests awaiting daemon response (keyed by task_id) pub pending_worktree_info: DashMap<Uuid, oneshot::Sender<WorktreeInfoResponse>>, + /// Lazily-loaded TTS engine (initialized on first Speak connection) + pub tts_engine: OnceCell<Box<dyn TtsEngine>>, } impl AppState { @@ -673,9 +676,28 @@ impl AppState { tool_keys: DashMap::new(), jwt_verifier, pending_worktree_info: DashMap::new(), + tts_engine: OnceCell::new(), } } + /// Get or initialize the TTS engine (lazy loading). + /// + /// The TTS engine is loaded on first Speak connection using the Qwen3 backend. + /// Returns a reference to the engine, or an error if loading fails. + pub async fn get_tts_engine(&self) -> Result<&dyn TtsEngine, Box<dyn std::error::Error + Send + Sync>> { + self.tts_engine.get_or_try_init(|| async { + tracing::info!("Lazy-loading TTS engine (Qwen3) on first Speak connection..."); + let engine = crate::tts::TtsEngineFactory::create( + crate::tts::TtsBackend::Qwen3, + None, // Use default model directory + ).map_err(|e| -> Box<dyn std::error::Error + Send + Sync> { + Box::new(e) + })?; + tracing::info!("TTS engine loaded successfully"); + Ok(engine) + }).await.map(|b| b.as_ref()) + } + /// Get or initialize ML models (lazy loading). /// /// Models are loaded on first call and cached for subsequent calls. |
