fix(runtime): harden session persistence and outbound dispatch
This commit is contained in:
parent
5079458f55
commit
b06bc4f025
@ -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
|
||||
|
||||
|
||||
@ -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<MessageBus>,
|
||||
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<String, mpsc::Sender<OutboundMessage>> = 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<dyn Channel + Send + Sync>,
|
||||
mut receiver: mpsc::Receiver<OutboundMessage>,
|
||||
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<Vec<String>>,
|
||||
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<MessageBus>) -> 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();
|
||||
}
|
||||
}
|
||||
|
||||
@ -15,6 +15,8 @@ type MessagePersistSnapshot = (
|
||||
crate::storage::session::SessionMeta,
|
||||
);
|
||||
|
||||
const SESSION_QUEUE_CAPACITY: usize = 32;
|
||||
|
||||
tokio::task_local! {
|
||||
static CURRENT_SOURCE_SESSION: Option<String>;
|
||||
}
|
||||
@ -84,7 +86,7 @@ pub struct Session {
|
||||
memory_manager: Arc<crate::memory::MemoryManager>,
|
||||
|
||||
/// Task queue for per-session serial agent processing
|
||||
agent_tx: Option<mpsc::UnboundedSender<AgentTask>>,
|
||||
agent_tx: Option<mpsc::Sender<AgentTask>>,
|
||||
/// Cancel signal for the currently executing agent task
|
||||
current_cancel: Option<oneshot::Sender<()>>,
|
||||
/// 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<ChatMessage> = 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<Arc<Mutex<Session>>, 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<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
|
||||
}
|
||||
|
||||
fn spawn_agent_worker(
|
||||
mut task_rx: mpsc::UnboundedReceiver<AgentTask>,
|
||||
mut task_rx: mpsc::Receiver<AgentTask>,
|
||||
session: Arc<Mutex<Session>>,
|
||||
bus: Arc<MessageBus>,
|
||||
memory_manager: Arc<crate::memory::MemoryManager>,
|
||||
@ -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<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);
|
||||
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 {
|
||||
|
||||
@ -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;
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user