potential fix for local-stt-server connection issues

Author: yujonglee

Cargo.lock

@@ -98,6 +98,7 @@ reqwest = "0.12"
 tokio = "1"
 tokio-stream = "0.1.17"
 tokio-tungstenite = "0.26.0"
+tokio-util = "0.7.15"
 
 anyhow = "1"
 approx = "0.5.1"

Cargo.toml

@@ -35,14 +35,6 @@ impl WebSocketClient {
                     .with_delay(std::time::Duration::from_millis(500)),
             )
             .when(|e| {
-                if let crate::Error::Connection(te) = e {
-                    if let tokio_tungstenite::tungstenite::Error::Http(res) = te {
-                        if res.status() == 429 {
-                            return true;
-                        }
-                    }
-                }
-
                 tracing::error!("ws_connect_failed: {:?}", e);
                 true
             })
@@ -105,7 +97,7 @@ impl WebSocketClient {
         tracing::info!("connect_async: {:?}", req.uri());
 
         let (ws_stream, _) =
-            tokio::time::timeout(std::time::Duration::from_secs(4), connect_async(req)).await??;
+            tokio::time::timeout(std::time::Duration::from_secs(8), connect_async(req)).await??;
 
         Ok(ws_stream)
     }

crates/ws/src/client.rs

@@ -44,9 +44,11 @@ thiserror = { workspace = true }
 tracing = { workspace = true }
 
 axum = { workspace = true, features = ["ws"] }
+tower-http = { workspace = true, features = ["cors", "trace"] }
+
 futures-util = { workspace = true }
 tokio = { workspace = true, features = ["rt", "macros"] }
-tower-http = { workspace = true, features = ["cors", "trace"] }
+tokio-util = { workspace = true }
 
 [target.'cfg(not(target_os = "macos"))'.dependencies]
 kalosm-sound = { workspace = true, default-features = false }

plugins/local-stt/Cargo.toml

