From ebbececdc93a2b7551fbe3796df27212a2fe5811 Mon Sep 17 00:00:00 2001
From: Andras Schmelczer <andras@schmelczer.dev>
Date: Sat, 7 Jun 2025 17:15:56 +0100
Subject: [PATCH] Nicer log line on ws disconnect

---
 backend/sync_server/src/errors.rs           | 23 ++++++--------
 backend/sync_server/src/server/websocket.rs | 35 +++++++++++++--------
 2 files changed, 31 insertions(+), 27 deletions(-)

diff --git a/backend/sync_server/src/errors.rs b/backend/sync_server/src/errors.rs
index a16f7137..aa1421ef 100644
--- a/backend/sync_server/src/errors.rs
+++ b/backend/sync_server/src/errors.rs
@@ -1,13 +1,11 @@
 use std::fmt::Display;
 
-use aide::OperationOutput;
 use axum::{
     Json,
     http::StatusCode,
     response::{IntoResponse, Response},
 };
-use log::{error, info};
-use schemars::JsonSchema;
+use log::{debug, error};
 use serde::Serialize;
 use thiserror::Error;
 
@@ -45,7 +43,7 @@ impl SyncServerError {
     }
 }
 
-#[derive(Debug, Clone, Serialize, JsonSchema)]
+#[derive(Debug, Clone, Serialize)]
 pub struct SerializedError {
     pub message: String,
     pub causes: Vec<String>,
@@ -96,35 +94,32 @@ impl From<&anyhow::Error> for SerializedError {
     }
 }
 
-impl OperationOutput for SyncServerError {
-    type Inner = Self;
-}
-
-pub const fn init_error(error: anyhow::Error) -> SyncServerError {
+pub fn init_error(error: anyhow::Error) -> SyncServerError {
+    debug!("Initialization error: {error:?}");
     SyncServerError::InitError(error)
 }
 
 pub fn server_error(error: anyhow::Error) -> SyncServerError {
-    error!("Server error: {error:?}");
+    debug!("Server error: {error:?}");
     SyncServerError::ServerError(error)
 }
 
 pub fn client_error(error: anyhow::Error) -> SyncServerError {
-    info!("Client error: {error:?}");
+    debug!("Client error: {error:?}");
     SyncServerError::ClientError(error)
 }
 
 pub fn not_found_error(error: anyhow::Error) -> SyncServerError {
-    info!("Not found: {error:?}");
+    debug!("Not found: {error:?}");
     SyncServerError::NotFound(error)
 }
 
 pub fn unauthenticated_error(error: anyhow::Error) -> SyncServerError {
-    info!("Unauthenticated user: {error:?}");
+    debug!("Unauthenticated user: {error:?}");
     SyncServerError::Unauthenticated(error)
 }
 
 pub fn permission_denied_error(error: anyhow::Error) -> SyncServerError {
-    info!("Permission denied: {error:?}");
+    debug!("Permission denied: {error:?}");
     SyncServerError::PermissionDeniedError(error)
 }
diff --git a/backend/sync_server/src/server/websocket.rs b/backend/sync_server/src/server/websocket.rs
index 4a7d7833..cfe7f8f9 100644
--- a/backend/sync_server/src/server/websocket.rs
+++ b/backend/sync_server/src/server/websocket.rs
@@ -7,8 +7,7 @@ use axum::{
     response::Response,
 };
 use futures::stream::StreamExt;
-use log::{error, info, warn};
-use schemars::JsonSchema;
+use log::{debug, info};
 use serde::Deserialize;
 
 use crate::{
@@ -29,8 +28,7 @@ use crate::{
     utils::normalize::normalize,
 };
 
-// This is required for aide to infer the path parameter types and names
-#[derive(Deserialize, JsonSchema)]
+#[derive(Deserialize)]
 pub struct WebSocketPathParams {
     #[serde(deserialize_with = "normalize")]
     vault_id: VaultId,
@@ -50,10 +48,8 @@ async fn websocket_wrapped(state: AppState, stream: WebSocket, vault_id: VaultId
     let result = websocket(state, stream, vault_id.clone()).await;
 
     if let Err(err) = result {
-        error!("WebSocket connection error on vault '{vault_id}': {err}");
+        debug!("WebSocket connection error on vault '{vault_id}': {err}");
     }
-
-    warn!("WebSocket connection closed on vault '{vault_id}'");
 }
 
 async fn websocket(
@@ -73,6 +69,11 @@ async fn websocket(
             .unwrap_or_default(),
     )?;
 
+    info!(
+        "WebSocket handshake successful for vault '{vault_id}' for '{}'",
+        handshake.device_id
+    );
+
     let mut broadcast_receiver = state.broadcasts.get_receiver(vault_id.clone()).await;
 
     send_update_over_websocket(
@@ -141,26 +142,34 @@ async fn websocket(
         _ = &mut receive_task => send_task.abort(),
     };
 
-    let result = {
+    let result: Result<(), SyncServerError> = (async {
         send_task
             .await
             .context("WebSocket send task failed")
-            .map_err(server_error)
-            .and_then(|x| x)?;
+            .map_err(client_error)
+            .and_then(|err| err)?;
 
         receive_task
             .await
             .context("WebSocket receive task failed")
-            .map_err(server_error)
-            .and_then(|x| x)?;
+            .map_err(client_error)
+            .and_then(|err| err)?;
 
         Ok(())
-    };
+    })
+    .await;
 
     state
         .cursors
         .remove_cursors_of_device(&vault_id, &handshake.device_id)
         .await;
 
+    if result.is_err() {
+        info!(
+            "WebSocket disconnected on vault '{vault_id}' for '{}'",
+            handshake.device_id
+        );
+    }
+
     result
 }