fix(listen): prevent panic on websocket keep-alive after close_stream (#144)

Closes #143

Author: lukeocodes

.gitignore

@@ -1,3 +1,4 @@
 /target
+.env
 your_output_file.wav
 .DS_Store

Cargo.toml

@@ -99,6 +99,11 @@ name = "microphone_stream"
 path = "examples/transcription/websocket/microphone_stream.rs"
 required-features = ["listen"]
 
+[[example]]
+name = "16_keepalive_close_stream"
+path = "examples/transcription/websocket/16_keepalive_close_stream.rs"
+required-features = ["listen"]
+
 [[example]]
 name = "simple_flux"
 path = "examples/transcription/flux/simple_flux.rs"

examples/transcription/websocket/16_keepalive_close_stream.rs

@@ -0,0 +1,63 @@
+/// Example: WebSocket Keep-Alive with Close Stream
+///
+/// Demonstrates that enabling keep_alive() and then calling close_stream()
+/// works correctly without panicking — even when the keep-alive timer fires
+/// after the channel has been closed.
+///
+/// This validates the fix for the race condition where close_stream() closes
+/// the internal channel, and a subsequent keep-alive send would previously
+/// panic on .expect().
+///
+/// Usage:
+///   DEEPGRAM_API_KEY=your-key cargo run --example 16_keepalive_close_stream
+use std::env;
+use std::time::Duration;
+
+use deepgram::{
+    common::options::{Encoding, Endpointing, Options},
+    Deepgram, DeepgramError,
+};
+
+#[tokio::main]
+async fn main() -> Result<(), DeepgramError> {
+    let deepgram_api_key =
+        env::var("DEEPGRAM_API_KEY").expect("DEEPGRAM_API_KEY environment variable");
+
+    let dg_client = Deepgram::new(&deepgram_api_key)?;
+
+    let options = Options::builder()
+        .query_params([("mip_opt_out".to_string(), "true".to_string())])
+        .build();
+
+    println!("Connecting to Deepgram with keep-alive enabled...");
+
+    let mut handle = dg_client
+        .transcription()
+        .stream_request_with_options(options)
+        .encoding(Encoding::Linear16)
+        .endpointing(Endpointing::Disabled)
+        .keep_alive()
+        .handle()
+        .await?;
+
+    println!("Connected. Request ID: {}", handle.request_id());
+
+    // Brief pause — no audio is sent, so the worker only has the keep-alive
+    // timer ticking. This simulates a connection that is idle before being
+    // torn down.
+    tokio::time::sleep(Duration::from_millis(100)).await;
+
+    println!("Closing stream...");
+    handle.close_stream().await?;
+    println!("Stream closed.");
+
+    // Wait longer than the 3s keep-alive interval so the keep-alive send
+    // path runs after close. Before the fix this would panic; now the send
+    // error is silently ignored and the worker exits cleanly.
+    println!("Waiting for keep-alive timer to fire (>3s)...");
+    tokio::time::sleep(Duration::from_secs(4)).await;
+
+    println!("No panic — keep-alive + close_stream works correctly.");
+
+    Ok(())
+}

src/listen/websocket.rs

@@ -481,7 +481,12 @@ async fn run_worker(
             _ = sleep.fuse() => {
                 // eprintln!("<worker> sleep");
                 if keep_alive && is_open {
-                    message_tx.send(WsMessage::ControlMessage(ControlMessage::KeepAlive)).await.expect("we hold the receiver, so we know it hasn't been dropped");
+                    // Ignore send errors: the channel may have been closed by
+                    // close_stream() (via close_channel()) before the worker
+                    // processes the pending CloseStream message. In that case
+                    // the next iteration will handle CloseStream, stop sending new
+                    // messages, and proceed toward shutdown.
+                    let _ = message_tx.send(WsMessage::ControlMessage(ControlMessage::KeepAlive)).await;
                     last_sent_message = tokio::time::Instant::now();
                 } else {
                     pending::<()>().await;
@@ -872,8 +877,10 @@ mod file_chunker {
 
 #[cfg(test)]
 mod tests {
+    use std::time::Duration;
+
     use super::ControlMessage;
-    use crate::common::options::Options;
+    use crate::common::options::{Encoding, Endpointing, Options};
 
     #[test]
     fn test_stream_url() {
@@ -910,4 +917,78 @@ mod tests {
             r#"{"type":"CloseStream"}"#
         );
     }
+
+    /// Reproduces the worker panic from issue #143: close_stream() calls
+    /// close_channel(), so when the worker's keep-alive sleep fires it sends
+    /// into a closed channel. Before the fix, .expect() would panic.
+    #[tokio::test]
+    #[ignore = "requires DEEPGRAM_API_KEY and network; run manually"]
+    async fn keepalive_then_close_stream_panic_repro() {
+        let Ok(api_key) = std::env::var("DEEPGRAM_API_KEY") else {
+            eprintln!("skipping: DEEPGRAM_API_KEY not set");
+            return;
+        };
+
+        let dg = crate::Deepgram::new(&api_key).expect("Deepgram::new");
+        let transcription = dg.transcription();
+
+        let options = Options::builder()
+            .query_params([("mip_opt_out".to_string(), "true".to_string())])
+            .build();
+        let mut handle = transcription
+            .stream_request_with_options(options)
+            .encoding(Encoding::Linear16)
+            .endpointing(Endpointing::Disabled)
+            .keep_alive()
+            .handle()
+            .await
+            .expect("handle");
+
+        // No audio sent: worker only has the keep-alive timer (3s interval).
+        // Close the channel before the keep-alive fires so the
+        // send(KeepAlive) hits a closed channel.
+        tokio::time::sleep(Duration::from_millis(100)).await;
+        let _ = handle.close_stream().await;
+
+        // Wait longer than the 3s keep-alive interval so the keep-alive
+        // send path runs after close. Before the fix this would panic.
+        tokio::time::sleep(Duration::from_secs(4)).await;
+    }
+
+    /// Runs the close_stream race in a loop to increase the probability of
+    /// hitting the timing window (scheduling variance).
+    #[tokio::test]
+    #[ignore = "requires DEEPGRAM_API_KEY and network; run manually"]
+    async fn keepalive_close_stream_panic_repro_loop() {
+        let Ok(api_key) = std::env::var("DEEPGRAM_API_KEY") else {
+            eprintln!("skipping: DEEPGRAM_API_KEY not set");
+            return;
+        };
+
+        const ITERATIONS: u32 = 3;
+
+        let options = Options::builder()
+            .query_params([("mip_opt_out".to_string(), "true".to_string())])
+            .build();
+        let dg = crate::Deepgram::new(&api_key).expect("Deepgram::new");
+        let transcription = dg.transcription();
+
+        for _iteration in 0..ITERATIONS {
+            let mut handle = transcription
+                .stream_request_with_options(options.clone())
+                .encoding(Encoding::Linear16)
+                .endpointing(Endpointing::Disabled)
+                .keep_alive()
+                .handle()
+                .await
+                .expect("handle");
+
+            tokio::time::sleep(Duration::from_millis(100)).await;
+            let _ = handle.close_stream().await;
+
+            // Wait longer than the 3s keep-alive interval so the keep-alive
+            // send path runs after close. Before the fix this would panic.
+            tokio::time::sleep(Duration::from_secs(4)).await;
+        }
+    }
 }

tests/flux_unknown_messages.rs

@@ -25,6 +25,7 @@ mod mock {
         tokio::spawn(async move {
             let (stream, _) = listener.accept().await.unwrap();
 
+            #[allow(clippy::result_large_err)]
             let callback =
                 |_req: &tungstenite::handshake::server::Request,
                  mut resp: tungstenite::handshake::server::Response| {