perf: fix memory leaks, optimize tracing logs, fix networking eof fragmentation bug in nvim-core and stub, and deduplicate json-rpc handlers

server/src/api/ws.rs

@@ -0,0 +1,144 @@
+use crate::AppState;
+use axum::extract::{
+    Query, State,
+    ws::{Message, WebSocket},
+};
+use axum::response::IntoResponse;
+use futures_util::{SinkExt, StreamExt};
+use std::collections::HashMap;
+use std::sync::Arc;
+use std::sync::atomic::Ordering;
+use tokio::sync::mpsc;
+
+pub async fn ws_handler(
+    ws: axum::extract::ws::WebSocketUpgrade,
+    _headers: axum::http::HeaderMap,
+    State(state): State<Arc<AppState>>,
+    Query(query): Query<HashMap<String, String>>,
+) -> axum::response::Response {
+    let client_type = query
+        .get("client")
+        .cloned()
+        .unwrap_or_else(|| "unknown".to_string());
+    ws.on_upgrade(move |socket| handle_socket(socket, state, client_type))
+        .into_response()
+}
+
+pub struct SessionCleanup {
+    pub session_id: String,
+    pub state: Arc<AppState>,
+    pub send_task: tokio::task::JoinHandle<()>,
+    pub recv_task: tokio::task::JoinHandle<()>,
+}
+
+impl Drop for SessionCleanup {
+    fn drop(&mut self) {
+        tracing::info!("Dropping session {}", self.session_id);
+        self.state
+            .clients
+            .write()
+            .unwrap_or_else(|e| e.into_inner())
+            .remove(&self.session_id);
+        self.send_task.abort();
+        self.recv_task.abort();
+    }
+}
+
+pub async fn handle_socket(socket: WebSocket, state: Arc<AppState>, _client_type: String) {
+    let session_id = format!("{}", state.next_id.fetch_add(1, Ordering::SeqCst));
+    let (tx, mut rx) = mpsc::channel::<String>(100);
+
+    state
+        .clients
+        .write()
+        .unwrap_or_else(|e| e.into_inner())
+        .insert(session_id.clone(), tx.clone());
+
+    let (mut sender, mut receiver) = socket.split();
+
+    let send_task = tokio::spawn(async move {
+        while let Some(msg) = rx.recv().await {
+            tracing::trace!(
+                "Sending message to websocket (length: {}): {}",
+                msg.len(),
+                msg
+            );
+            if sender.send(Message::Text(msg.into())).await.is_err() {
+                tracing::error!("Failed to send message to websocket");
+                break;
+            }
+        }
+    });
+
+    let handler = Arc::clone(&state.handler);
+    let state_clone = Arc::clone(&state);
+    let session_id_clone = session_id.clone();
+
+    let recv_task = tokio::spawn(async move {
+        while let Some(msg_result) = receiver.next().await {
+            match msg_result {
+                Ok(Message::Text(text)) => {
+                    tracing::info!(
+                        "Received text message from websocket (length: {})",
+                        text.len()
+                    );
+                    tracing::trace!("Message content: {}", text);
+                    if let Ok(payload) = serde_json::from_str::<serde_json::Value>(&text) {
+                        // Process MCP request
+                        if let Some(response) = handler.handle_request(payload).await {
+                            let res_str = serde_json::to_string(&response).unwrap_or_default();
+                            let tx_opt = state_clone
+                                .clients
+                                .read()
+                                .unwrap_or_else(|e| e.into_inner())
+                                .get(&session_id_clone)
+                                .cloned();
+                            if let Some(client_tx) = tx_opt {
+                                if let Err(e) = client_tx.send(res_str).await {
+                                    tracing::error!(
+                                        "Failed to send response to client channel for session {}: {}",
+                                        session_id_clone,
+                                        e
+                                    );
+                                }
+                            } else {
+                                tracing::warn!(
+                                    "Could not find client_tx for session_id {} when trying to send response",
+                                    session_id_clone
+                                );
+                            }
+                        }
+                    } else {
+                        tracing::warn!(
+                            "Failed to parse payload as JSON from websocket message: {}",
+                            text
+                        );
+                    }
+                }
+                Ok(other) => {
+                    tracing::info!("Received non-text message from websocket: {:?}", other);
+                }
+                Err(e) => {
+                    tracing::error!("Websocket receive error: {}", e);
+                    break;
+                }
+            }
+        }
+    });
+
+    let mut cleanup = SessionCleanup {
+        session_id: session_id.clone(),
+        state: Arc::clone(&state),
+        send_task,
+        recv_task,
+    };
+
+    tokio::select! {
+        _ = &mut cleanup.send_task => {
+            tracing::info!("Websocket send task finished for session {}", session_id);
+        },
+        _ = &mut cleanup.recv_task => {
+            tracing::info!("Websocket recv task finished for session {}", session_id);
+        },
+    };
+}