ooodc/PicoBot
ooodc/PicoBot
1
0
Fork 0
Code Pull Requests Activity

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

Browse Source
This commit is contained in:
oudecheng 2026-07-03 16:32:43 +08:00
parent 7e24e57af4
commit 7640a6e5b7
3 changed files with 113 additions and 58 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
Cargo.toml
View File

@ -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"] }

63
src/channels/cli.rs
View File

@ -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());
}
}

107
src/gateway/ws.rs
View File

@ -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,50 +234,63 @@ async fn handle_socket(ws: WebSocket, state: Arc<GatewayState>) {
}
});
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(