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

fix(messaging): make persistence and delivery explicit

Browse Source
This commit is contained in:
xiaoxixi 2026-07-14 11:49:02 +08:00
parent c2d4fc5f09
commit 954bfd1d75
6 changed files with 259 additions and 145 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

81
src/bus/dispatcher.rs
View File

@ -56,10 +56,14 @@ impl OutboundDispatcher {
if sender.as_ref().is_none_or(mpsc::Sender::is_closed) {
let Some(channel) = self.channel_manager.get_channel(&msg.channel).await else {
tracing::warn!(channel = %msg.channel, "No channel found for message");
msg.complete_delivery(Err(format!("channel not found: {}", msg.channel)));
continue;
};
let (new_sender, receiver) = mpsc::channel(LANE_CAPACITY);
self.spawn_lane(channel, receiver, msg.channel.clone(), msg.chat_id.clone());
if !self.spawn_lane(channel, receiver, msg.channel.clone(), msg.chat_id.clone()) {
msg.complete_delivery(Err("dispatcher is shutting down".to_string()));
continue;
}
lanes.insert(lane_key.clone(), new_sender.clone());
sender = Some(new_sender);
}
@ -77,6 +81,7 @@ impl OutboundDispatcher {
capacity = LANE_CAPACITY,
"Outbound lane full; rejecting message instead of blocking other destinations"
);
msg.complete_delivery(Err("outbound lane is full".to_string()));
}
Err(mpsc::error::TrySendError::Closed(msg)) => {
// The lane may have expired between the closed check and
@ -84,12 +89,24 @@ impl OutboundDispatcher {
lanes.remove(&lane_key);
let Some(channel) = self.channel_manager.get_channel(&msg.channel).await else {
tracing::warn!(channel = %msg.channel, "No channel found for message");
msg.complete_delivery(Err(format!("channel not found: {}", msg.channel)));
continue;
};
let (new_sender, receiver) = mpsc::channel(LANE_CAPACITY);
self.spawn_lane(channel, receiver, msg.channel.clone(), msg.chat_id.clone());
if new_sender.try_send(msg).is_ok() {
lanes.insert(lane_key, new_sender);
if !self.spawn_lane(channel, receiver, msg.channel.clone(), msg.chat_id.clone())
{
msg.complete_delivery(Err("dispatcher is shutting down".to_string()));
continue;
}
match new_sender.try_send(msg) {
Ok(()) => {
lanes.insert(lane_key, new_sender);
}
Err(error) => {
error.into_inner().complete_delivery(Err(
"outbound lane could not be restarted during shutdown".to_string(),
));
}
}
}
}
@ -102,7 +119,7 @@ impl OutboundDispatcher {
mut receiver: mpsc::Receiver<OutboundMessage>,
channel_name: String,
chat_id: String,
) {
) -> bool {
self.task_supervisor.spawn(
format!("outbound-lane:{channel_name}:{chat_id}"),
async move {
@ -111,7 +128,8 @@ impl OutboundDispatcher {
Ok(Some(msg)) => msg,
Ok(None) | Err(_) => break,
};
if let Err(error) = Self::send_with_retry(&*channel, msg).await {
let result = Self::send_with_retry(&*channel, &msg).await;
if let Err(error) = &result {
tracing::error!(
channel = %channel_name,
chat_id = %chat_id,
@ -119,14 +137,15 @@ impl OutboundDispatcher {
"Failed to send message after retries"
);
}
msg.complete_delivery(result.map_err(|error| error.to_string()));
}
},
);
)
}
async fn send_with_retry(
channel: &dyn Channel,
msg: OutboundMessage,
msg: &OutboundMessage,
) -> Result<(), ChannelError> {
const DELAYS: &[u64] = &[1, 2, 4];
@ -201,6 +220,7 @@ mod tests {
reply_to: None,
media: vec![],
metadata: HashMap::new(),
delivery: None,
}
}
@ -247,4 +267,49 @@ mod tests {
task.abort();
supervisor.shutdown(Duration::from_secs(1)).await;
}
#[tokio::test]
async fn confirmed_delivery_reports_missing_channel() {
let bus = MessageBus::new(8);
let manager = ChannelManager::with_bus(
Arc::new(crate::channels::CliChatChannel::new()),
bus.clone(),
);
let supervisor = TaskSupervisor::new();
let dispatcher = OutboundDispatcher::new(bus.clone(), manager, supervisor.clone());
let task = tokio::spawn(async move { dispatcher.run().await });
let mut message = outbound("missing", "not delivered");
message.channel = "missing".to_string();
let error = bus.deliver_outbound(message).await.unwrap_err();
assert!(matches!(error, crate::bus::BusError::DeliveryFailed(_)));
task.abort();
supervisor.shutdown(Duration::from_secs(1)).await;
}
#[tokio::test]
async fn confirmed_delivery_waits_for_channel_send() {
let bus = MessageBus::new(8);
let manager = ChannelManager::with_bus(
Arc::new(crate::channels::CliChatChannel::new()),
bus.clone(),
);
let channel = Arc::new(RecordingChannel {
sent: Mutex::new(Vec::new()),
notify: Notify::new(),
});
manager.register_channel("recording", channel.clone()).await;
let supervisor = TaskSupervisor::new();
let dispatcher = OutboundDispatcher::new(bus.clone(), manager, supervisor.clone());
let task = tokio::spawn(async move { dispatcher.run().await });
bus.deliver_outbound(outbound("confirmed", "delivered"))
.await
.unwrap();
assert_eq!(channel.sent.lock().await.as_slice(), &["delivered"]);
task.abort();
supervisor.shutdown(Duration::from_secs(1)).await;
}
}

