PicoBot/src/gateway/router.rs

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<MessageBus>,
    session_manager: Arc<SessionManager>,
    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<MessageBus>,
    session_manager: Arc<SessionManager>,
    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<String, mpsc::Sender<AdmittedInbound>> = 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<MessageBus>,
    session_manager: Arc<SessionManager>,
    channel: String,
    chat_id: String,
    receiver: mpsc::Receiver<AdmittedInbound>,
) -> 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<T, F, Fut>(
    mut receiver: mpsc::Receiver<T>,
    idle_timeout: Duration,
    mut handler: F,
) where
    T: Send + 'static,
    F: FnMut(T) -> Fut,
    Fut: Future<Output = ()>,
{
    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<MessageBus>,
    session_manager: Arc<SessionManager>,
    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<MessageBus>,
    session_manager: Arc<SessionManager>,
    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<SessionEvent, ChannelError> = 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();
    }
}