From f42e9d44cc392a73b48eb737e1db63441f3fe5cd Mon Sep 17 00:00:00 2001 From: xiaoxixi Date: Tue, 14 Jul 2026 11:56:19 +0800 Subject: [PATCH] refactor(session): isolate messaging and persistence --- src/session/messenger.rs | 113 +++++++++++++++++++ src/session/mod.rs | 2 + src/session/persistence.rs | 65 +++++++++++ src/session/session.rs | 219 +++++-------------------------------- src/session/session_id.rs | 22 ---- 5 files changed, 205 insertions(+), 216 deletions(-) create mode 100644 src/session/messenger.rs create mode 100644 src/session/persistence.rs diff --git a/src/session/messenger.rs b/src/session/messenger.rs new file mode 100644 index 0000000..104c161 --- /dev/null +++ b/src/session/messenger.rs @@ -0,0 +1,113 @@ +use std::collections::HashMap; + +use crate::bus::{ChatMessage, MediaItem, MessageSource, OutboundMessage}; +use crate::session::UnifiedSessionId; +use crate::tools::OutboundMessenger; + +use super::persistence::append_persisted_messages; +use super::session::{CURRENT_SOURCE_SESSION, SessionManager}; + +#[async_trait::async_trait] +impl OutboundMessenger for SessionManager { + async fn send_message( + &self, + channel: &str, + chat_id: &str, + dialog_id: Option<&str>, + content: &str, + mut source: MessageSource, + media: Vec, + ) -> Result<(), String> { + if source.from_session.is_none() { + source.from_session = CURRENT_SOURCE_SESSION + .try_with(|value| value.clone()) + .ok() + .flatten(); + } + + let (target_sid, session) = if let Some(dialog_id) = dialog_id { + let session_id = UnifiedSessionId::new(channel, chat_id, dialog_id); + let session = self + .get_or_activate_session(&session_id) + .await + .map_err(|error| error.to_string())?; + (session_id, session) + } else { + let session_id = self + .resolve_dialog_id(channel, chat_id) + .await + .map_err(|error| error.to_string())?; + let session = self + .get_or_create_session(&session_id) + .await + .map_err(|error| error.to_string())?; + (session_id, session) + }; + + let origin = source.from_session.as_deref().unwrap_or("unknown"); + let origin_id = source.from_session.clone(); + let same_session = source.from_session.as_deref() == Some(target_sid.to_string().as_str()); + let marked_content = if content.trim().is_empty() && !media.is_empty() && same_session { + String::new() + } else { + format!("[message from {origin}] \n{content}") + }; + + let message = outbound_history_message(marked_content.clone(), source, &media); + append_persisted_messages(&session, vec![message]) + .await + .map_err(|error| error.to_string())?; + + if let Some(origin_id) = origin_id { + self.restore_origin_dialog(&origin_id, &target_sid).await; + } + + self.bus + .deliver_outbound(OutboundMessage { + channel: channel.to_string(), + chat_id: chat_id.to_string(), + content: marked_content, + reply_to: None, + media, + metadata: HashMap::new(), + delivery: None, + }) + .await + .map_err(|error| error.to_string()) + } +} + +fn outbound_history_message( + content: impl Into, + source: MessageSource, + media: &[MediaItem], +) -> ChatMessage { + let mut message = ChatMessage::assistant_with_source(content, source); + message.media_refs = media.iter().map(MediaItem::to_media_ref).collect(); + message +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::bus::SourceKind; + + #[test] + fn outbound_history_preserves_delivered_media() { + let source = MessageSource { + kind: SourceKind::CrossChannel, + from_channel: Some("cli_chat".into()), + from_session: Some("cli_chat:source:dialog".into()), + from_user_id: None, + system_name: None, + task_id: None, + }; + let media = vec![MediaItem::new("/tmp/report.pdf", "file")]; + + let message = outbound_history_message("report", source, &media); + + assert_eq!(message.media_refs.len(), 1); + assert_eq!(message.media_refs[0].path, "/tmp/report.pdf"); + assert_eq!(message.media_refs[0].media_type, "file"); + } +} diff --git a/src/session/mod.rs b/src/session/mod.rs index c3686f4..0445afa 100644 --- a/src/session/mod.rs +++ b/src/session/mod.rs @@ -1,6 +1,8 @@ pub mod commands; pub mod error; pub mod events; +mod messenger; +mod persistence; // The public `session::session` path is retained for API compatibility. #[allow(clippy::module_inception)] pub mod session; diff --git a/src/session/persistence.rs b/src/session/persistence.rs new file mode 100644 index 0000000..303925a --- /dev/null +++ b/src/session/persistence.rs @@ -0,0 +1,65 @@ +use std::sync::Arc; + +use tokio::sync::Mutex; + +use super::session::{MessagePersistSnapshot, Session}; +use crate::bus::ChatMessage; +use crate::storage::StorageError; + +async fn persist_added_messages( + snapshots: Vec>, +) -> Result<(), StorageError> { + let mut storage = None; + let mut session_id = None; + let mut messages = Vec::new(); + let mut final_meta = None; + + for snapshot in snapshots.into_iter().flatten() { + let (snapshot_storage, snapshot_session_id, message, meta) = snapshot; + if let Some(ref expected) = session_id + && expected != &snapshot_session_id + { + return Err(StorageError::Serialization( + "attempted to persist messages from different sessions in one turn".to_string(), + )); + } + storage = Some(snapshot_storage); + session_id = Some(snapshot_session_id); + messages.push(message); + final_meta = Some(meta); + } + + let (Some(storage), Some(session_id), Some(final_meta)) = (storage, session_id, final_meta) + else { + return Ok(()); + }; + storage + .persist_message_batch_with_retry(&session_id, &messages, &final_meta) + .await +} + +pub(super) async fn append_persisted_messages( + session: &Arc>, + messages: Vec, +) -> Result<(), StorageError> { + if messages.is_empty() { + return Ok(()); + } + + let persistence_lock = { session.lock().await.persistence_lock.clone() }; + let _persistence_guard = persistence_lock.lock().await; + let message_ids: Vec<_> = messages.iter().map(|message| message.id.clone()).collect(); + let snapshots = { + let mut guard = session.lock().await; + messages + .into_iter() + .map(|message| guard.add_message_in_memory(message, true)) + .collect() + }; + + if let Err(error) = persist_added_messages(snapshots).await { + session.lock().await.rollback_message_suffix(&message_ids); + return Err(error); + } + Ok(()) +} diff --git a/src/session/session.rs b/src/session/session.rs index 447a6c4..3fb5c13 100644 --- a/src/session/session.rs +++ b/src/session/session.rs @@ -3,12 +3,13 @@ use std::sync::Arc; use tokio::sync::{Mutex, mpsc, oneshot}; +use super::persistence::append_persisted_messages; use crate::bus::{ChatMessage, MediaItem, MediaRef, MessageSource, OutboundMessage, SourceKind}; use crate::mcp::get_mcp_status; use crate::storage::{Storage, StorageError}; use std::sync::Arc as StdArc; -type MessagePersistSnapshot = ( +pub(super) type MessagePersistSnapshot = ( StdArc, String, crate::storage::message::MessageMeta, @@ -18,7 +19,7 @@ type MessagePersistSnapshot = ( const SESSION_QUEUE_CAPACITY: usize = 32; tokio::task_local! { - static CURRENT_SOURCE_SESSION: Option; + pub(super) static CURRENT_SOURCE_SESSION: Option; } /// Result of handling a message - either an AI response or a command output @@ -100,7 +101,7 @@ pub struct Session { state_version: u64, /// Serializes durable mutations while allowing the session state mutex to /// be released during SQLite I/O. - persistence_lock: Arc>, + pub(super) persistence_lock: Arc>, } /// A task to be processed by the per-session agent worker @@ -359,7 +360,7 @@ impl Session { self.id.to_string() } - fn add_message_in_memory( + pub(super) fn add_message_in_memory( &mut self, message: ChatMessage, persist: bool, @@ -439,7 +440,7 @@ impl Session { /// Roll back messages that were appended in memory but whose atomic /// persistence failed. This is only called while holding the session lock, /// so the suffix check also protects against removing unrelated messages. - fn rollback_message_suffix(&mut self, message_ids: &[String]) { + pub(super) fn rollback_message_suffix(&mut self, message_ids: &[String]) { if message_ids.is_empty() || self.messages.len() < message_ids.len() { return; } @@ -886,7 +887,7 @@ pub struct SessionManager { tools: Arc, skills_loader: Arc, storage: Arc, - bus: Arc, + pub(super) bus: Arc, memory_manager: Arc, sub_agent_manager: Arc, task_supervisor: crate::task_supervisor::TaskSupervisor, @@ -1747,7 +1748,7 @@ impl SessionManager { } } - async fn resolve_dialog_id( + pub(super) async fn resolve_dialog_id( &self, channel: &str, chat_id: &str, @@ -1798,6 +1799,23 @@ impl SessionManager { } } + pub(super) async fn restore_origin_dialog(&self, origin_id: &str, target: &UnifiedSessionId) { + let Some(origin) = UnifiedSessionId::parse(origin_id) else { + tracing::warn!(origin_id, "Ignoring malformed source session id"); + return; + }; + if origin.channel == target.channel + && origin.chat_id == target.chat_id + && origin.dialog_id != target.dialog_id + { + self.inner + .lock() + .await + .current_sessions + .insert(target.chat_scope(), origin_id.to_string()); + } + } + /// Send a system notification (no LLM triggered). /// /// Flow: @@ -2004,64 +2022,6 @@ async fn maybe_generate_title_outside_lock(session: Arc>) -> Resu Ok(()) } -async fn persist_added_messages( - snapshots: Vec>, -) -> Result<(), StorageError> { - let mut storage = None; - let mut session_id = None; - let mut messages = Vec::new(); - let mut final_meta = None; - - for snapshot in snapshots.into_iter().flatten() { - let (snapshot_storage, snapshot_session_id, message, meta) = snapshot; - if let Some(ref expected) = session_id - && expected != &snapshot_session_id - { - return Err(StorageError::Serialization( - "attempted to persist messages from different sessions in one turn".to_string(), - )); - } - storage = Some(snapshot_storage); - session_id = Some(snapshot_session_id); - messages.push(message); - final_meta = Some(meta); - } - - let (Some(storage), Some(session_id), Some(final_meta)) = (storage, session_id, final_meta) - else { - return Ok(()); - }; - storage - .persist_message_batch_with_retry(&session_id, &messages, &final_meta) - .await -} - -async fn append_persisted_messages( - session: &Arc>, - messages: Vec, -) -> Result<(), StorageError> { - if messages.is_empty() { - return Ok(()); - } - - let persistence_lock = { session.lock().await.persistence_lock.clone() }; - let _persistence_guard = persistence_lock.lock().await; - let message_ids: Vec<_> = messages.iter().map(|message| message.id.clone()).collect(); - let snapshots = { - let mut guard = session.lock().await; - messages - .into_iter() - .map(|message| guard.add_message_in_memory(message, true)) - .collect() - }; - - if let Err(error) = persist_added_messages(snapshots).await { - session.lock().await.rollback_message_suffix(&message_ids); - return Err(error); - } - Ok(()) -} - fn spawn_agent_worker( mut task_rx: mpsc::Receiver, session: Arc>, @@ -2548,111 +2508,6 @@ impl SessionManager { } } -#[async_trait::async_trait] -impl OutboundMessenger for SessionManager { - async fn send_message( - &self, - channel: &str, - chat_id: &str, - dialog_id: Option<&str>, - content: &str, - mut source: MessageSource, - media: Vec, - ) -> Result<(), String> { - // Fill origin from current source session if not provided - if source.from_session.is_none() { - source.from_session = CURRENT_SOURCE_SESSION - .try_with(|v| v.clone()) - .ok() - .flatten(); - } - - let (target_sid, session) = if let Some(did) = dialog_id { - let sid = UnifiedSessionId::new(channel, chat_id, did); - let session = self - .get_or_activate_session(&sid) - .await - .map_err(|e| e.to_string())?; - (sid, session) - } else { - let sid = self - .resolve_dialog_id(channel, chat_id) - .await - .map_err(|e| e.to_string())?; - let session = self - .get_or_create_session(&sid) - .await - .map_err(|e| e.to_string())?; - (sid, session) - }; - - // Build message prefix: [message from ] - // Skip prefix for pure file messages to the same session (no cross-session redirect). - let origin = source.from_session.as_deref().unwrap_or("unknown"); - let origin_id = source.from_session.clone(); - let same_session = source - .from_session - .as_deref() - .map(|src| src == target_sid.to_string().as_str()) - .unwrap_or(false); - let marked_content = if content.trim().is_empty() && !media.is_empty() && same_session { - String::new() - } else { - format!("[message from {}] \n{}", origin, content) - }; - - // Write the same text and media delivered to the target into history. - let msg = outbound_history_message(marked_content.clone(), source, &media); - append_persisted_messages(&session, vec![msg]) - .await - .map_err(|e| e.to_string())?; - - // Restore active dialog if source and target share channel:chat_id but differ in dialog_id - if let Some(ref origin_id) = origin_id { - let parts: Vec<&str> = origin_id.split(':').collect(); - if parts.len() == 3 - && parts[0] == channel - && parts[1] == chat_id - && parts[2] != target_sid.dialog_id - { - let scope = format!("{}:{}", channel, chat_id); - self.inner - .lock() - .await - .current_sessions - .insert(scope, origin_id.clone()); - } - } - - // Publish OutboundMessage via bus to target channel - let outbound = OutboundMessage { - channel: channel.to_string(), - chat_id: chat_id.to_string(), - content: marked_content, - reply_to: None, - media, - metadata: HashMap::new(), - delivery: None, - }; - self.bus - .deliver_outbound(outbound) - .await - .map_err(|e| e.to_string())?; - - Ok(()) - } -} - -fn outbound_history_message( - content: impl Into, - source: MessageSource, - media: &[MediaItem], -) -> ChatMessage { - let mut message = ChatMessage::assistant_with_source(content, source); - message.media_refs = media.iter().map(MediaItem::to_media_ref).collect(); - message -} - fn format_task_notification( task_id: &str, status: &crate::agent::TaskStatus, @@ -2670,27 +2525,3 @@ fn format_task_notification( crate::agent::TaskStatus::TimedOut => format!("📋 后台任务超时\n\n任务 ID: {}", task_id), } } - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn outbound_history_preserves_delivered_media() { - let source = MessageSource { - kind: SourceKind::CrossChannel, - from_channel: Some("cli_chat".into()), - from_session: Some("cli_chat:source:dialog".into()), - from_user_id: None, - system_name: None, - task_id: None, - }; - let media = vec![MediaItem::new("/tmp/report.pdf", "file")]; - - let message = outbound_history_message("report", source, &media); - - assert_eq!(message.media_refs.len(), 1); - assert_eq!(message.media_refs[0].path, "/tmp/report.pdf"); - assert_eq!(message.media_refs[0].media_type, "file"); - } -} diff --git a/src/session/session_id.rs b/src/session/session_id.rs index 8a0f6eb..6d464f1 100644 --- a/src/session/session_id.rs +++ b/src/session/session_id.rs @@ -5,13 +5,8 @@ /// Examples: /// - CLI: `"cli_chat:sid_abc123:dialog_xyz"` /// - Feishu: `"feishu:oc_123456:dialog_xyz"` -/// -/// For simple cases where only one dialog exists per chat: -/// - `dialog_id` defaults to `"default"` use serde::{Deserialize, Serialize}; -pub const DEFAULT_DIALOG_ID: &str = "default"; - #[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)] pub struct UnifiedSessionId { pub channel: String, @@ -33,15 +28,6 @@ impl UnifiedSessionId { } } - /// Create with default dialog_id ("default") - pub fn with_default_dialog(channel: impl Into, chat_id: impl Into) -> Self { - Self { - channel: channel.into(), - chat_id: chat_id.into(), - dialog_id: DEFAULT_DIALOG_ID.to_string(), - } - } - /// Parse from string format "channel:chat_id:dialog_id" pub fn parse(s: &str) -> Option { let parts: Vec<&str> = s.split(':').collect(); @@ -82,14 +68,6 @@ mod tests { assert_eq!(id.dialog_id, "dialog456"); } - #[test] - fn test_with_default_dialog() { - let id = UnifiedSessionId::with_default_dialog("feishu", "oc123"); - assert_eq!(id.channel, "feishu"); - assert_eq!(id.chat_id, "oc123"); - assert_eq!(id.dialog_id, "default"); - } - #[test] fn test_parse() { let id = UnifiedSessionId::parse("cli_chat:sid123:dialog456").unwrap();