9
src/bus/message.rs
View File

@ -276,6 +276,15 @@ pub struct OutboundMessage {
pub reply_to: Option<String>,
pub media: Vec<MediaItem>,
pub metadata: HashMap<String, String>,
pub(crate) delivery: Option<tokio::sync::watch::Sender<Option<Result<(), String>>>>,
}
impl OutboundMessage {
pub(crate) fn complete_delivery(&self, result: Result<(), String>) {
if let Some(delivery) = &self.delivery {
delivery.send_replace(Some(result));
}
}
}
// ============================================================================

23
src/bus/mod.rs
View File

@ -68,6 +68,25 @@ impl MessageBus {
.map_err(|_| BusError::Closed)
}
/// Publish an outbound message and wait for the dispatcher to report the
/// actual channel delivery result.
pub async fn deliver_outbound(&self, mut msg: OutboundMessage) -> Result<(), BusError> {
let (delivery_tx, mut delivery_rx) = tokio::sync::watch::channel(None);
msg.delivery = Some(delivery_tx);
self.publish_outbound(msg).await?;
tokio::time::timeout(std::time::Duration::from_secs(120), async {
loop {
delivery_rx.changed().await.map_err(|_| BusError::Closed)?;
if let Some(result) = delivery_rx.borrow().clone() {
return result.map_err(BusError::DeliveryFailed);
}
}
})
.await
.map_err(|_| BusError::DeliveryTimedOut)?
}
/// Consume an outbound message (Dispatcher -> Bus)
pub async fn consume_outbound(&self) -> Option<OutboundMessage> {
self.outbound_rx.lock().await.recv().await
@ -95,12 +114,16 @@ impl MessageBus {
#[derive(Debug)]
pub enum BusError {
Closed,
DeliveryFailed(String),
DeliveryTimedOut,
}
impl std::fmt::Display for BusError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
BusError::Closed => write!(f, "Bus channel closed"),
BusError::DeliveryFailed(error) => write!(f, "Outbound delivery failed: {error}"),
BusError::DeliveryTimedOut => write!(f, "Outbound delivery confirmation timed out"),
}
}
}

1
src/channels/cli_chat.rs
View File

@ -650,6 +650,7 @@ mod tests {
reply_to: None,
media: Vec::new(),
metadata: Default::default(),
delivery: None,
})
.await
.unwrap();

2
src/gateway/mod.rs
View File

