From 7640a6e5b7395526bb7c05ff09552d33660da033 Mon Sep 17 00:00:00 2001 From: oudecheng <13802883547@139.com> Date: Fri, 3 Jul 2026 16:32:43 +0800 Subject: [PATCH] =?UTF-8?q?fix:=20=E4=BF=AE=E5=A4=8D=20WebSocket=20?= =?UTF-8?q?=E8=BF=9E=E6=8E=A5=E6=B3=84=E6=BC=8F=EF=BC=8C=E4=BD=BF=E7=94=A8?= =?UTF-8?q?=20CancellationToken=20=E5=AE=9E=E7=8E=B0=E4=BC=98=E9=9B=85?= =?UTF-8?q?=E9=87=8D=E5=90=AF?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- Cargo.toml | 1 + src/channels/cli.rs | 63 +++++++++++++++++++++----- src/gateway/ws.rs | 107 +++++++++++++++++++++++++------------------- 3 files changed, 113 insertions(+), 58 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 2f2596e..99ef423 100644 --- a/Cargo.toml +++ b/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"] } diff --git a/src/channels/cli.rs b/src/channels/cli.rs index ee7703c..411ad4e 100644 --- a/src/channels/cli.rs +++ b/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>>, + /// 全局关闭信号：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, connection_id: impl Into, sender: mpsc::Sender, - ) { - 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()); + } } diff --git a/src/gateway/ws.rs b/src/gateway/ws.rs index 7510615..3352125 100644 --- a/src/gateway/ws.rs +++ b/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) { }; 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 = None; - state + let shutdown_token = state .channel_manager .websocket_channel() .register_connection( @@ -232,50 +234,63 @@ async fn handle_socket(ws: WebSocket, state: Arc) { } }); - while let Some(msg) = ws_receiver.next().await { - match msg { - Ok(WsMessage::Text(text)) => { - let text = text.to_string(); - match parse_inbound(&text) { - Ok(inbound) => { - if let Err(e) = handle_inbound( - &state, - &sender, - &runtime_session_id, - &mut current_session_id, - &mut current_topic_id, - inbound, - ) - .await - { - tracing::warn!(error = %e, session_id = %current_session_id, "Failed to handle inbound message"); - let _ = sender - .send(WsOutbound::Error { - timestamp: Some(crate::protocol::now_timestamp()), - code:"SESSION_ERROR".to_string(), - message: e.to_string(), - }) - .await; - } - } - Err(e) => { - tracing::warn!(error = %e, "Failed to parse inbound message"); - let _ = sender - .send(WsOutbound::Error { - timestamp: Some(crate::protocol::now_timestamp()), - code:"PARSE_ERROR".to_string(), - message: e.to_string(), - }) - .await; - } - } - } - Ok(WsMessage::Close(_)) | Err(_) => { - #[cfg(debug_assertions)] - tracing::debug!(session_id = %runtime_session_id, "WebSocket closed"); + 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(); + match parse_inbound(&text) { + Ok(inbound) => { + if let Err(e) = handle_inbound( + &state, + &sender, + &runtime_session_id, + &mut current_session_id, + &mut current_topic_id, + inbound, + ) + .await + { + tracing::warn!(error = %e, session_id = %current_session_id, "Failed to handle inbound message"); + let _ = sender + .send(WsOutbound::Error { + timestamp: Some(crate::protocol::now_timestamp()), + code:"SESSION_ERROR".to_string(), + message: e.to_string(), + }) + .await; + } + } + Err(e) => { + tracing::warn!(error = %e, "Failed to parse inbound message"); + let _ = sender + .send(WsOutbound::Error { + timestamp: Some(crate::protocol::now_timestamp()), + code:"PARSE_ERROR".to_string(), + message: e.to_string(), + }) + .await; + } + } + } + Ok(WsMessage::Close(_)) | Err(_) => { + #[cfg(debug_assertions)] + tracing::debug!(session_id = %runtime_session_id, "WebSocket closed"); + break; + } + _ => {} + } + } } } @@ -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(