PicoBot/src/channels/cli.rs

use async_trait::async_trait;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::{RwLock, mpsc};
use tokio_util::sync::CancellationToken;

use crate::bus::{MessageBus, OutboundMessage};
use crate::protocol::WsOutbound;
use crate::protocol::ws_adapter::ws_outbound_from_outbound_message;

use super::base::{Channel, ChannelError};

#[derive(Clone)]
struct CliConnection {
    connection_id: String,
    sender: mpsc::Sender<WsOutbound>,
}

#[derive(Clone)]
pub struct CliChannel {
    connections: Arc<RwLock<HashMap<String, CliConnection>>>,
    /// 全局关闭信号：stop() 时 cancel，所有 handle_socket 同时收到信号
    shutdown_token: CancellationToken,
}

impl CliChannel {
    pub fn new() -> Self {
        Self {
            connections: Arc::new(RwLock::new(HashMap::new())),
            shutdown_token: CancellationToken::new(),
        }
    }

    /// 注册连接，返回 CancellationToken 供 handle_socket 监听
    ///
    /// 注意：此方法不会向被替换的旧连接发送关闭信号。
    /// 关闭信号仅由 stop() 统一触发，避免正常消息处理时误关连接。
    pub async fn register_connection(
        &self,
        session_id: impl Into<String>,
        connection_id: impl Into<String>,
        sender: mpsc::Sender<WsOutbound>,
    ) -> CancellationToken {
        self.connections.write().await.insert(
            session_id.into(),
            CliConnection {
                connection_id: connection_id.into(),
                sender,
            },
        );
        self.shutdown_token.clone()
    }

    pub async fn unregister_connection(&self, connection_id: &str) {
        self.connections
            .write()
            .await
            .retain(|_, connection| connection.connection_id != connection_id);
    }
}

impl Default for CliChannel {
    fn default() -> Self {
        Self::new()
    }
}

#[async_trait]
impl Channel for CliChannel {
    fn name(&self) -> &str {
        "websocket"
    }

    fn is_running(&self) -> bool {
        true
    }

    async fn start(&self, _bus: Arc<MessageBus>) -> Result<(), ChannelError> {
        Ok(())
    }

    async fn stop(&self) -> Result<(), ChannelError> {
        // 先 cancel 所有 handle_socket，再清空 connections
        self.shutdown_token.cancel();
        let count = self.connections.read().await.len();
        self.connections.write().await.clear();
        tracing::info!(connection_count = count, "CliChannel stopped, all connections signaled to close");
        Ok(())
    }

    async fn send(&self, msg: OutboundMessage) -> Result<(), ChannelError> {
        let connection = self.connections.read().await.get(&msg.chat_id).cloned();
        let Some(connection) = connection else {
            return Err(ChannelError::SendError(format!(
                "No active CLI connection for session {}",
                msg.chat_id
            )));
        };

        for outbound in ws_outbound_from_outbound_message(&msg) {
            connection
                .sender
                .send(outbound)
                .await
                .map_err(|_| ChannelError::SendError("CLI websocket sender closed".to_string()))?;
        }

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::bus::OutboundMessage;

    #[tokio::test]
    async fn test_cli_channel_sends_to_registered_session() {
        let channel = CliChannel::new();
        let (sender, mut receiver) = mpsc::channel(4);
        let _token = channel
            .register_connection("session-1", "conn-1", sender)
            .await;

        channel
            .send(OutboundMessage::assistant(
                "cli",
                "session-1",
                None, // session_id
                "hello",
                None,
                HashMap::new(),
            ))
            .await
            .unwrap();

        let outbound = receiver.recv().await.unwrap();
        assert!(matches!(outbound, WsOutbound::AssistantResponse { .. }));
    }

    #[tokio::test]
    async fn test_cli_channel_unregisters_connection_sessions() {
        let channel = CliChannel::new();
        let (sender, _receiver) = mpsc::channel(4);
        let _token = channel
            .register_connection("session-1", "conn-1", sender)
            .await;
        channel.unregister_connection("conn-1").await;

        let error = channel
            .send(OutboundMessage::assistant(
                "cli",
                "session-1",
                None, // session_id
                "hello",
                None,
                HashMap::new(),
            ))
            .await
            .unwrap_err();

        assert!(error.to_string().contains("No active CLI connection"));
    }

    #[tokio::test]
    async fn test_stop_cancels_token() {
        let channel = CliChannel::new();
        let (sender, _receiver) = mpsc::channel(4);
        let token = channel
            .register_connection("session-1", "conn-1", sender)
            .await;

        assert!(!token.is_cancelled());
        channel.stop().await.unwrap();
        assert!(token.is_cancelled());
    }

    #[tokio::test]
    async fn test_register_connection_does_not_cancel_old_token() {
        // 验证：register_connection 替换旧连接时不会 cancel token
        // 只有 stop() 才会 cancel
        let channel = CliChannel::new();
        let (sender1, _receiver1) = mpsc::channel(4);
        let token1 = channel
            .register_connection("session-1", "conn-1", sender1)
            .await;

        let (sender2, _receiver2) = mpsc::channel(4);
        let _token2 = channel
            .register_connection("session-1", "conn-2", sender2)
            .await;

        // 旧 token 不应被 cancel（只有 stop() 才 cancel）
        assert!(!token1.is_cancelled());
    }
}