From 7d837e9f395bc202d3562379ea1a4289e4b74aa2 Mon Sep 17 00:00:00 2001
From: Damien Murphy <damien.murphy@deepgram.com>
Date: Wed, 24 Jul 2024 18:51:02 -0700
Subject: [PATCH] Add close message handling

---
 src/listen/websocket.rs | 24 ++++++++++++++++++------
 1 file changed, 18 insertions(+), 6 deletions(-)

diff --git a/src/listen/websocket.rs b/src/listen/websocket.rs
index 64584332..e7c41265 100644
--- a/src/listen/websocket.rs
+++ b/src/listen/websocket.rs
@@ -371,16 +371,28 @@ where
                 match read.next().await {
                     None => break,
                     Some(Ok(msg)) => {
-                        if let Message::Text(txt) = msg {
-                            let resp = serde_json::from_str(&txt).map_err(DeepgramError::from);
-                            if let Err(e) = tx.send(resp).await {
-                                eprintln!("Failed to send message: {:?}", e);
-                                // TODO Handle the error appropriately, e.g., log it, retry, or break the loop
+                        match msg {
+                            Message::Text(txt) => {
+                                let resp = serde_json::from_str(&txt).map_err(DeepgramError::from);
+                                if let Err(e) = tx.send(resp).await {
+                                    eprintln!("Failed to send message: {:?}", e);
+                                    // Handle the error appropriately, e.g., log it, retry, or break the loop
+                                    break;
+                                }
+                            }
+                            Message::Close(close_frame) => {
+                                println!("Received close frame: {:?}", close_frame);
+                                // Send a close frame back to acknowledge the close request
+                                let mut write = write.lock().await;
+                                if let Err(e) = write.send(Message::Close(None)).await {
+                                    eprintln!("Failed to send close frame: {:?}", e);
+                                }
                                 break;
                             }
+                            _ => {}
                         }
                     }
-                    Some(e) => {
+                    Some(Err(e)) => {
                         let _ = dbg!(e);
                         break;
                     }