fix ws connection state (#175)

Co-authored-by: Gregory Schier <gschier1990@gmail.com>

src-tauri/yaak-ws/src/commands.rs

@@ -2,7 +2,6 @@ use crate::error::Error::GenericError;
 use crate::error::Result;
 use crate::manager::WebsocketManager;
 use crate::render::render_request;
-use chrono::Utc;
 use log::{info, warn};
 use std::str::FromStr;
 use tauri::http::{HeaderMap, HeaderName};
@@ -147,42 +146,21 @@ pub(crate) async fn close<R: Runtime>(
     ws_manager: State<'_, Mutex<WebsocketManager>>,
 ) -> Result<WebsocketConnection> {
     let connection = get_websocket_connection(&window, connection_id).await?;
-    let request = get_websocket_request(&window, &connection.request_id)
-        .await?
-        .ok_or(GenericError("WebSocket Request not found".to_string()))?;
-
-    let mut ws_manager = ws_manager.lock().await;
-    if let Err(e) = ws_manager.send(&connection.id, Message::Close(None)).await {
-        warn!("Failed to close WebSocket connection: {e:?}");
-    };
-    upsert_websocket_event(
+    let connection = upsert_websocket_connection(
         &window,
-        WebsocketEvent {
-            connection_id: connection.id.clone(),
-            request_id: request.id.clone(),
-            workspace_id: request.workspace_id.clone(),
-            is_server: false,
-            message_type: WebsocketEventType::Close,
-            ..Default::default()
+        &WebsocketConnection {
+            state: WebsocketConnectionState::Closing,
+            ..connection
         },
         &UpdateSource::Window,
     )
     .await
     .unwrap();
 
-    let connection = upsert_websocket_connection(
-        &window,
-        &WebsocketConnection {
-            state: WebsocketConnectionState::Closed,
-            elapsed: Utc::now()
-                .naive_utc()
-                .signed_duration_since(connection.created_at)
-                .num_milliseconds() as i32,
-            ..connection.clone()
-        },
-        &UpdateSource::Window,
-    )
-    .await?;
+    let mut ws_manager = ws_manager.lock().await;
+    if let Err(e) = ws_manager.close(&connection.id).await {
+        warn!("Failed to close WebSocket connection: {e:?}");
+    };
 
     Ok(connection)
 }
@@ -264,42 +242,6 @@ pub(crate) async fn connect<R: Runtime>(
     let (receive_tx, mut receive_rx) = mpsc::channel::<Message>(128);
     let mut ws_manager = ws_manager.lock().await;
 
-    {
-        let connection_id = connection.id.clone();
-        let request_id = request.id.to_string();
-        let workspace_id = request.workspace_id.clone();
-        let window = window.clone();
-        tokio::spawn(async move {
-            while let Some(message) = receive_rx.recv().await {
-                upsert_websocket_event(
-                    &window,
-                    WebsocketEvent {
-                        connection_id: connection_id.clone(),
-                        request_id: request_id.clone(),
-                        workspace_id: workspace_id.clone(),
-                        is_server: true,
-                        message_type: match message {
-                            Message::Text(_) => WebsocketEventType::Text,
-                            Message::Binary(_) => WebsocketEventType::Binary,
-                            Message::Ping(_) => WebsocketEventType::Ping,
-                            Message::Pong(_) => WebsocketEventType::Pong,
-                            Message::Close(_) => WebsocketEventType::Close,
-                            // Raw frame will never happen during a read
-                            Message::Frame(_) => WebsocketEventType::Frame,
-                        },
-                        message: message.into_data().into(),
-                        ..Default::default()
-                    },
-                    &UpdateSource::Window,
-                )
-                .await
-                .unwrap();
-            }
-            info!("Websocket connection closed");
-        });
-    }
-
-
     let (url, url_parameters) = apply_path_placeholders(&request.url, request.url_parameters);
 
     // Add URL parameters to URL
@@ -331,6 +273,21 @@ pub(crate) async fn connect<R: Runtime>(
         }
     };
 
+    upsert_websocket_event(
+        &window,
+        WebsocketEvent {
+            connection_id: connection.id.clone(),
+            request_id: request.id.clone(),
+            workspace_id: connection.workspace_id.clone(),
+            is_server: false,
+            message_type: WebsocketEventType::Open,
+            ..Default::default()
+        },
+        &UpdateSource::Window,
+    )
+    .await
+    .unwrap();
+
     let response_headers = response
         .headers()
         .into_iter()
@@ -353,5 +310,74 @@ pub(crate) async fn connect<R: Runtime>(
     )
     .await?;
 
+    {
+        let connection_id = connection.id.clone();
+        let request_id = request.id.to_string();
+        let workspace_id = request.workspace_id.clone();
+        let window = window.clone();
+        let connection = connection.clone();
+        let mut has_written_close = false;
+        tokio::spawn(async move {
+            while let Some(message) = receive_rx.recv().await {
+                if let Message::Close(_) = message {
+                    has_written_close = true;
+                }
+
+                upsert_websocket_event(
+                    &window,
+                    WebsocketEvent {
+                        connection_id: connection_id.clone(),
+                        request_id: request_id.clone(),
+                        workspace_id: workspace_id.clone(),
+                        is_server: true,
+                        message_type: match message {
+                            Message::Text(_) => WebsocketEventType::Text,
+                            Message::Binary(_) => WebsocketEventType::Binary,
+                            Message::Ping(_) => WebsocketEventType::Ping,
+                            Message::Pong(_) => WebsocketEventType::Pong,
+                            Message::Close(_) => WebsocketEventType::Close,
+                            // Raw frame will never happen during a read
+                            Message::Frame(_) => WebsocketEventType::Frame,
+                        },
+                        message: message.into_data().into(),
+                        ..Default::default()
+                    },
+                    &UpdateSource::Window,
+                )
+                .await
+                .unwrap();
+            }
+            info!("Websocket connection closed");
+            if !has_written_close {
+                upsert_websocket_event(
+                    &window,
+                    WebsocketEvent {
+                        connection_id: connection_id.clone(),
+                        request_id: request_id.clone(),
+                        workspace_id: workspace_id.clone(),
+                        is_server: true,
+                        message_type: WebsocketEventType::Close,
+                        ..Default::default()
+                    },
+                    &UpdateSource::Window,
+                )
+                .await
+                .unwrap();
+            }
+            upsert_websocket_connection(
+                &window,
+                &WebsocketConnection {
+                    workspace_id: request.workspace_id.clone(),
+                    request_id: request_id.to_string(),
+                    state: WebsocketConnectionState::Closed,
+                    ..connection
+                },
+                &UpdateSource::Window,
+            )
+            .await
+            .unwrap();
+        });
+    }
+
     Ok(connection)
 }

src-tauri/yaak-ws/src/connect.rs

@@ -41,40 +41,4 @@ pub(crate) async fn ws_connect(
     )
     .await?;
     Ok((stream, response))
-}
-
-#[cfg(test)]
-mod tests {
-    use crate::connect::ws_connect;
-    use crate::error::Result;
-    use futures_util::{SinkExt, StreamExt};
-    use std::time::Duration;
-    use tokio::time::timeout;
-    use tokio_tungstenite::tungstenite::Message;
-
-    #[tokio::test]
-    async fn test_connection() -> Result<()> {
-        let (stream, response) = ws_connect("wss://echo.websocket.org/", Default::default()).await?;
-        assert_eq!(response.status(), 101);
-
-        let (mut write, mut read) = stream.split();
-
-        let task = tokio::spawn(async move {
-            while let Some(Ok(message)) = read.next().await {
-                if message.is_text() && message.to_text().unwrap() == "Hello" {
-                    return message;
-                }
-            }
-            panic!("Didn't receive text message");
-        });
-
-        write.send(Message::Text("Hello".into())).await?;
-
-        let task = timeout(Duration::from_secs(3), task);
-        let message = task.await.unwrap().unwrap();
-
-        assert_eq!(message.into_text().unwrap(), "Hello");
-
-        Ok(())
-    }
-}
+}