@ -226,6 +226,7 @@ impl GatewayState {
reply_to: None,
media: vec![],
metadata: inbound.forwarded_metadata,
delivery: None,
};
if let Err(e) = bus.publish_outbound(outbound).await {
tracing::error!(error = %e, "Failed to publish outbound");
@ -239,6 +240,7 @@ impl GatewayState {
reply_to: None,
media: vec![],
metadata: inbound.forwarded_metadata,
delivery: None,
};
if let Err(e) = bus.publish_outbound(outbound).await {
tracing::error!(error = %e, "Failed to publish outbound");

288
src/session/session.rs
View File

@ -98,6 +98,9 @@ pub struct Session {
/// not overwrite a session that was changed by a command such as /clear or
/// /delete while the slow work was in flight.
state_version: u64,
/// Serializes durable mutations while allowing the session state mutex to
/// be released during SQLite I/O.
persistence_lock: Arc<Mutex<()>>,
}
/// A task to be processed by the per-session agent worker
@ -171,6 +174,7 @@ impl Session {
current_cancel: None,
worker_generation: 0,
state_version: 0,
persistence_lock: Arc::new(Mutex::new(())),
})
}
@ -346,6 +350,7 @@ impl Session {
current_cancel: None,
worker_generation: 0,
state_version: 0,
persistence_lock: Arc::new(Mutex::new(())),
})
}
@ -354,22 +359,6 @@ impl Session {
self.id.to_string()
}
/// 添加消息到历史并持久化到 Storage
/// 如果 `persist` 为 false，只更新内存（用于 compaction 场景）
pub async fn add_message(
&mut self,
message: ChatMessage,
persist: bool,
) -> Result<(), StorageError> {
let message_id = message.id.clone();
let snapshot = self.add_message_in_memory(message, persist);
if let Err(error) = persist_added_message(snapshot).await {
self.rollback_message_suffix(std::slice::from_ref(&message_id));
return Err(error);
}
Ok(())
}
fn add_message_in_memory(
&mut self,
message: ChatMessage,
@ -480,30 +469,6 @@ impl Session {
&self.messages
}
/// 清除历史消息
pub fn clear_history(&mut self) {
let len = self.messages.len();
self.messages.clear();
self.seq_counter = 1;
self.total_message_count = 0;
self.message_count = 0;
self.state_version = self.state_version.wrapping_add(1);
#[cfg(debug_assertions)]
tracing::debug!(session_id = %self.id, previous_len = len, "Chat history cleared");
}
/// 重置对话上下文
pub fn reset_context(&mut self) {
let len = self.messages.len();
self.messages.clear();
self.seq_counter = 1;
self.total_message_count = 0;
self.message_count = 0;
self.state_version = self.state_version.wrapping_add(1);
#[cfg(debug_assertions)]
tracing::debug!(session_id = %self.id, previous_len = len, "Chat context reset in memory");
}
pub fn create_user_message(&self, content: &str, media_refs: Vec<MediaRef>) -> ChatMessage {
if media_refs.is_empty() {
ChatMessage::user(content)
@ -535,14 +500,6 @@ impl Session {
message
}
/// 将 session 元数据写回 Storage
pub async fn persist_session_meta(&self) -> Result<(), StorageError> {
if let Some((storage, meta)) = self.session_meta_snapshot() {
storage.upsert_session(&meta).await?;
}
Ok(())
}
fn session_meta_snapshot(
&self,
) -> Option<(StdArc<Storage>, crate::storage::session::SessionMeta)> {
@ -1079,6 +1036,7 @@ impl SessionManager {
reply_to: None,
media: vec![],
metadata: std::collections::HashMap::new(),
delivery: None,
};
let _ = sm_bus.publish_outbound(outbound).await;
}
@ -1701,17 +1659,28 @@ impl SessionManager {
) -> Result<(), AgentError> {
// Update in-memory session
let session = self.get_or_create_session(session_id).await?;
let mut session_guard = session.lock().await;
session_guard.title = title.to_string();
session_guard
.persist_session_meta()
.await
.map_err(|e| AgentError::Other(format!("failed to rename dialog: {}", e)))?;
let persistence_lock = { session.lock().await.persistence_lock.clone() };
let _persistence_guard = persistence_lock.lock().await;
let meta_snapshot = {
let mut session_guard = session.lock().await;
session_guard.title = title.to_string();
session_guard.state_version = session_guard.state_version.wrapping_add(1);
session_guard.session_meta_snapshot()
};
if let Some((storage, meta)) = meta_snapshot {
storage
.upsert_session(&meta)
.await
.map_err(|e| AgentError::Other(format!("failed to rename dialog: {}", e)))?;
}
Ok(())
}
pub async fn delete_dialog(&self, session_id: &UnifiedSessionId) -> Result<(), AgentError> {
let session_id_str = session_id.to_string();
let session = self.get_or_create_session(session_id).await?;
let persistence_lock = { session.lock().await.persistence_lock.clone() };
let _persistence_guard = persistence_lock.lock().await;
// Soft delete from Storage
self.storage
@ -1730,6 +1699,9 @@ impl SessionManager {
pub async fn archive_dialog(&self, session_id: &UnifiedSessionId) -> Result<(), AgentError> {
let session_id_str = session_id.to_string();
let session = self.get_or_create_session(session_id).await?;
let persistence_lock = { session.lock().await.persistence_lock.clone() };
let _persistence_guard = persistence_lock.lock().await;
self.storage
.archive_session(&session_id_str)
.await
@ -1842,22 +1814,18 @@ impl SessionManager {
) -> Result<(), AgentError> {
let unified_id = self.resolve_dialog_id(channel, chat_id).await?;
let session = self.get_or_create_session(&unified_id).await?;
{
let mut guard = session.lock().await;
let source = MessageSource {
kind: SourceKind::SystemNotification,
from_channel: None,
from_session: None,
from_user_id: None,
system_name: Some(system_name.to_string()),
task_id: task_id.map(|s| s.to_string()),
};
let msg = ChatMessage::assistant_with_source(content, source);
guard
.add_message(msg, true)
.await
.map_err(|e| AgentError::Other(format!("persist error: {}", e)))?;
}
let source = MessageSource {
kind: SourceKind::SystemNotification,
from_channel: None,
from_session: None,
from_user_id: None,
system_name: Some(system_name.to_string()),
task_id: task_id.map(|s| s.to_string()),
};
let msg = ChatMessage::assistant_with_source(content, source);
append_persisted_messages(&session, vec![msg])
.await
.map_err(|e| AgentError::Other(format!("persist error: {}", e)))?;
let outbound = OutboundMessage {
channel: channel.to_string(),
@ -1866,9 +1834,10 @@ impl SessionManager {
reply_to: None,
media: vec![],
metadata: HashMap::new(),
delivery: None,
};
self.bus
.publish_outbound(outbound)
.deliver_outbound(outbound)
.await
.map_err(|e| AgentError::Other(format!("bus publish error: {}", e)))?;
@ -2014,6 +1983,8 @@ async fn maybe_generate_title_outside_lock(session: Arc<Mutex<Session>>) -> Resu
.map_err(|e| AgentError::Other(format!("LLM call failed: {}", e)))?;
let title = response.content.trim().to_string();
let persistence_lock = { session.lock().await.persistence_lock.clone() };
let _persistence_guard = persistence_lock.lock().await;
let meta_snapshot = {
let mut guard = session.lock().await;
if guard.apply_generated_title(title) {
@ -2033,22 +2004,6 @@ async fn maybe_generate_title_outside_lock(session: Arc<Mutex<Session>>) -> Resu
Ok(())
}
async fn persist_added_message(
snapshot: Option<MessagePersistSnapshot>,
) -> Result<(), StorageError> {
let Some((storage, session_id, msg_meta, session_meta)) = snapshot else {
return Ok(());
};
storage
.persist_message_batch_with_retry(
&session_id,
std::slice::from_ref(&msg_meta),
&session_meta,
)
.await
}
async fn persist_added_messages(
snapshots: Vec<Option<MessagePersistSnapshot>>,
) -> Result<(), StorageError> {
@ -2081,6 +2036,32 @@ async fn persist_added_messages(
.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>>,
@ -2122,6 +2103,7 @@ fn spawn_agent_worker(
reply_to: None,
media: vec![],
metadata,
delivery: None,
};
let _ = bus.publish_outbound(outbound).await;
}
@ -2134,6 +2116,30 @@ fn spawn_agent_worker(
// /stop and other commands are not blocked behind slow I/O or
// LLM-backed compaction.
let skills_prompt = skills_loader.build_skills_prompt();
let user_message = {
let guard = session.lock().await;
if guard.worker_generation != worker_gen {
return;
}
let media_refs: Vec<MediaRef> =
task.media.iter().map(MediaItem::to_media_ref).collect();
guard.create_user_message(&task.content, media_refs)
};
if let Err(e) = append_persisted_messages(&session, vec![user_message]).await {
tracing::error!(error = %e, "Failed to persist user message");
let err_outbound = OutboundMessage {
channel: task_chan.clone(),
chat_id: task_cid.clone(),
content: "Failed to save your message, please try again.".to_string(),
reply_to: None,
media: vec![],
metadata: HashMap::new(),
delivery: None,
};
let _ = bus.publish_outbound(err_outbound).await;
continue 'tasks;
}
let (agent, history_raw, mut compressor, base_version, cancel_rx) = {
let mut guard = session.lock().await;
@ -2141,31 +2147,6 @@ fn spawn_agent_worker(
return; // stale worker
}
let media_refs: Vec<MediaRef> =
task.media.iter().map(|m| m.to_media_ref()).collect();
let user_message = guard.create_user_message(&task.content, media_refs);
let user_message_id = user_message.id.clone();
let user_persist = guard.add_message_in_memory(user_message, true);
if let Err(e) = persist_added_message(user_persist).await {
guard.rollback_message_suffix(std::slice::from_ref(&user_message_id));
drop(guard);
tracing::error!(error = %e, "Failed to persist user message");
let err_outbound = OutboundMessage {
channel: task_chan.clone(),
chat_id: task_cid.clone(),
content: "Failed to save your message, please try again."
.to_string(),
reply_to: None,
media: vec![],
metadata: HashMap::new(),
};
let _ = bus.publish_outbound(err_outbound).await;
continue 'tasks;
}
if guard.worker_generation != worker_gen {
return;
}
let history_raw = guard.get_history().to_vec();
let agent = match guard.create_agent_with_notify(notify_tx) {
@ -2180,6 +2161,7 @@ fn spawn_agent_worker(
reply_to: None,
media: vec![],
metadata: HashMap::new(),
delivery: None,
};
let _ = bus.publish_outbound(err_outbound).await;
continue 'tasks;
@ -2332,6 +2314,7 @@ fn spawn_agent_worker(
reply_to: None,
media: vec![],
metadata: HashMap::new(),
delivery: None,
};
let _ = bus2.publish_outbound(err_outbound).await;
return;
@ -2390,6 +2373,7 @@ fn spawn_agent_worker(
reply_to: None,
media: vec![],
metadata: HashMap::new(),
delivery: None,
};
let _ = bus2.publish_outbound(err_outbound).await;
return;
@ -2405,30 +2389,25 @@ fn spawn_agent_worker(
reply_to: None,
media: vec![],
metadata: HashMap::new(),
delivery: None,
};
let _ = bus2.publish_outbound(err_outbound).await;
return;
}
};
let response = {
let mut guard = session2.lock().await;
let mut persist_snapshots = Vec::new();
let mut message_ids = Vec::new();
for msg in result.emitted_messages {
message_ids.push(msg.id.clone());
persist_snapshots.push(guard.add_message_in_memory(msg, true));
}
let sent_count = guard.messages.len();
guard.compressor.set_last_api_info(sent_count, result.total_tokens);
if let Err(e) = persist_added_messages(persist_snapshots).await {
guard.rollback_message_suffix(&message_ids);
let response_content = result.final_response.content;
let total_tokens = result.total_tokens;
let response =
if let Err(e) = append_persisted_messages(&session2, result.emitted_messages).await {
tracing::error!(error = %e, "Failed to atomically persist agent turn");
None
} else {
Some(result.final_response.content)
}
};
let mut guard = session2.lock().await;
let sent_count = guard.messages.len();
guard.compressor.set_last_api_info(sent_count, total_tokens);
Some(response_content)
};
let Some(response) = response else {
let err_outbound = OutboundMessage {
@ -2439,6 +2418,7 @@ fn spawn_agent_worker(
reply_to: None,
media: vec![],
metadata: HashMap::new(),
delivery: None,
};
let _ = bus2.publish_outbound(err_outbound).await;
return;
@ -2455,6 +2435,7 @@ fn spawn_agent_worker(
reply_to: None,
media: vec![],
metadata: HashMap::new(),
delivery: None,
};
let _ = bus2.publish_outbound(outbound).await;
};
@ -2533,6 +2514,8 @@ impl SessionManager {
unified_id: &UnifiedSessionId,
) -> Result<(), AgentError> {
let session = self.get_or_create_session(unified_id).await?;
let persistence_lock = { session.lock().await.persistence_lock.clone() };
let _persistence_guard = persistence_lock.lock().await;
let (storage, session_id, meta_snapshot) = {
let mut session_guard = session.lock().await;
// Clear in-memory
@ -2618,15 +2601,11 @@ impl OutboundMessenger for SessionManager {
format!("[message from {}] \n{}", origin, content)
};
// Write source-tagged assistant message to target session history
{
let mut guard = session.lock().await;
let msg = ChatMessage::assistant_with_source(marked_content.clone(), source);
guard
.add_message(msg, true)
.await
.map_err(|e| e.to_string())?;
}
// 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 {
@ -2653,9 +2632,10 @@ impl OutboundMessenger for SessionManager {
reply_to: None,
media,
metadata: HashMap::new(),
delivery: None,
};
self.bus
.publish_outbound(outbound)
.deliver_outbound(outbound)
.await
.map_err(|e| e.to_string())?;
@ -2663,6 +2643,16 @@ impl OutboundMessenger for SessionManager {
}
}
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,
@ -2680,3 +2670,27 @@ 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");
}
}