fix: 修复 WebSocket 连接泄漏，使用 CancellationToken 实现优雅重启

Cargo.toml

@@ -13,6 +13,7 @@ serde_yaml = "0.9"
 async-trait = "0.1"
 thiserror = "2.0.18"
 tokio = { version = "1.0", features = ["full"] }
+tokio-util = { version = "0.7", features = ["rt"] }
 uuid = { version = "1.0", features = ["v4"] }
 axum = { version = "0.8", features = ["ws"] }
 tokio-tungstenite = { version = "0.29.0", features = ["rustls-tls-webpki-roots", "rustls"] }

src/channels/cli.rs

@@ -2,6 +2,7 @@ use async_trait::async_trait;
 use std::collections::HashMap;
 use std::sync::Arc;
 use tokio::sync::{RwLock, mpsc};
+use tokio_util::sync::CancellationToken;
 
 use crate::bus::{MessageBus, OutboundMessage};
 use crate::protocol::WsOutbound;
@@ -18,34 +19,36 @@ struct CliConnection {
 #[derive(Clone)]
 pub struct CliChannel {
     connections: Arc<RwLock<HashMap<String, CliConnection>>>,
+    /// 全局关闭信号：stop() 时 cancel，所有 handle_socket 同时收到信号
+    shutdown_token: CancellationToken,
 }
 
 impl CliChannel {
     pub fn new() -> Self {
         Self {
             connections: Arc::new(RwLock::new(HashMap::new())),
+            shutdown_token: CancellationToken::new(),
         }
     }
 
+    /// 注册连接，返回 CancellationToken 供 handle_socket 监听
+    ///
+    /// 注意：此方法不会向被替换的旧连接发送关闭信号。
+    /// 关闭信号仅由 stop() 统一触发，避免正常消息处理时误关连接。
     pub async fn register_connection(
         &self,
         session_id: impl Into<String>,
         connection_id: impl Into<String>,
         sender: mpsc::Sender<WsOutbound>,
-    ) {
-        let session_id = session_id.into();
-        let connection_id = connection_id.into();
-        let previous = self.connections.write().await.insert(
-            session_id.clone(),
+    ) -> CancellationToken {
+        self.connections.write().await.insert(
+            session_id.into(),
             CliConnection {
-                connection_id: connection_id.clone(),
+                connection_id: connection_id.into(),
                 sender,
             },
         );
-
-        if previous.is_some() {
-            tracing::info!(session_id = %session_id, connection_id = %connection_id, "CLI session sender replaced");
-        }
+        self.shutdown_token.clone()
     }
 
     pub async fn unregister_connection(&self, connection_id: &str) {
@@ -77,7 +80,11 @@ impl Channel for CliChannel {
     }
 
     async fn stop(&self) -> Result<(), ChannelError> {
+        // 先 cancel 所有 handle_socket，再清空 connections
+        self.shutdown_token.cancel();
+        let count = self.connections.read().await.len();
         self.connections.write().await.clear();
+        tracing::info!(connection_count = count, "CliChannel stopped, all connections signaled to close");
         Ok(())
     }
 
@@ -111,7 +118,7 @@ mod tests {
     async fn test_cli_channel_sends_to_registered_session() {
         let channel = CliChannel::new();
         let (sender, mut receiver) = mpsc::channel(4);
-        channel
+        let _token = channel
             .register_connection("session-1", "conn-1", sender)
             .await;
 
@@ -135,7 +142,7 @@ mod tests {
     async fn test_cli_channel_unregisters_connection_sessions() {
         let channel = CliChannel::new();
         let (sender, _receiver) = mpsc::channel(4);
-        channel
+        let _token = channel
             .register_connection("session-1", "conn-1", sender)
             .await;
         channel.unregister_connection("conn-1").await;
@@ -154,4 +161,36 @@ mod tests {
 
         assert!(error.to_string().contains("No active CLI connection"));
     }
+
+    #[tokio::test]
+    async fn test_stop_cancels_token() {
+        let channel = CliChannel::new();
+        let (sender, _receiver) = mpsc::channel(4);
+        let token = channel
+            .register_connection("session-1", "conn-1", sender)
+            .await;
+
+        assert!(!token.is_cancelled());
+        channel.stop().await.unwrap();
+        assert!(token.is_cancelled());
+    }
+
+    #[tokio::test]
+    async fn test_register_connection_does_not_cancel_old_token() {
+        // 验证：register_connection 替换旧连接时不会 cancel token
+        // 只有 stop() 才会 cancel
+        let channel = CliChannel::new();
+        let (sender1, _receiver1) = mpsc::channel(4);
+        let token1 = channel
+            .register_connection("session-1", "conn-1", sender1)
+            .await;
+
+        let (sender2, _receiver2) = mpsc::channel(4);
+        let _token2 = channel
+            .register_connection("session-1", "conn-2", sender2)
+            .await;
+
+        // 旧 token 不应被 cancel（只有 stop() 才 cancel）
+        assert!(!token1.is_cancelled());
+    }
 }

src/gateway/ws.rs

@@ -28,6 +28,7 @@ use crate::command::handlers::switch_topic::SwitchTopicCommandHandler;
 use crate::gateway::agent_prompt_provider::AgentPromptProvider;
 use crate::protocol::{WsInbound, WsOutbound, MediaSummary, parse_inbound, serialize_outbound};
 use crate::skills::SkillPromptProvider;
+use crate::storage::persistent_session_id;
 use crate::tools::task::repository::TaskRepository;
 use crate::tools::task::types::TaskSessionState;
 use axum::extract::State;
@@ -157,9 +158,10 @@ async fn handle_socket(ws: WebSocket, state: Arc<GatewayState>) {
     };
 
     let runtime_session_id = uuid::Uuid::new_v4().to_string();
-    let mut current_session_id = initial_record.id.clone();
+    // 清理数据库中可能已被污染的多重前缀 session_id（幂等处理）
+    let mut current_session_id = persistent_session_id(WS_CHANNEL_NAME, &initial_record.id);
     let mut current_topic_id: Option<String> = None;
-    state
+    let shutdown_token = state
         .channel_manager
         .websocket_channel()
         .register_connection(
@@ -232,7 +234,18 @@ async fn handle_socket(ws: WebSocket, state: Arc<GatewayState>) {
         }
     });
 
-    while let Some(msg) = ws_receiver.next().await {
+    loop {
+        tokio::select! {
+            // 监听 shutdown 信号（来自 CliChannel::stop()）
+            _ = shutdown_token.cancelled() => {
+                tracing::info!(session_id = %current_session_id, "WebSocket shutdown signal received, closing connection");
+                break;
+            }
+            // 监听 WebSocket 消息
+            msg = ws_receiver.next() => {
+                let Some(msg) = msg else {
+                    break;
+                };
                 match msg {
                     Ok(WsMessage::Text(text)) => {
                         let text = text.to_string();
@@ -278,6 +291,8 @@ async fn handle_socket(ws: WebSocket, state: Arc<GatewayState>) {
                     _ => {}
                 }
             }
+        }
+    }
 
     state
         .channel_manager
@@ -307,7 +322,7 @@ async fn handle_inbound(
             let chat_id = chat_id.unwrap_or_else(|| current_session_id.clone());
             let sender_id = resolve_ws_sender_id(sender_id.as_deref(), runtime_session_id);
 
-            state
+            let _ = state
                 .channel_manager
                 .websocket_channel()
                 .register_connection(
@@ -479,7 +494,7 @@ async fn handle_inbound(
                         "Updating current_session_id"
                     );
                     *current_session_id = session_id.clone();
-                    state
+                    let _ = state
                         .channel_manager
                         .websocket_channel()
                         .register_connection(