fix(server): Fix memory leak in handle_socket on abrupt disconnects

- Used a SessionCleanup Drop guard (RAII pattern) to guarantee that the websocket session is evicted from the global clients map and background read/write tasks are forcibly aborted even when Axum terminates the handler mid-execution due to a TCP connection loss.

server/src/main.rs

@@ -557,22 +557,45 @@ async fn handle_socket(socket: WebSocket, state: Arc<AppState>, client_type: Str
         );
     });
 
-    tokio::select! {
+    struct SessionCleanup {
-        _ = (&mut send_task) => {
+        session_id: String,
-            tracing::info!("Websocket send task finished for session {}", session_id);
+        state: Arc<AppState>,
-            recv_task.abort();
+        send_task: Option<tokio::task::JoinHandle<()>>,
-        },
+        recv_task: Option<tokio::task::JoinHandle<()>>,
-        _ = (&mut recv_task) => {
+    }
-            tracing::info!("Websocket recv task finished for session {}", session_id);
+
-            send_task.abort();
+    impl Drop for SessionCleanup {
-        },
+        fn drop(&mut self) {
+            self.state.clients.write().unwrap().remove(&self.session_id);
+            if let Some(task) = self.send_task.take() {
+                task.abort();
+            }
+            if let Some(task) = self.recv_task.take() {
+                task.abort();
+            }
+            tracing::info!(
+                "Websocket session {} closed and cleaned up",
+                self.session_id
+            );
+        }
+    }
+
+    let mut cleanup = SessionCleanup {
+        session_id: session_id.clone(),
+        state: Arc::clone(&state),
+        send_task: Some(send_task),
+        recv_task: Some(recv_task),
     };
 
-    state.clients.write().unwrap().remove(&session_id);
+    tokio::select! {
-    tracing::info!(
+        _ = cleanup.send_task.as_mut().unwrap() => {
-        "Websocket session {} closed and removed from state",
+            tracing::info!("Websocket send task finished for session {}", session_id);
-        session_id
+        },
-    );
+        _ = cleanup.recv_task.as_mut().unwrap() => {
+            tracing::info!("Websocket recv task finished for session {}", session_id);
+        },
+    };
+    // Drop guard automatically handles removal and aborts the other task.
 }
 
 #[derive(serde::Deserialize, serde::Serialize, Debug)]

stub/src/main.rs

@@ -117,20 +117,23 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
                     });
 
                     tokio::select! {
-                        _ = shutdown_rx.recv() => { tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
+                        _ = shutdown_rx.recv() => { 
+                            tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
                             tracing::info!("Shutdown received while connected");
-                            
+                            break;
                         }
                         _ = &mut send_task => {
                             tracing::error!("Send task exited");
-                            tokio::time::sleep(tokio::time::Duration::from_millis(50)).await; recv_task.abort();
+                            tokio::time::sleep(tokio::time::Duration::from_millis(50)).await; 
+                            recv_task.abort();
+                            break;
                         }
                         _ = &mut recv_task => {
                             tracing::error!("Recv task exited");
                             send_task.abort();
                         }
                     }
-                    tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
+                    tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
                 }
                 Err(e) => {
                     tracing::error!("Failed to connect via WSS: {}", e);