use std::collections::HashMap; use std::future::Future; use std::sync::Arc; use std::time::Duration; use tokio::sync::{Semaphore, mpsc}; use crate::bus::{ControlMessage, InboundMessage, MessageBus, OutboundMessage}; use crate::channels::ChannelError; use crate::channels::parse_slash_command; use crate::gateway::reload::{ActivityGuard, RuntimeAdmission}; use crate::session::{SessionCommand, SessionEvent, SessionManager}; use crate::task_supervisor::TaskSupervisor; const INBOUND_LANE_CAPACITY: usize = 32; const INBOUND_LANE_IDLE_TIMEOUT: Duration = Duration::from_secs(300); const CONTROL_MAX_IN_FLIGHT: usize = 64; pub(super) fn spawn_message_routers( bus: Arc, session_manager: Arc, supervisor: TaskSupervisor, admission: RuntimeAdmission, ) { spawn_inbound_router( bus.clone(), session_manager.clone(), supervisor.clone(), admission, ); spawn_control_router(bus, session_manager, supervisor); } fn spawn_inbound_router( bus: Arc, session_manager: Arc, supervisor: TaskSupervisor, admission: RuntimeAdmission, ) { let lane_supervisor = supervisor.clone(); supervisor.spawn("inbound-router", async move { tracing::info!(lane_capacity = INBOUND_LANE_CAPACITY, "Inbound router started"); let mut lanes: HashMap> = HashMap::new(); let mut messages_seen = 0_u64; while let Some(inbound) = bus.consume_inbound().await { messages_seen = messages_seen.wrapping_add(1); if messages_seen.is_multiple_of(128) { lanes.retain(|_, sender| !sender.is_closed()); } let Some(activity) = admission.try_enter() else { publish_command_output( &bus, inbound, "Gateway 正在重新加载配置,请稍后重试。".to_string(), ) .await; continue; }; let inbound = AdmittedInbound { inbound, activity }; // Stop must be able to invalidate a running worker even when an // earlier slow slash command occupies this conversation's lane. if is_priority_stop(&inbound.inbound.content) { let request_bus = bus.clone(); let request_manager = session_manager.clone(); let task_name = format!( "inbound-stop:{}:{}", inbound.inbound.channel, inbound.inbound.chat_id ); if !lane_supervisor.spawn(task_name, async move { process_inbound(request_bus, request_manager, inbound).await; }) { break; } continue; } let key = conversation_key(&inbound.inbound.channel, &inbound.inbound.chat_id); let mut sender = lanes.get(&key).cloned(); if sender.as_ref().is_none_or(mpsc::Sender::is_closed) { let (new_sender, receiver) = mpsc::channel(INBOUND_LANE_CAPACITY); if !spawn_inbound_lane( &lane_supervisor, bus.clone(), session_manager.clone(), inbound.inbound.channel.clone(), inbound.inbound.chat_id.clone(), receiver, ) { tracing::warn!("Inbound router is stopping"); break; } lanes.insert(key.clone(), new_sender.clone()); sender = Some(new_sender); } let Some(sender) = sender else { tracing::error!("Inbound lane creation did not produce a sender"); continue; }; match sender.try_send(inbound) { Ok(()) => {} Err(mpsc::error::TrySendError::Full(inbound)) => { tracing::warn!(channel = %inbound.inbound.channel, chat_id = %inbound.inbound.chat_id, "Inbound conversation lane is full"); publish_command_output( &bus, inbound.inbound, "当前对话入口队列已满,请稍后重试。".to_string(), ) .await; } Err(mpsc::error::TrySendError::Closed(inbound)) => { // The lane may have exited on its idle boundary between the // closed check and try_send. Recreate it once without // dropping this input. let (new_sender, receiver) = mpsc::channel(INBOUND_LANE_CAPACITY); if !spawn_inbound_lane( &lane_supervisor, bus.clone(), session_manager.clone(), inbound.inbound.channel.clone(), inbound.inbound.chat_id.clone(), receiver, ) { break; } lanes.insert(key, new_sender.clone()); if new_sender.try_send(inbound).is_err() { tracing::error!("Failed to enqueue input into replacement lane"); } } } } tracing::warn!("Inbound router stopped because inbound bus closed"); }); } fn spawn_inbound_lane( supervisor: &TaskSupervisor, bus: Arc, session_manager: Arc, channel: String, chat_id: String, receiver: mpsc::Receiver, ) -> bool { supervisor.spawn(format!("inbound-lane:{channel}:{chat_id}"), async move { run_ordered_lane(receiver, INBOUND_LANE_IDLE_TIMEOUT, move |inbound| { process_inbound(bus.clone(), session_manager.clone(), inbound) }) .await; }) } async fn run_ordered_lane( mut receiver: mpsc::Receiver, idle_timeout: Duration, mut handler: F, ) where T: Send + 'static, F: FnMut(T) -> Fut, Fut: Future, { loop { let item = match tokio::time::timeout(idle_timeout, receiver.recv()).await { Ok(Some(item)) => item, Ok(None) | Err(_) => break, }; handler(item).await; } } async fn process_inbound( bus: Arc, session_manager: Arc, admitted: AdmittedInbound, ) { let AdmittedInbound { inbound, activity: _activity, } = admitted; let result = session_manager.handle_message(&inbound).await; match result { Ok(crate::session::session::HandleResult::AgentResponse(content)) => { publish_assistant_output(&bus, inbound, content).await; } Ok(crate::session::session::HandleResult::CommandOutput(content)) => { publish_command_output(&bus, inbound, content).await; } Ok(crate::session::session::HandleResult::AgentProcessing) => {} Err(error) => { tracing::error!(channel = %inbound.channel, chat_id = %inbound.chat_id, error = %error, "Failed to handle inbound message"); publish_command_output(&bus, inbound, "消息处理失败,请稍后重试。".to_string()).await; } } } async fn publish_assistant_output(bus: &MessageBus, inbound: InboundMessage, content: String) { publish_output(bus, inbound, content, false, false).await; } async fn publish_command_output(bus: &MessageBus, inbound: InboundMessage, content: String) { publish_output(bus, inbound, content, true, true).await; } async fn publish_output( bus: &MessageBus, inbound: InboundMessage, content: String, command: bool, confirmed: bool, ) { let mut metadata = inbound.channel_context.private; if command { metadata.insert("_type".to_string(), "command".to_string()); } let outbound = OutboundMessage { channel: inbound.channel, chat_id: inbound.chat_id, content, reply_to: inbound.channel_context.reply_to, media: vec![], metadata, delivery: None, }; let result = if confirmed { bus.deliver_outbound(outbound).await } else { bus.publish_outbound(outbound).await }; if let Err(error) = result { tracing::error!(error = %error, "Failed to publish routed outbound message"); } } struct AdmittedInbound { inbound: InboundMessage, activity: ActivityGuard, } fn spawn_control_router( bus: Arc, session_manager: Arc, supervisor: TaskSupervisor, ) { let request_supervisor = supervisor.clone(); supervisor.spawn("control-router", async move { tracing::info!( max_in_flight = CONTROL_MAX_IN_FLIGHT, "Control router started" ); let permits = Arc::new(Semaphore::new(CONTROL_MAX_IN_FLIGHT)); loop { let permit = match permits.clone().acquire_owned().await { Ok(permit) => permit, Err(_) => break, }; let Some(message) = bus.consume_control().await else { break; }; let manager = session_manager.clone(); if !request_supervisor.spawn("control-request", async move { let _permit = permit; handle_control_message(&manager, message).await; }) { break; } } tracing::warn!("Control router stopped because control bus closed"); }); } async fn handle_control_message(session_manager: &SessionManager, message: ControlMessage) { use SessionCommand::*; let reply_tx = message.reply_tx; let result: Result = match message.op { CreateDialog { channel, chat_id, title, } => session_manager .create_dialog(&channel, &chat_id, title.as_deref()) .await .map(|(session_id, title)| SessionEvent::DialogCreated { session_id, title }) .map_err(|error| ChannelError::Other(error.to_string())), ListDialogs { channel, chat_id, include_archived, } => session_manager .list_dialogs(&channel, &chat_id, include_archived) .await .map(|(dialogs, current_dialog_id)| SessionEvent::DialogList { dialogs, current_dialog_id, }) .map_err(|error| ChannelError::Other(error.to_string())), GetCurrentDialog { channel, chat_id } => session_manager .get_current_dialog(&channel, &chat_id) .await .map(|session_id| SessionEvent::CurrentDialog { session_id }) .map_err(|error| ChannelError::Other(error.to_string())), SwitchDialog { channel, chat_id, dialog_id, } => session_manager .switch_dialog(&channel, &chat_id, &dialog_id) .await .map(|session_id| SessionEvent::DialogSwitched { session_id }) .map_err(|error| ChannelError::Other(error.to_string())), GetDialogHistory { session_id, limit } => session_manager .get_dialog_history(&session_id, limit) .await .map(|messages| SessionEvent::DialogHistory { session_id, messages, }) .map_err(|error| ChannelError::Other(error.to_string())), GetTaskPlan { session_id } => session_manager .get_task_plan(&session_id) .await .map(|plan| SessionEvent::TaskPlan { session_id, plan }) .map_err(|error| ChannelError::Other(error.to_string())), RenameDialog { session_id, title } => session_manager .rename_dialog(&session_id, &title) .await .map(|()| SessionEvent::DialogRenamed { session_id, title }) .map_err(|error| ChannelError::Other(error.to_string())), ArchiveDialog { session_id } => session_manager .archive_dialog(&session_id) .await .map(|()| SessionEvent::DialogArchived { session_id }) .map_err(|error| ChannelError::Other(error.to_string())), DeleteDialog { session_id } => session_manager .delete_dialog(&session_id) .await .map(|()| SessionEvent::DialogDeleted { session_id }) .map_err(|error| ChannelError::Other(error.to_string())), ClearHistory { session_id } => session_manager .clear_dialog_history(&session_id) .await .map(|()| SessionEvent::HistoryCleared { session_id }) .map_err(|error| ChannelError::Other(error.to_string())), GetSlashCommands { .. } => Ok(SessionEvent::SlashCommandsList { commands: session_manager.get_slash_commands().to_vec(), }), ExecuteSlashCommand { command, args, channel, chat_id, current_session_id, } => session_manager .execute_slash_command( &command, args.as_deref(), &channel, &chat_id, current_session_id.as_ref(), ) .await .map( |(new_session_id, message)| SessionEvent::SlashCommandExecuted { new_session_id, message, }, ) .map_err(|error| ChannelError::Other(error.to_string())), }; let _ = reply_tx.send(result).await; } fn conversation_key(channel: &str, chat_id: &str) -> String { format!("{channel}\0{chat_id}") } fn is_priority_stop(content: &str) -> bool { parse_slash_command(content).is_some_and(|(command, _)| command == "stop") } #[cfg(test)] mod tests { use super::*; use crate::bus::ChannelContext; use std::collections::HashSet; use tokio::sync::Notify; #[test] fn only_stop_bypasses_a_conversation_lane() { assert!(is_priority_stop("/stop")); assert!(is_priority_stop(" /stop ")); assert!(!is_priority_stop("/compact")); assert!(!is_priority_stop("normal message")); } #[test] fn conversation_keys_do_not_alias() { let keys = HashSet::from([ conversation_key("a", "bc"), conversation_key("ab", "c"), conversation_key("a", "bd"), ]); assert_eq!(keys.len(), 3); } #[tokio::test] async fn routed_output_preserves_reply_target_and_private_context() { let bus = MessageBus::new(2); let inbound = InboundMessage { channel: "test".to_string(), sender_id: "user".to_string(), chat_id: "chat".to_string(), content: "hello".to_string(), received_at: 123, media: vec![], channel_context: ChannelContext { reply_to: Some("parent".to_string()), private: HashMap::from([("opaque".to_string(), "value".to_string())]), }, }; let publish_task = tokio::spawn({ let bus = bus.clone(); async move { publish_command_output(&bus, inbound, "done".to_string()).await } }); let output = bus.consume_outbound().await.unwrap(); assert_eq!(output.reply_to.as_deref(), Some("parent")); assert_eq!( output.metadata.get("opaque").map(String::as_str), Some("value") ); assert_eq!( output.metadata.get("_type").map(String::as_str), Some("command") ); assert!(!publish_task.is_finished()); output.complete_delivery(Ok(())); publish_task.await.unwrap(); } #[tokio::test] async fn slow_conversation_lane_does_not_block_another_lane() { let (slow_tx, slow_rx) = mpsc::channel(2); let (fast_tx, fast_rx) = mpsc::channel(2); let slow_started = Arc::new(Notify::new()); let release_slow = Arc::new(Notify::new()); let fast_finished = Arc::new(Notify::new()); let slow_task = tokio::spawn({ let slow_started = slow_started.clone(); let release_slow = release_slow.clone(); async move { run_ordered_lane(slow_rx, Duration::from_secs(1), move |_| { let slow_started = slow_started.clone(); let release_slow = release_slow.clone(); async move { slow_started.notify_one(); release_slow.notified().await; } }) .await; } }); let fast_task = tokio::spawn({ let fast_finished = fast_finished.clone(); async move { run_ordered_lane(fast_rx, Duration::from_secs(1), move |_| { let fast_finished = fast_finished.clone(); async move { fast_finished.notify_one() } }) .await; } }); slow_tx.send("slow").await.unwrap(); slow_started.notified().await; fast_tx.send("fast").await.unwrap(); tokio::time::timeout(Duration::from_millis(100), fast_finished.notified()) .await .expect("fast lane was blocked by unrelated slow lane"); release_slow.notify_one(); drop(slow_tx); drop(fast_tx); slow_task.await.unwrap(); fast_task.await.unwrap(); } }