From b06bc4f02501d52dd259c3cabd97ebfee48043b1 Mon Sep 17 00:00:00 2001 From: xiaoxixi Date: Tue, 14 Jul 2026 10:08:55 +0800 Subject: [PATCH] fix(runtime): harden session persistence and outbound dispatch --- AGENTS.md | 2 +- src/bus/dispatcher.rs | 221 +++++++++++++++++++++++++---- src/session/session.rs | 315 +++++++++++++++++++++++++++++------------ src/storage/mod.rs | 139 ++++++++++++++++++ 4 files changed, 560 insertions(+), 117 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 28357fa..3edbf13 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -21,7 +21,7 @@ ## Reference -- `reference/` — third-party reference implementations (nanobot, Mini-Agent, zeroclaw); not part of this project; do not modify +- `reference/` — third-party reference implementations; not part of this project; do not modify ## Architecture diff --git a/src/bus/dispatcher.rs b/src/bus/dispatcher.rs index 82b607c..dbc6ebe 100644 --- a/src/bus/dispatcher.rs +++ b/src/bus/dispatcher.rs @@ -1,11 +1,20 @@ +use std::collections::HashMap; use std::sync::Arc; +use std::time::Duration; + +use tokio::sync::mpsc; use crate::bus::{MessageBus, OutboundMessage}; use crate::channels::ChannelManager; use crate::channels::base::{Channel, ChannelError}; -/// OutboundDispatcher consumes outbound messages from the MessageBus -/// and dispatches them to the appropriate Channel +const LANE_CAPACITY: usize = 64; +const LANE_IDLE_TIMEOUT: Duration = Duration::from_secs(300); +const SEND_TIMEOUT: Duration = Duration::from_secs(30); + +/// Dispatches outbound messages through independent per-conversation lanes. +/// Messages to the same channel/chat remain ordered, while a slow destination +/// cannot block delivery to unrelated conversations. pub struct OutboundDispatcher { bus: Arc, channel_manager: ChannelManager, @@ -19,9 +28,10 @@ impl OutboundDispatcher { } } - /// Run the dispatcher loop - consumes from bus and dispatches to channels pub async fn run(&self) { - tracing::info!("OutboundDispatcher started"); + tracing::info!(lane_capacity = LANE_CAPACITY, "OutboundDispatcher started"); + let mut lanes: HashMap> = HashMap::new(); + let mut messages_seen = 0_u64; loop { let Some(msg) = self.bus.consume_outbound().await else { @@ -29,46 +39,199 @@ impl OutboundDispatcher { break; }; - let channel_name = msg.channel.clone(); - let channel = self.channel_manager.get_channel(&channel_name).await; + messages_seen = messages_seen.wrapping_add(1); + if messages_seen.is_multiple_of(128) { + lanes.retain(|_, sender| !sender.is_closed()); + } - match channel { - Some(ch) => { - if let Err(e) = self.send_with_retry(&*ch, msg).await { - tracing::error!(channel = %channel_name, error = %e, "Failed to send message after retries"); - } + let lane_key = format!("{}\0{}", msg.channel, msg.chat_id); + let mut sender = lanes.get(&lane_key).cloned(); + 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"); + continue; + }; + let (new_sender, receiver) = mpsc::channel(LANE_CAPACITY); + Self::spawn_lane(channel, receiver, msg.channel.clone(), msg.chat_id.clone()); + lanes.insert(lane_key.clone(), new_sender.clone()); + sender = Some(new_sender); + } + + let Some(sender) = sender else { + tracing::error!("Outbound lane creation did not produce a sender"); + continue; + }; + match sender.try_send(msg) { + Ok(()) => {} + Err(mpsc::error::TrySendError::Full(msg)) => { + tracing::error!( + channel = %msg.channel, + chat_id = %msg.chat_id, + capacity = LANE_CAPACITY, + "Outbound lane full; rejecting message instead of blocking other destinations" + ); } - None => { - tracing::warn!(channel = %channel_name, "No channel found for message"); + Err(mpsc::error::TrySendError::Closed(msg)) => { + // The lane may have expired between the closed check and + // enqueue. Recreate it once and preserve this message. + 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"); + 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); + } } } } } - /// Send a message with exponential retry + fn spawn_lane( + channel: Arc, + mut receiver: mpsc::Receiver, + channel_name: String, + chat_id: String, + ) { + tokio::spawn(async move { + loop { + let msg = match tokio::time::timeout(LANE_IDLE_TIMEOUT, receiver.recv()).await { + Ok(Some(msg)) => msg, + Ok(None) | Err(_) => break, + }; + if let Err(error) = Self::send_with_retry(&*channel, msg).await { + tracing::error!( + channel = %channel_name, + chat_id = %chat_id, + error = %error, + "Failed to send message after retries" + ); + } + } + }); + } + async fn send_with_retry( - &self, channel: &dyn Channel, msg: OutboundMessage, ) -> Result<(), ChannelError> { const DELAYS: &[u64] = &[1, 2, 4]; - for (i, &delay) in DELAYS.iter().enumerate() { - match channel.send(msg.clone()).await { - Ok(()) => return Ok(()), - Err(e) if i < DELAYS.len() - 1 => { - tracing::warn!( - attempt = i + 1, - delay = delay, - error = %e, - "Send failed, retrying" - ); - tokio::time::sleep(tokio::time::Duration::from_secs(delay)).await; + for (attempt, &delay) in DELAYS.iter().enumerate() { + let result = tokio::time::timeout(SEND_TIMEOUT, channel.send(msg.clone())).await; + match result { + Ok(Ok(())) => return Ok(()), + Ok(Err(error)) if attempt < DELAYS.len() - 1 => { + tracing::warn!(attempt = attempt + 1, delay, error = %error, "Send failed, retrying"); + } + Ok(Err(error)) => return Err(error), + Err(_) if attempt < DELAYS.len() - 1 => { + tracing::warn!(attempt = attempt + 1, delay, "Send timed out, retrying"); + } + Err(_) => { + return Err(ChannelError::Other(format!( + "send timed out after {} seconds", + SEND_TIMEOUT.as_secs() + ))); } - Err(e) => return Err(e), } + tokio::time::sleep(Duration::from_secs(delay)).await; } - // All retries exhausted - should not reach here as last iteration returns - Ok(()) + unreachable!() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use async_trait::async_trait; + use tokio::sync::{Mutex, Notify}; + + struct RecordingChannel { + sent: Mutex>, + notify: Notify, + } + + #[async_trait] + impl Channel for RecordingChannel { + fn name(&self) -> &str { + "recording" + } + + fn is_running(&self) -> bool { + true + } + + async fn start(&self, _bus: Arc) -> Result<(), ChannelError> { + Ok(()) + } + + async fn stop(&self) -> Result<(), ChannelError> { + Ok(()) + } + + async fn send(&self, msg: OutboundMessage) -> Result<(), ChannelError> { + if msg.chat_id == "slow" { + tokio::time::sleep(Duration::from_millis(50)).await; + } + self.sent.lock().await.push(msg.content); + self.notify.notify_waiters(); + Ok(()) + } + } + + fn outbound(chat_id: &str, content: &str) -> OutboundMessage { + OutboundMessage { + channel: "recording".to_string(), + chat_id: chat_id.to_string(), + content: content.to_string(), + reply_to: None, + media: vec![], + metadata: HashMap::new(), + } + } + + #[tokio::test] + async fn slow_conversation_does_not_block_other_conversations() { + 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 dispatcher = OutboundDispatcher::new(bus.clone(), manager); + let task = tokio::spawn(async move { dispatcher.run().await }); + bus.publish_outbound(outbound("slow", "slow-1")) + .await + .unwrap(); + bus.publish_outbound(outbound("slow", "slow-2")) + .await + .unwrap(); + bus.publish_outbound(outbound("fast", "fast-1")) + .await + .unwrap(); + + tokio::time::timeout(Duration::from_secs(1), async { + loop { + if channel.sent.lock().await.len() == 3 { + break; + } + tokio::time::sleep(Duration::from_millis(5)).await; + } + }) + .await + .unwrap(); + + let sent = channel.sent.lock().await.clone(); + assert_eq!(sent[0], "fast-1"); + assert_eq!(&sent[1..], &["slow-1", "slow-2"]); + task.abort(); } } diff --git a/src/session/session.rs b/src/session/session.rs index 6f387c6..7e8b870 100644 --- a/src/session/session.rs +++ b/src/session/session.rs @@ -15,6 +15,8 @@ type MessagePersistSnapshot = ( crate::storage::session::SessionMeta, ); +const SESSION_QUEUE_CAPACITY: usize = 32; + tokio::task_local! { static CURRENT_SOURCE_SESSION: Option; } @@ -84,7 +86,7 @@ pub struct Session { memory_manager: Arc, /// Task queue for per-session serial agent processing - agent_tx: Option>, + agent_tx: Option>, /// Cancel signal for the currently executing agent task current_cancel: Option>, /// Monotonic counter to detect stale workers @@ -203,10 +205,9 @@ impl Session { let timelines = storage .load_session_timelines(&id.to_string(), 4) .await - .unwrap_or_else(|e| { - tracing::warn!(error = %e, "Failed to load session timelines"); - Vec::new() - }); + .map_err(|e| { + AgentError::Other(format!("failed to load session timelines: {}", e)) + })?; let has_more_timelines = timelines.len() > 3; @@ -229,10 +230,9 @@ impl Session { let tail = storage .load_messages_after_timestamp(&id.to_string(), after_ts) .await - .unwrap_or_else(|e| { - tracing::warn!(error = %e, "Failed to load messages after timestamp"); - Vec::new() - }); + .map_err(|e| { + AgentError::Other(format!("failed to load messages after timestamp: {}", e)) + })?; let mut tail_msgs: Vec = tail .into_iter() @@ -312,7 +312,7 @@ impl Session { let max_seq = storage .get_max_message_seq(&id.to_string()) .await - .unwrap_or(0); + .map_err(|e| AgentError::Other(format!("failed to load message sequence: {}", e)))?; let seq_counter = max_seq + 1; let total_message_count = session_meta.message_count; @@ -354,8 +354,13 @@ impl Session { message: ChatMessage, persist: bool, ) -> Result<(), StorageError> { + let message_id = message.id.clone(); let snapshot = self.add_message_in_memory(message, persist); - persist_added_message(snapshot).await + 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( @@ -435,6 +440,34 @@ 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]) { + if message_ids.is_empty() || self.messages.len() < message_ids.len() { + return; + } + let start = self.messages.len() - message_ids.len(); + if self.messages[start..] + .iter() + .zip(message_ids) + .any(|(message, id)| &message.id != id) + { + tracing::error!(session_id = %self.id, "Refusing to roll back a non-matching message suffix"); + return; + } + + let removed_user_messages = self.messages[start..] + .iter() + .filter(|message| message.role == "user") + .count() as i64; + self.messages.truncate(start); + self.seq_counter -= message_ids.len() as i64; + self.total_message_count -= message_ids.len() as i64; + self.message_count -= removed_user_messages; + self.state_version = self.state_version.wrapping_add(1); + } + /// 获取消息历史 pub fn get_history(&self) -> &[ChatMessage] { &self.messages @@ -1481,52 +1514,78 @@ impl SessionManager { unified_id: &UnifiedSessionId, ) -> Result>, AgentError> { let session_id_str = unified_id.to_string(); - let inner = &mut *self.inner.lock().await; - - if let Some(session) = inner.sessions.get(&session_id_str) { - return Ok(session.clone()); + if let Some(session) = self + .inner + .lock() + .await + .sessions + .get(&session_id_str) + .cloned() + { + return Ok(session); } - // Try to restore from Storage - match self.storage.get_session(&session_id_str).await { + // Perform storage/provider I/O without holding the global registry lock. + let session = match self.storage.get_session(&session_id_str).await { Ok(meta) => { tracing::debug!(session_id = %session_id_str, last_active_at = %meta.last_active_at, message_count = %meta.message_count, "Restoring session from Storage"); - let session = Session::from_storage( + Session::from_storage( unified_id.clone(), self.provider_config.clone(), self.tools.clone(), self.storage.clone(), self.memory_manager.clone(), ) - .await?; - - let arc = Arc::new(Mutex::new(session)); - inner.sessions.insert(session_id_str.clone(), arc.clone()); - // Set as current session - let chat_scope = format!("{}:{}", unified_id.channel, unified_id.chat_id); - inner.current_sessions.insert(chat_scope, session_id_str); - return Ok(arc); + .await? } - Err(_) => { - // Session not in Storage, create new + Err(StorageError::NotFound(_)) => { + let now = chrono::Utc::now().timestamp_millis(); + let meta = crate::storage::session::SessionMeta { + id: session_id_str.clone(), + channel: unified_id.channel.clone(), + chat_id: unified_id.chat_id.clone(), + dialog_id: unified_id.dialog_id.clone(), + title: "新对话".to_string(), + created_at: now, + last_active_at: now, + message_count: 0, + routing_info: None, + archived_at: None, + deleted_at: None, + last_consolidated_at: None, + last_compressed_message_at: None, + }; + self.storage.upsert_session(&meta).await.map_err(|e| { + AgentError::Other(format!("failed to create session in storage: {}", e)) + })?; + Session::new( + unified_id.clone(), + self.provider_config.clone(), + self.tools.clone(), + Some(self.storage.clone()), + String::new(), + "新对话".to_string(), + self.memory_manager.clone(), + ) + .await? } - } - - // Create new session - let session = Session::new( - unified_id.clone(), - self.provider_config.clone(), - self.tools.clone(), - Some(self.storage.clone()), - String::new(), - "新对话".to_string(), - self.memory_manager.clone(), - ) - .await?; + Err(e) => { + return Err(AgentError::Other(format!( + "failed to look up session in storage: {}", + e + ))); + } + }; let arc = Arc::new(Mutex::new(session)); + // Another caller may have completed the same load while I/O was in + // progress. Keep the already-published instance as the single source + // of truth. + let inner = &mut *self.inner.lock().await; + if let Some(existing) = inner.sessions.get(&session_id_str) { + return Ok(existing.clone()); + } inner.sessions.insert(session_id_str.clone(), arc.clone()); - // Set as current session let chat_scope = format!("{}:{}", unified_id.channel, unified_id.chat_id); inner.current_sessions.insert(chat_scope, session_id_str); Ok(arc) @@ -1713,11 +1772,20 @@ impl SessionManager { .cloned() }; - if let Some(ref current_id) = current_id - && let Ok(_) = self.storage.get_session(current_id).await - { - if let Some(parsed) = UnifiedSessionId::parse(current_id) { - return Ok(parsed); + if let Some(ref current_id) = current_id { + match self.storage.get_session(current_id).await { + Ok(_) => { + if let Some(parsed) = UnifiedSessionId::parse(current_id) { + return Ok(parsed); + } + } + Err(StorageError::NotFound(_)) => {} + Err(e) => { + return Err(AgentError::Other(format!( + "failed to resolve current session: {}", + e + ))); + } } } @@ -1727,12 +1795,16 @@ impl SessionManager { .await { Ok(Some(meta)) => Ok(UnifiedSessionId::new(channel, chat_id, &meta.dialog_id)), - _ => { + Ok(None) => { let (new_id, _) = self .create_session(channel, chat_id, None, String::new()) .await?; Ok(new_id) } + Err(e) => Err(AgentError::Other(format!( + "failed to find recent session: {}", + e + ))), } } @@ -1752,7 +1824,7 @@ 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 persist_snapshot = { + { let mut guard = session.lock().await; let source = MessageSource { kind: SourceKind::SystemNotification, @@ -1763,11 +1835,11 @@ impl SessionManager { task_id: task_id.map(|s| s.to_string()), }; let msg = ChatMessage::assistant_with_source(content, source); - guard.add_message_in_memory(msg, true) - }; - persist_added_message(persist_snapshot) - .await - .map_err(|e| AgentError::Other(format!("persist error: {}", e)))?; + guard + .add_message(msg, true) + .await + .map_err(|e| AgentError::Other(format!("persist error: {}", e)))?; + } let outbound = OutboundMessage { channel: channel.to_string(), @@ -1841,7 +1913,7 @@ impl SessionManager { guard.current_cancel = None; guard.worker_generation = guard.worker_generation.wrapping_add(1); let generation = guard.worker_generation; - let (tx, rx) = mpsc::unbounded_channel(); + let (tx, rx) = mpsc::channel(SESSION_QUEUE_CAPACITY); guard.agent_tx = Some(tx); spawn_agent_worker( rx, @@ -1853,14 +1925,25 @@ impl SessionManager { unified_str.clone(), ); } - if let Err(e) = guard.agent_tx.as_ref().unwrap().send(task) { + let Some(agent_tx) = guard.agent_tx.as_ref() else { + return Err(AgentError::Other( + "agent worker queue was not initialized".to_string(), + )); + }; + if let Err(e) = agent_tx.try_send(task) { + if matches!(e, mpsc::error::TrySendError::Full(_)) { + tracing::warn!(session_id = %unified_str, capacity = SESSION_QUEUE_CAPACITY, "Session queue is full"); + return Ok(HandleResult::CommandOutput( + "当前对话消息队列已满，请稍后重试。".to_string(), + )); + } // Worker died after we just spawned it — respawn with the recovered task - let task = e.0; + let task = e.into_inner(); guard.agent_tx = None; guard.current_cancel = None; guard.worker_generation = guard.worker_generation.wrapping_add(1); let generation = guard.worker_generation; - let (tx, rx) = mpsc::unbounded_channel(); + let (tx, rx) = mpsc::channel(SESSION_QUEUE_CAPACITY); guard.agent_tx = Some(tx); spawn_agent_worker( rx, @@ -1874,11 +1957,17 @@ impl SessionManager { guard .agent_tx .as_ref() - .unwrap() - .send(task) - .unwrap_or_else(|_| { - tracing::error!("Agent worker spawn+send failed irrecoverably"); - }); + .ok_or_else(|| { + AgentError::Other( + "agent worker queue was not initialized".to_string(), + ) + })? + .try_send(task) + .map_err(|_| { + AgentError::Other( + "agent worker spawn+send failed irrecoverably".to_string(), + ) + })?; } } Ok(HandleResult::AgentProcessing) @@ -1938,14 +2027,48 @@ async fn persist_added_message( }; storage - .append_message_with_retry(&session_id, &msg_meta) - .await?; - storage.upsert_session(&session_meta).await?; - Ok(()) + .persist_message_batch_with_retry( + &session_id, + std::slice::from_ref(&msg_meta), + &session_meta, + ) + .await +} + +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 } fn spawn_agent_worker( - mut task_rx: mpsc::UnboundedReceiver, + mut task_rx: mpsc::Receiver, session: Arc>, bus: Arc, memory_manager: Arc, @@ -1956,7 +2079,7 @@ fn spawn_agent_worker( tokio::spawn(async move { let unified_for_source = unified_str.clone(); let _scope = CURRENT_SOURCE_SESSION.scope(Some(unified_for_source), async { - while let Some(task) = task_rx.recv().await { + 'tasks: while let Some(task) = task_rx.recv().await { let task_chan = task.channel.clone(); let task_cid = task.chat_id.clone(); @@ -1999,9 +2122,11 @@ fn spawn_agent_worker( let media_refs: Vec = 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); - drop(guard); 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(), @@ -2013,9 +2138,8 @@ fn spawn_agent_worker( metadata: HashMap::new(), }; let _ = bus.publish_outbound(err_outbound).await; - return; + continue 'tasks; } - let mut guard = session.lock().await; if guard.worker_generation != worker_gen { return; } @@ -2036,7 +2160,7 @@ fn spawn_agent_worker( metadata: HashMap::new(), }; let _ = bus.publish_outbound(err_outbound).await; - return; + continue 'tasks; } }; @@ -2103,7 +2227,8 @@ fn spawn_agent_worker( session_id = %guard.id, "Session changed while preparing agent history; dropping stale task" ); - return; + guard.current_cancel = None; + continue 'tasks; } if result.created_timelines { guard.last_compressed_message_at = @@ -2267,19 +2392,35 @@ fn spawn_agent_worker( 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); - (result.final_response.content, persist_snapshots) - }; - let (response, persist_snapshots) = response; - for snapshot in persist_snapshots { - if let Err(e) = persist_added_message(snapshot).await { - tracing::error!(error = %e, "Failed to persist message"); + if let Err(e) = persist_added_messages(persist_snapshots).await { + guard.rollback_message_suffix(&message_ids); + tracing::error!(error = %e, "Failed to atomically persist agent turn"); + None + } else { + Some(result.final_response.content) } - } + }; + + let Some(response) = response else { + let err_outbound = OutboundMessage { + channel: chan2, + chat_id: cid2, + content: "Failed to save the agent response, please try again." + .to_string(), + reply_to: None, + media: vec![], + metadata: HashMap::new(), + }; + let _ = bus2.publish_outbound(err_outbound).await; + return; + }; if let Err(e) = maybe_generate_title_outside_lock(session2.clone()).await { tracing::warn!("failed to generate title: {}", e); @@ -2456,14 +2597,14 @@ impl OutboundMessenger for SessionManager { }; // Write source-tagged assistant message to target session history - let persist_snapshot = { + { let mut guard = session.lock().await; let msg = ChatMessage::assistant_with_source(marked_content.clone(), source); - guard.add_message_in_memory(msg, true) - }; - persist_added_message(persist_snapshot) - .await - .map_err(|e| e.to_string())?; + guard + .add_message(msg, true) + .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 { diff --git a/src/storage/mod.rs b/src/storage/mod.rs index 10f4a5a..78221a5 100644 --- a/src/storage/mod.rs +++ b/src/storage/mod.rs @@ -620,6 +620,97 @@ impl Storage { Ok(seqs) } + /// Atomically persist all messages produced by one logical turn together + /// with the resulting session metadata. A turn is either fully visible + /// after restart or not visible at all. + pub async fn persist_message_batch( + &self, + session_id: &str, + msgs: &[crate::storage::message::MessageMeta], + meta: &crate::storage::session::SessionMeta, + ) -> Result<(), StorageError> { + let mut tx = self.pool.begin().await?; + + for msg in msgs { + sqlx::query( + r#" + INSERT INTO messages (id, session_id, seq, role, content, reasoning_content, media_refs, tool_call_id, tool_name, tool_calls, source, created_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + "#, + ) + .bind(&msg.id) + .bind(session_id) + .bind(msg.seq) + .bind(&msg.role) + .bind(&msg.content) + .bind(&msg.reasoning_content) + .bind(&msg.media_refs) + .bind(&msg.tool_call_id) + .bind(&msg.tool_name) + .bind(&msg.tool_calls) + .bind(&msg.source) + .bind(msg.created_at) + .execute(&mut *tx) + .await?; + } + + sqlx::query( + r#" + INSERT INTO sessions (id, channel, chat_id, dialog_id, title, created_at, last_active_at, message_count, routing_info, archived_at, deleted_at, last_consolidated_at, last_compressed_message_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT(id) DO UPDATE SET + title = excluded.title, + last_active_at = excluded.last_active_at, + message_count = excluded.message_count, + routing_info = excluded.routing_info, + archived_at = excluded.archived_at, + deleted_at = excluded.deleted_at, + last_consolidated_at = excluded.last_consolidated_at, + last_compressed_message_at = excluded.last_compressed_message_at + "#, + ) + .bind(&meta.id) + .bind(&meta.channel) + .bind(&meta.chat_id) + .bind(&meta.dialog_id) + .bind(&meta.title) + .bind(meta.created_at) + .bind(meta.last_active_at) + .bind(meta.message_count) + .bind(&meta.routing_info) + .bind(meta.archived_at) + .bind(meta.deleted_at) + .bind(meta.last_consolidated_at) + .bind(meta.last_compressed_message_at) + .execute(&mut *tx) + .await?; + + tx.commit().await?; + Ok(()) + } + + /// Persist a turn with bounded retry. Retrying the whole transaction keeps + /// message rows and metadata consistent on transient SQLite failures. + pub async fn persist_message_batch_with_retry( + &self, + session_id: &str, + msgs: &[crate::storage::message::MessageMeta], + meta: &crate::storage::session::SessionMeta, + ) -> Result<(), StorageError> { + let delays = [100, 200, 300]; + for (attempt, delay) in delays.iter().enumerate() { + match self.persist_message_batch(session_id, msgs, meta).await { + Ok(()) => return Ok(()), + Err(error) if attempt < delays.len() - 1 => { + tracing::warn!(attempt = attempt + 1, error = %error, "Turn persistence failed; retrying"); + sleep(Duration::from_millis(*delay)).await; + } + Err(error) => return Err(error), + } + } + unreachable!() + } + pub async fn load_messages( &self, session_id: &str, @@ -1197,6 +1288,54 @@ mod tests { assert_eq!(loaded[0].content, "你好"); } + #[tokio::test] + async fn test_persist_message_batch_is_atomic() { + let (storage, _dir) = create_test_storage().await; + let session_meta = crate::storage::session::SessionMeta { + id: "cli_chat:atomic:dialog1".to_string(), + channel: "cli_chat".to_string(), + chat_id: "atomic".to_string(), + dialog_id: "dialog1".to_string(), + title: "Atomic turn".to_string(), + created_at: 1000, + last_active_at: 2000, + message_count: 1, + routing_info: None, + archived_at: None, + deleted_at: None, + last_consolidated_at: None, + last_compressed_message_at: None, + }; + storage.upsert_session(&session_meta).await.unwrap(); + + let message = crate::storage::message::MessageMeta { + id: "duplicate-id".to_string(), + session_id: session_meta.id.clone(), + seq: 1, + role: "assistant".to_string(), + content: "must roll back".to_string(), + reasoning_content: None, + media_refs: None, + tool_call_id: None, + tool_name: None, + tool_calls: None, + source: None, + created_at: 2000, + }; + + let result = storage + .persist_message_batch(&session_meta.id, &[message.clone(), message], &session_meta) + .await; + assert!(result.is_err()); + assert!( + storage + .load_messages(&session_meta.id, 0) + .await + .unwrap() + .is_empty() + ); + } + #[tokio::test] async fn test_touch_session() { let (storage, _dir) = create_test_storage().await;