@@ -1,43 +1,41 @@
-use std::{
-    sync::atomic::{AtomicUsize, Ordering},
-    sync::Arc,
-};
+use std::sync::{Arc, Mutex};
+use tokio_util::sync::CancellationToken;
 
 #[derive(Clone)]
 pub struct ConnectionManager {
-    num_connections: Arc<AtomicUsize>,
+    inner: Arc<Mutex<Option<CancellationToken>>>,
 }
 
 impl Default for ConnectionManager {
     fn default() -> Self {
         Self {
-            num_connections: Arc::new(AtomicUsize::new(0)),
+            inner: Arc::new(Mutex::new(None)),
         }
     }
 }
 
 impl ConnectionManager {
-    pub fn try_acquire_connection(&self) -> Option<ConnectionGuard> {
-        let current = self.num_connections.load(Ordering::SeqCst);
-        if current >= 1 {
-            return None;
-        }
+    pub fn acquire_connection(&self) -> ConnectionGuard {
+        let mut slot = self.inner.lock().unwrap();
 
-        match self
-            .num_connections
-            .compare_exchange(0, 1, Ordering::SeqCst, Ordering::SeqCst)
-        {
-            Ok(_) => Some(ConnectionGuard(self.num_connections.clone())),
-            Err(_) => None,
+        if let Some(old) = slot.take() {
+            old.cancel();
         }
+
+        let token = CancellationToken::new();
+        *slot = Some(token.clone());
+
+        ConnectionGuard { token }
     }
 }
 
-pub struct ConnectionGuard(Arc<AtomicUsize>);
+pub struct ConnectionGuard {
+    token: CancellationToken,
+}
 
-impl Drop for ConnectionGuard {
-    fn drop(&mut self) {
-        self.0.fetch_sub(1, Ordering::SeqCst);
+impl ConnectionGuard {
+    pub async fn cancelled(&self) {
+        self.token.cancelled().await
     }
 }
 
@@ -70,9 +68,7 @@ mod tests {
         ws: WebSocketUpgrade,
         AxumState(manager): AxumState<ConnectionManager>,
     ) -> Result<impl IntoResponse, StatusCode> {
-        let guard = manager
-            .try_acquire_connection()
-            .ok_or(StatusCode::TOO_MANY_REQUESTS)?;
+        let guard = manager.acquire_connection();
 
         Ok(ws.on_upgrade(move |socket| handle_socket(socket, guard)))
     }

plugins/local-stt/src/manager.rs

@@ -56,12 +56,6 @@ pub struct ServerState {
     connection_manager: ConnectionManager,
 }
 
-impl ServerState {
-    pub fn try_acquire_connection(&self) -> Option<ConnectionGuard> {
-        self.connection_manager.try_acquire_connection()
-    }
-}
-
 #[derive(Clone)]
 pub struct ServerHandle {
     pub addr: SocketAddr,
@@ -114,13 +108,25 @@ async fn listen(
     ws: WebSocketUpgrade,
     AxumState(state): AxumState<ServerState>,
 ) -> Result<impl IntoResponse, StatusCode> {
-    let guard = state
-        .try_acquire_connection()
-        .ok_or(StatusCode::TOO_MANY_REQUESTS)?;
+    let guard = state.connection_manager.acquire_connection();
 
-    let model_path = state.model_type.model_path(&state.model_cache_dir);
+    Ok(ws.on_upgrade(move |socket| async move {
+        websocket_with_model(socket, params, state, guard).await
+    }))
+}
+
+async fn websocket_with_model(
+    socket: WebSocket,
+    params: ListenParams,
+    state: ServerState,
+    guard: ConnectionGuard,
+) {
+    let model_type = state.model_type;
+    let model_cache_dir = state.model_cache_dir.clone();
+
+    let model_path = model_type.model_path(&model_cache_dir);
     let language = params.language.try_into().unwrap_or_else(|e| {
-        tracing::error!("convert_to_whisper_language: {:?}", e);
+        tracing::error!("convert_to_whisper_language: {e:?}");
         hypr_whisper::Language::En
     });
 
@@ -131,15 +137,11 @@ async fn listen(
         .dynamic_prompt(&params.dynamic_prompt)
         .build();
 
-    Ok(ws.on_upgrade(move |socket| websocket(socket, model, guard)))
+    websocket(socket, model, guard).await;
 }
 
 #[tracing::instrument(skip_all)]
-async fn websocket(
-    socket: WebSocket,
-    model: hypr_whisper::local::Whisper,
-    _guard: ConnectionGuard,
-) {
+async fn websocket(socket: WebSocket, model: hypr_whisper::local::Whisper, guard: ConnectionGuard) {
     let (mut ws_sender, ws_receiver) = socket.split();
     let mut stream = {
         let audio_source = WebSocketAudioSource::new(ws_receiver, 16 * 1000);
@@ -148,35 +150,44 @@ async fn websocket(
         hypr_whisper::local::TranscribeChunkedAudioStreamExt::transcribe(chunked, model)
     };
 
-    while let Some(chunk) = stream.next().await {
-        let text = chunk.text().to_string();
-        let start = chunk.start() as u64;
-        let duration = chunk.duration() as u64;
-        let confidence = chunk.confidence();
-
-        if confidence < 0.45 {
-            tracing::warn!(confidence, "skipping_transcript: {}", text);
-            continue;
-        }
-
-        let data = ListenOutputChunk {
-            words: text
-                .split_whitespace()
-                .filter(|w| !w.is_empty())
-                .map(|w| Word {
-                    text: w.trim().to_string(),
-                    speaker: None,
-                    start_ms: Some(start),
-                    end_ms: Some(start + duration),
-                    confidence: Some(confidence),
-                })
-                .collect(),
-        };
-
-        let msg = Message::Text(serde_json::to_string(&data).unwrap().into());
-        if let Err(e) = ws_sender.send(msg).await {
-            tracing::warn!("websocket_send_error: {}", e);
-            break;
+    loop {
+        tokio::select! {
+            _ = guard.cancelled() => {
+                tracing::info!("websocket_cancelled_by_new_connection");
+                break;
+            }
+            chunk_opt = stream.next() => {
+                let Some(chunk) = chunk_opt else { break };
+                let text = chunk.text().to_string();
+                let start = chunk.start() as u64;
+                let duration = chunk.duration() as u64;
+                let confidence = chunk.confidence();
+
+                if confidence < 0.45 {
+                    tracing::warn!(confidence, "skipping_transcript: {}", text);
+                    continue;
+                }
+
+                let data = ListenOutputChunk {
+                    words: text
+                        .split_whitespace()
+                        .filter(|w| !w.is_empty())
+                        .map(|w| Word {
+                            text: w.trim().to_string(),
+                            speaker: None,
+                            start_ms: Some(start),
+                            end_ms: Some(start + duration),
+                            confidence: Some(confidence),
+                        })
+                        .collect(),
+                };
+
+                let msg = Message::Text(serde_json::to_string(&data).unwrap().into());
+                if let Err(e) = ws_sender.send(msg).await {
+                    tracing::warn!("websocket_send_error: {}", e);
+                    break;
+                }
+            }
         }
     }