xiaoski/PicoBot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor(session): isolate messaging and persistence

Browse Source
This commit is contained in:
xiaoxixi 2026-07-14 11:56:19 +08:00
parent 24d3e26b43
commit f42e9d44cc
5 changed files with 205 additions and 216 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

113
src/session/messenger.rs Normal file
View File

@ -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<MediaItem>,
) -> 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<String>,
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");
}
}

2
src/session/mod.rs
View File

@ -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;

65
src/session/persistence.rs Normal file
View File

@ -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<Option<MessagePersistSnapshot>>,
) -> 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<Mutex<Session>>,
messages: Vec<ChatMessage>,
) -> 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(())
}

219
src/session/session.rs
View File

@ -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<Storage>,
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<String>;
pub(super) static CURRENT_SOURCE_SESSION: Option<String>;
}
/// 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<Mutex<()>>,
pub(super) persistence_lock: Arc<Mutex<()>>,
}
/// 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<ToolRegistry>,
skills_loader: Arc<SkillsLoader>,
storage: Arc<Storage>,
bus: Arc<MessageBus>,
pub(super) bus: Arc<MessageBus>,
memory_manager: Arc<crate::memory::MemoryManager>,
sub_agent_manager: Arc<crate::agent::SubAgentManager>,
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<Mutex<Session>>) -> Resu
Ok(())
}
async fn persist_added_messages(
snapshots: Vec<Option<MessagePersistSnapshot>>,
) -> 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<Mutex<Session>>,
messages: Vec<ChatMessage>,
) -> 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<AgentTask>,
session: Arc<Mutex<Session>>,
@ -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<MediaItem>,
) -> 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 <origin>]
// 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<String>,
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");
}
}

22
src/session/session_id.rs
View File

@ -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<String>, chat_id: impl Into<String>) -> 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<Self> {
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();