src-tauri/yaak-ws/src/lib.rs

@@ -5,7 +5,7 @@ mod manager;
 mod render;
 
 use crate::commands::{
-    connect, close, delete_connection, delete_connections, delete_request, duplicate_request,
+    close, connect, delete_connection, delete_connections, delete_request, duplicate_request,
     list_connections, list_events, list_requests, send, upsert_request,
 };
 use crate::manager::WebsocketManager;
@@ -31,7 +31,6 @@ pub fn init<R: Runtime>() -> TauriPlugin<R> {
         .setup(|app, _api| {
             let manager = WebsocketManager::new();
             app.manage(Mutex::new(manager));
-
             Ok(())
         })
         .build()

src-tauri/yaak-ws/src/manager.rs

@@ -2,7 +2,7 @@ use crate::connect::ws_connect;
 use crate::error::Result;
 use futures_util::stream::SplitSink;
 use futures_util::{SinkExt, StreamExt};
-use log::debug;
+use log::{debug, warn};
 use std::collections::HashMap;
 use std::sync::Arc;
 use tokio::net::TcpStream;
@@ -32,19 +32,27 @@ impl WebsocketManager {
         headers: HeaderMap<HeaderValue>,
         receive_tx: mpsc::Sender<Message>,
     ) -> Result<Response> {
+        let connections = self.connections.clone();
+        let connection_id = id.to_string();
+        let tx = receive_tx.clone();
+
         let (stream, response) = ws_connect(url, headers).await?;
         let (write, mut read) = stream.split();
-        self.connections.lock().await.insert(id.to_string(), write);
 
-        let tx = receive_tx.clone();
+        connections.lock().await.insert(id.to_string(), write);
+
         tauri::async_runtime::spawn(async move {
-            while let Some(Ok(message)) = read.next().await {
-                debug!("Received websocket message {message:?}");
-                if message.is_close() {
-                    return;
+            while let Some(msg) = read.next().await {
+                match msg {
+                    Err(e) => {
+                        warn!("Broken websocket connection: {}", e);
+                        break;
+                    }
+                    Ok(message) => tx.send(message).await.unwrap(),
                 }
-                tx.send(message).await.unwrap();
             }
+            debug!("Connection {} closed", connection_id);
+            connections.lock().await.remove(&connection_id);
         });
         Ok(response)
     }
@@ -59,4 +67,15 @@ impl WebsocketManager {
         connection.send(msg).await?;
         Ok(())
     }
+
+    pub async fn close(&mut self, id: &str) -> Result<()> {
+        debug!("Closing websocket");
+        let mut connections = self.connections.lock().await;
+        let connection = match connections.get_mut(id) {
+            None => return Ok(()),
+            Some(c) => c,
+        };
+        connection.close().await?;
+        Ok(())
+    }
 }