320 lines
10 KiB
Rust
320 lines
10 KiB
Rust
use std::collections::{HashMap, HashSet};
|
||
use std::sync::Arc;
|
||
use tokio::sync::RwLock;
|
||
|
||
use crate::bus::MessageBus;
|
||
use crate::channels::base::{Channel, ChannelError};
|
||
use crate::channels::cli::CliChannel;
|
||
use crate::channels::feishu::FeishuChannel;
|
||
use crate::channels::wechat::WechatChannel;
|
||
use crate::config::{Config, TaggedChannelConfig};
|
||
use crate::protocol::Channel as ProtocolChannel;
|
||
|
||
/// ChannelManager manages all Channel instances and the MessageBus
|
||
#[derive(Clone)]
|
||
pub struct ChannelManager {
|
||
channels: Arc<RwLock<HashMap<String, Arc<dyn Channel + Send + Sync>>>>,
|
||
bus: Arc<MessageBus>,
|
||
websocket_channel: Arc<CliChannel>,
|
||
}
|
||
|
||
impl ChannelManager {
|
||
pub fn new() -> Self {
|
||
let websocket_channel = Arc::new(CliChannel::new());
|
||
let mut channels: HashMap<String, Arc<dyn Channel + Send + Sync>> = HashMap::new();
|
||
channels.insert("websocket".to_string(), websocket_channel.clone());
|
||
|
||
Self {
|
||
channels: Arc::new(RwLock::new(channels)),
|
||
bus: MessageBus::new(100),
|
||
websocket_channel,
|
||
}
|
||
}
|
||
|
||
/// Get a reference to the MessageBus
|
||
pub fn bus(&self) -> Arc<MessageBus> {
|
||
self.bus.clone()
|
||
}
|
||
|
||
pub fn websocket_channel(&self) -> Arc<CliChannel> {
|
||
self.websocket_channel.clone()
|
||
}
|
||
|
||
/// Initialize all Channel instances from config
|
||
pub async fn init(
|
||
&self,
|
||
config: &Config,
|
||
provider_config: crate::config::LLMProviderConfig,
|
||
) -> Result<(), ChannelError> {
|
||
for (name, channel_config) in &config.channels {
|
||
match channel_config {
|
||
crate::config::ChannelConfig::Tagged(TaggedChannelConfig::Feishu(feishu_config))
|
||
| crate::config::ChannelConfig::LegacyFeishu(feishu_config) => {
|
||
if feishu_config.enabled {
|
||
let channel = FeishuChannel::new(
|
||
name.clone(),
|
||
feishu_config.clone(),
|
||
provider_config.clone(),
|
||
)
|
||
.map_err(|e| {
|
||
ChannelError::Other(format!(
|
||
"Failed to create Feishu channel '{}': {}",
|
||
name, e
|
||
))
|
||
})?;
|
||
|
||
self.channels
|
||
.write()
|
||
.await
|
||
.insert(name.clone(), Arc::new(channel));
|
||
tracing::info!(channel = %name, kind = channel_config.kind(), "Channel registered");
|
||
} else {
|
||
tracing::info!(channel = %name, kind = channel_config.kind(), "Channel disabled in config");
|
||
}
|
||
}
|
||
crate::config::ChannelConfig::Tagged(TaggedChannelConfig::Wechat(wechat_config)) => {
|
||
if wechat_config.enabled {
|
||
let channel = WechatChannel::new(
|
||
name.clone(),
|
||
wechat_config.clone(),
|
||
provider_config.clone(),
|
||
)
|
||
.map_err(|e| {
|
||
ChannelError::Other(format!(
|
||
"Failed to create WeChat channel '{}': {}",
|
||
name, e
|
||
))
|
||
})?;
|
||
|
||
self.channels
|
||
.write()
|
||
.await
|
||
.insert(name.clone(), Arc::new(channel));
|
||
tracing::info!(channel = %name, kind = channel_config.kind(), "Channel registered");
|
||
} else {
|
||
tracing::info!(channel = %name, kind = channel_config.kind(), "Channel disabled in config");
|
||
}
|
||
}
|
||
}
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
pub async fn start_all(&self) -> Result<(), ChannelError> {
|
||
let channels = self.channels.read().await;
|
||
let bus = self.bus.clone();
|
||
for (name, channel) in channels.iter() {
|
||
tracing::info!(channel = %name, "Starting channel");
|
||
if let Err(e) = channel.start(bus.clone()).await {
|
||
tracing::error!(channel = %name, error = %e, "Failed to start channel");
|
||
}
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
pub async fn stop_all(&self) -> Result<(), ChannelError> {
|
||
let mut channels = self.channels.write().await;
|
||
for (name, channel) in channels.iter() {
|
||
tracing::info!(channel = %name, "Stopping channel");
|
||
if let Err(e) = channel.stop().await {
|
||
tracing::error!(channel = %name, error = %e, "Error stopping channel");
|
||
}
|
||
}
|
||
channels.clear();
|
||
Ok(())
|
||
}
|
||
|
||
pub async fn get_channel(&self, name: &str) -> Option<Arc<dyn Channel + Send + Sync>> {
|
||
self.channels.read().await.get(name).cloned()
|
||
}
|
||
|
||
pub async fn channels(&self) -> Vec<(String, Arc<dyn Channel + Send + Sync>)> {
|
||
self.channels
|
||
.read()
|
||
.await
|
||
.iter()
|
||
.map(|(name, channel)| (name.clone(), channel.clone()))
|
||
.collect()
|
||
}
|
||
|
||
/// 构建面向前端的通道列表（合并 websocket + 动态注册的通道）
|
||
pub async fn build_channel_list(&self) -> Vec<ProtocolChannel> {
|
||
let mut seen = HashSet::new();
|
||
let mut channels: Vec<ProtocolChannel> = Vec::new();
|
||
|
||
// 所有注册的通道（websocket, feishu, wechat 等）
|
||
for (name, _channel) in self.channels().await {
|
||
if seen.contains(&name) {
|
||
continue;
|
||
}
|
||
seen.insert(name.clone());
|
||
channels.push(ProtocolChannel {
|
||
id: name.clone(),
|
||
name: ChannelManager::channel_display_name(&name),
|
||
description: ChannelManager::channel_description(&name),
|
||
is_writable: ChannelManager::is_channel_writable(&name),
|
||
});
|
||
}
|
||
|
||
channels
|
||
}
|
||
|
||
/// 通道名称 → 显示名称
|
||
fn channel_display_name(name: &str) -> String {
|
||
match name {
|
||
"websocket" => "WebSocket".to_string(),
|
||
"feishu" => "飞书".to_string(),
|
||
"wechat" => "微信".to_string(),
|
||
other => other.to_string(),
|
||
}
|
||
}
|
||
|
||
/// 通道名称 → 描述
|
||
fn channel_description(name: &str) -> Option<String> {
|
||
match name {
|
||
"websocket" => Some("Web 前端通道".to_string()),
|
||
"cli" => Some("命令行终端通道".to_string()),
|
||
"feishu" => Some("飞书消息通道".to_string()),
|
||
"wechat" => Some("微信消息通道".to_string()),
|
||
_ => None,
|
||
}
|
||
}
|
||
|
||
/// 判断通道是否可写（从 Web 前端视角）
|
||
fn is_channel_writable(name: &str) -> bool {
|
||
// 只有 WebSocket 通道可写，其他通道（CLI、飞书、微信等）均为只读
|
||
name == "websocket"
|
||
}
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
|
||
fn write_test_config() -> tempfile::NamedTempFile {
|
||
let file = tempfile::NamedTempFile::new().unwrap();
|
||
std::fs::write(
|
||
file.path(),
|
||
r#"{
|
||
"providers": {
|
||
"aliyun": {
|
||
"type": "openai",
|
||
"base_url": "https://example.invalid/v1",
|
||
"api_key": "test-key",
|
||
"extra_headers": {}
|
||
}
|
||
},
|
||
"models": {
|
||
"qwen-plus": {
|
||
"model_id": "qwen-plus"
|
||
}
|
||
},
|
||
"agents": {
|
||
"default": {
|
||
"provider": "aliyun",
|
||
"model": "qwen-plus"
|
||
}
|
||
},
|
||
"channels": {
|
||
"primary": {
|
||
"type": "feishu",
|
||
"enabled": true,
|
||
"app_id": "app-id-1",
|
||
"app_secret": "secret-1"
|
||
},
|
||
"backup": {
|
||
"type": "feishu",
|
||
"enabled": true,
|
||
"app_id": "app-id-2",
|
||
"app_secret": "secret-2"
|
||
}
|
||
}
|
||
}"#,
|
||
)
|
||
.unwrap();
|
||
file
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn init_registers_all_configured_channels_by_instance_name() {
|
||
let file = write_test_config();
|
||
let config = Config::load(file.path().to_str().unwrap()).unwrap();
|
||
let provider_config = config.get_provider_config("default").unwrap();
|
||
let manager = ChannelManager::new();
|
||
|
||
manager.init(&config, provider_config).await.unwrap();
|
||
|
||
let mut names = manager
|
||
.channels()
|
||
.await
|
||
.into_iter()
|
||
.map(|(name, _)| name)
|
||
.collect::<Vec<_>>();
|
||
names.sort();
|
||
|
||
assert_eq!(names, vec!["backup", "primary", "websocket"]);
|
||
assert_eq!(manager.get_channel("primary").await.unwrap().name(), "primary");
|
||
assert_eq!(manager.get_channel("backup").await.unwrap().name(), "backup");
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn init_registers_wechat_channel_by_instance_name() {
|
||
let file = tempfile::NamedTempFile::new().unwrap();
|
||
// 使用临时目录确保跨平台兼容
|
||
let temp_dir = tempfile::tempdir().unwrap();
|
||
let cred_path = temp_dir.path().join("wechat-creds.json");
|
||
// JSON 中的路径需要转义反斜杠
|
||
let cred_path_json = cred_path.display().to_string().replace('\\', "\\\\");
|
||
|
||
std::fs::write(
|
||
file.path(),
|
||
r#"{
|
||
"providers": {
|
||
"aliyun": {
|
||
"type": "openai",
|
||
"base_url": "https://example.invalid/v1",
|
||
"api_key": "test-key",
|
||
"extra_headers": {}
|
||
}
|
||
},
|
||
"models": {
|
||
"qwen-plus": {
|
||
"model_id": "qwen-plus"
|
||
}
|
||
},
|
||
"agents": {
|
||
"default": {
|
||
"provider": "aliyun",
|
||
"model": "qwen-plus"
|
||
}
|
||
},
|
||
"channels": {
|
||
"wechat_main": {
|
||
"type": "wechat",
|
||
"enabled": true,
|
||
"cred_path": "<CRED_PATH>"
|
||
}
|
||
}
|
||
}"#.replace("<CRED_PATH>", &cred_path_json),
|
||
)
|
||
.unwrap();
|
||
|
||
let config = Config::load(file.path().to_str().unwrap()).unwrap();
|
||
let provider_config = config.get_provider_config("default").unwrap();
|
||
let manager = ChannelManager::new();
|
||
|
||
manager.init(&config, provider_config).await.unwrap();
|
||
|
||
let mut names = manager
|
||
.channels()
|
||
.await
|
||
.into_iter()
|
||
.map(|(name, _)| name)
|
||
.collect::<Vec<_>>();
|
||
names.sort();
|
||
|
||
assert_eq!(names, vec!["websocket", "wechat_main"]);
|
||
assert_eq!(manager.get_channel("wechat_main").await.unwrap().name(), "wechat_main");
|
||
}
|
||
}
|